{From,With}Context for storing a logger in a context.Context

WithContext also supports additional arguments to set fields since this
is a very common operation.

Author: mitchellh

context.go

@@ -0,0 +1,33 @@
+package hclog
+
+import (
+	"context"
+)
+
+// WithContext inserts a logger into the context and is retrievable
+// with FromContext. The optional args can be set with the same syntax as
+// Logger.With to set fields on the inserted logger. This will not modify
+// the logger argument in-place.
+func WithContext(ctx context.Context, logger Logger, args ...interface{}) context.Context {
+	// While we could call logger.With even with zero args, we have this
+	// check to avoid unnecessary allocations around creating a copy of a
+	// logger.
+	if len(args) > 0 {
+		logger = logger.With(args...)
+	}
+
+	return context.WithValue(ctx, contextKey, logger)
+}
+
+// FromContext returns a logger from the context. This will return nil
+// if there is no logger in the context.
+func FromContext(ctx context.Context) Logger {
+	logger, _ := ctx.Value(contextKey).(Logger)
+	return logger
+}
+
+// Unexported new type so that our context key never collides with another.
+type contextKeyType struct{}
+
+// contextKey is the key used for the context to store the logger.
+var contextKey = contextKeyType{}

context_test.go

@@ -0,0 +1,32 @@
+package hclog
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestContext_simpleLogger(t *testing.T) {
+	l := L()
+	ctx := WithContext(context.Background(), l)
+	require.Equal(t, l, FromContext(ctx))
+}
+
+func TestContext_fields(t *testing.T) {
+	var buf bytes.Buffer
+	l := New(&LoggerOptions{
+		Level:  Debug,
+		Output: &buf,
+	})
+
+	// Insert the logger with fields
+	ctx := WithContext(context.Background(), l, "hello", "world")
+	l = FromContext(ctx)
+	require.NotNil(t, l)
+
+	// Log something so we can test the output that the field is there
+	l.Debug("test")
+	require.Contains(t, buf.String(), "hello")
+}