From deb2f5e7938274deaf100230a2566063cc851e11 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josh@tailscale.com>
Date: Tue, 30 Nov 2021 14:32:30 -0800
Subject: [PATCH] types/logger: add Context and related helpers

We often need both a log function and a context.
We can do this by adding the log function as a context value.
This commit adds helper glue to make that easy.
It is designed to allow incremental adoption.

Updates tailscale/corp#3138

Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>
---
 types/logger/logger.go      | 23 +++++++++++++++++++++++
 types/logger/logger_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/types/logger/logger.go b/types/logger/logger.go
index 2f8973abe..619eba1e1 100644
--- a/types/logger/logger.go
+++ b/types/logger/logger.go
@@ -18,6 +18,8 @@
 	"strings"
 	"sync"
 	"time"
+
+	"context"
 )
 
 // Logf is the basic Tailscale logger type: a printf-like func.
@@ -25,6 +27,27 @@
 // Logf functions must be safe for concurrent use.
 type Logf func(format string, args ...interface{})
 
+// A Context is a context.Context that should contain a custom log function, obtainable from FromContext.
+// If no log function is present, FromContext will return log.Printf.
+// To construct a Context, use Add
+type Context context.Context
+
+type logfKey struct{}
+
+// FromContext extracts a log function from ctx.
+func FromContext(ctx Context) Logf {
+	v := ctx.Value(logfKey{})
+	if v == nil {
+		return log.Printf
+	}
+	return v.(Logf)
+}
+
+// Ctx constructs a Context from ctx with fn as its custom log function.
+func Ctx(ctx context.Context, fn Logf) Context {
+	return context.WithValue(ctx, logfKey{}, fn)
+}
+
 // WithPrefix wraps f, prefixing each format with the provided prefix.
 func WithPrefix(f Logf, prefix string) Logf {
 	return func(format string, args ...interface{}) {
diff --git a/types/logger/logger_test.go b/types/logger/logger_test.go
index fc9502105..14945105f 100644
--- a/types/logger/logger_test.go
+++ b/types/logger/logger_test.go
@@ -7,11 +7,14 @@
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"fmt"
 	"log"
 	"sync"
 	"testing"
 	"time"
+
+	qt "github.com/frankban/quicktest"
 )
 
 func TestFuncWriter(t *testing.T) {
@@ -183,3 +186,28 @@ func TestRateLimitedFnReentrancy(t *testing.T) {
 		rlogf("boom") // this used to deadlock
 	}))
 }
+
+func TestContext(t *testing.T) {
+	c := qt.New(t)
+	ctx := context.Background()
+
+	// Test that FromContext returns log.Printf when the context has no custom log function.
+	defer log.SetOutput(log.Writer())
+	defer log.SetFlags(log.Flags())
+	var buf bytes.Buffer
+	log.SetOutput(&buf)
+	log.SetFlags(0)
+	logf := FromContext(ctx)
+	logf("a")
+	c.Assert(buf.String(), qt.Equals, "a\n")
+
+	// Test that FromContext and Ctx work together.
+	var called bool
+	markCalled := func(string, ...interface{}) {
+		called = true
+	}
+	ctx = Ctx(ctx, markCalled)
+	logf = FromContext(ctx)
+	logf("a")
+	c.Assert(called, qt.IsTrue)
+}