Context package to support context-aware logging

The new context package supports context-aware logging, integrating with
logrus. Several utilities are provided to associate http requests with a
context, ensuring that one can trace log messages all the way through a
context-aware call stack.

A full description of this functionality is available in doc.go.

Signed-off-by: Stephen J Day <stephen.day@docker.com>

context/doc.go

@@ -0,0 +1,76 @@
+// Package context provides several utilities for working with
+// golang.org/x/net/context in http requests. Primarily, the focus is on
+// logging relevent request information but this package is not limited to
+// that purpose.
+//
+// Logging
+//
+// The most useful aspect of this package is GetLogger. This function takes
+// any context.Context interface and returns the current logger from the
+// context. Canonical usage looks like this:
+//
+// 	GetLogger(ctx).Infof("something interesting happened")
+//
+// GetLogger also takes optional key arguments. The keys will be looked up in
+// the context and reported with the logger. The following example would
+// return a logger that prints the version with each log message:
+//
+// 	ctx := context.Context(context.Background(), "version", version)
+// 	GetLogger(ctx, "version").Infof("this log message has a version field")
+//
+// The above would print out a log message like this:
+//
+// 	INFO[0000] this log message has a version field        version=v2.0.0-alpha.2.m
+//
+// When used with WithLogger, we gain the ability to decorate the context with
+// loggers that have information from disparate parts of the call stack.
+// Following from the version example, we can build a new context with the
+// configured logger such that we always print the version field:
+//
+// 	ctx = WithLogger(ctx, GetLogger(ctx, "version"))
+//
+// Since the logger has been pushed to the context, we can now get the version
+// field for free with our log messages. Future calls to GetLogger on the new
+// context will have the version field:
+//
+// 	GetLogger(ctx).Infof("this log message has a version field")
+//
+// This becomes more powerful when we start stacking loggers. Let's say we
+// have the version logger from above but also want a request id. Using the
+// context above, in our request scoped function, we place another logger in
+// the context:
+//
+// 	ctx = context.WithValue(ctx, "http.request.id", "unique id") // called when building request context
+// 	ctx = WithLogger(ctx, GetLogger(ctx, "http.request.id"))
+//
+// When GetLogger is called on the new context, "http.request.id" will be
+// included as a logger field, along with the original "version" field:
+//
+// 	INFO[0000] this log message has a version field        http.request.id=unique id version=v2.0.0-alpha.2.m
+//
+// Note that this only affects the new context, the previous context, with the
+// version field, can be used independently. Put another way, the new logger,
+// added to the request context, is unique to that context and can have
+// request scoped varaibles.
+//
+// HTTP Requests
+//
+// This package also contains several methods for working with http requests.
+// The concepts are very similar to those described above. We simply place the
+// request in the context using WithRequest. This makes the request variables
+// available. GetRequestLogger can then be called to get request specific
+// variables in a log line:
+//
+// 	ctx = WithRequest(ctx, req)
+// 	GetRequestLogger(ctx).Infof("request variables")
+//
+// Like above, if we want to include the request data in all log messages in
+// the context, we push the logger to a new context and use that one:
+//
+// 	ctx = WithLogger(ctx, GetRequestLogger(ctx))
+//
+// The concept is fairly powerful and ensures that calls throughout the stack
+// can be traced in log messages. Using the fields like "http.request.id", one
+// can analyze call flow for a particular request with a simple grep of the
+// logs.
+package context

context/http.go

@@ -0,0 +1,270 @@
+package context
+
+import (
+	"errors"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"code.google.com/p/go-uuid/uuid"
+	"github.com/gorilla/mux"
+	"golang.org/x/net/context"
+)
+
+// Common errors used with this package.
+var (
+	ErrNoRequestContext = errors.New("no http request in context")
+)
+
+// WithRequest places the request on the context. The context of the request
+// is assigned a unique id, available at "http.request.id". The request itself
+// is available at "http.request". Other common attributes are available under
+// the prefix "http.request.". If a request is already present on the context,
+// this method will panic.
+func WithRequest(ctx context.Context, r *http.Request) context.Context {
+	if ctx.Value("http.request") != nil {
+		// NOTE(stevvooe): This needs to be considered a programming error. It
+		// is unlikely that we'd want to have more than one request in
+		// context.
+		panic("only one request per context")
+	}
+
+	return &httpRequestContext{
+		Context:   ctx,
+		startedAt: time.Now(),
+		id:        uuid.New(), // assign the request a unique.
+		r:         r,
+	}
+}
+
+// GetRequest returns the http request in the given context. Returns
+// ErrNoRequestContext if the context does not have an http request associated
+// with it.
+func GetRequest(ctx context.Context) (*http.Request, error) {
+	if r, ok := ctx.Value("http.request").(*http.Request); r != nil && ok {
+		return r, nil
+	}
+	return nil, ErrNoRequestContext
+}
+
+// GetRequestID attempts to resolve the current request id, if possible. An
+// error is return if it is not available on the context.
+func GetRequestID(ctx context.Context) string {
+	return GetStringValue(ctx, "http.request.id")
+}
+
+// WithResponseWriter returns a new context and response writer that makes
+// interesting response statistics available within the context.
+func WithResponseWriter(ctx context.Context, w http.ResponseWriter) (context.Context, http.ResponseWriter) {
+	irw := &instrumentedResponseWriter{
+		ResponseWriter: w,
+		Context:        ctx,
+	}
+
+	return irw, irw
+}
+
+// getVarsFromRequest let's us change request vars implementation for testing
+// and maybe future changes.
+var getVarsFromRequest = mux.Vars
+
+// WithVars extracts gorilla/mux vars and makes them available on the returned
+// context. Variables are available at keys with the prefix "vars.". For
+// example, if looking for the variable "name", it can be accessed as
+// "vars.name". Implementations that are accessing values need not know that
+// the underlying context is implemented with gorilla/mux vars.
+func WithVars(ctx context.Context, r *http.Request) context.Context {
+	return &muxVarsContext{
+		Context: ctx,
+		vars:    getVarsFromRequest(r),
+	}
+}
+
+// GetRequestLogger returns a logger that contains fields from the request in
+// the current context. If the request is not available in the context, no
+// fields will display. Request loggers can safely be pushed onto the context.
+func GetRequestLogger(ctx context.Context) Logger {
+	return GetLogger(ctx,
+		"http.request.id",
+		"http.request.method",
+		"http.request.host",
+		"http.request.uri",
+		"http.request.referer",
+		"http.request.useragent",
+		"http.request.remoteaddr",
+		"http.request.contenttype")
+}
+
+// GetResponseLogger reads the current response stats and builds a logger.
+// Because the values are read at call time, pushing a logger returned from
+// this function on the context will lead to missing or invalid data. Only
+// call this at the end of a request, after the response has been written.
+func GetResponseLogger(ctx context.Context) Logger {
+	l := getLogrusLogger(ctx,
+		"http.response.written",
+		"http.response.status",
+		"http.response.contenttype")
+
+	duration := Since(ctx, "http.request.startedat")
+
+	if duration > 0 {
+		l = l.WithField("http.response.duration", duration)
+	}
+
+	return l
+}
+
+// httpRequestContext makes information about a request available to context.
+type httpRequestContext struct {
+	context.Context
+
+	startedAt time.Time
+	id        string
+	r         *http.Request
+}
+
+// Value returns a keyed element of the request for use in the context. To get
+// the request itself, query "request". For other components, access them as
+// "request.<component>". For example, r.RequestURI
+func (ctx *httpRequestContext) Value(key interface{}) interface{} {
+	if keyStr, ok := key.(string); ok {
+		if keyStr == "http.request" {
+			return ctx.r
+		}
+
+		parts := strings.Split(keyStr, ".")
+
+		if len(parts) != 3 {
+			goto fallback
+		}
+
+		switch parts[2] {
+		case "uri":
+			return ctx.r.RequestURI
+		case "remoteaddr":
+			return ctx.r.RemoteAddr
+		case "method":
+			return ctx.r.Method
+		case "host":
+			return ctx.r.Host
+		case "referer":
+			referer := ctx.r.Referer()
+			if referer != "" {
+				return referer
+			}
+		case "useragent":
+			return ctx.r.UserAgent()
+		case "id":
+			return ctx.id
+		case "startedat":
+			return ctx.startedAt
+		case "contenttype":
+			ct := ctx.r.Header.Get("Content-Type")
+			if ct != "" {
+				return ct
+			}
+		}
+	}
+
+fallback:
+	return ctx.Context.Value(key)
+}
+
+type muxVarsContext struct {
+	context.Context
+	vars map[string]string
+}
+
+func (ctx *muxVarsContext) Value(key interface{}) interface{} {
+	if keyStr, ok := key.(string); ok {
+		if keyStr == "vars" {
+			return ctx.vars
+		}
+
+		if strings.HasPrefix(keyStr, "vars.") {
+			keyStr = strings.TrimPrefix(keyStr, "vars.")
+		}
+
+		if v, ok := ctx.vars[keyStr]; ok {
+			return v
+		}
+	}
+
+	return ctx.Context.Value(key)
+}
+
+// instrumentedResponseWriter provides response writer information in a
+// context.
+type instrumentedResponseWriter struct {
+	http.ResponseWriter
+	context.Context
+
+	mu      sync.Mutex
+	status  int
+	written int64
+}
+
+func (irw *instrumentedResponseWriter) Write(p []byte) (n int, err error) {
+	n, err = irw.ResponseWriter.Write(p)
+
+	irw.mu.Lock()
+	irw.written += int64(n)
+
+	// Guess the likely status if not set.
+	if irw.status == 0 {
+		irw.status = http.StatusOK
+	}
+
+	irw.mu.Unlock()
+
+	return
+}
+
+func (irw *instrumentedResponseWriter) WriteHeader(status int) {
+	irw.ResponseWriter.WriteHeader(status)
+
+	irw.mu.Lock()
+	irw.status = status
+	irw.mu.Unlock()
+}
+
+func (irw *instrumentedResponseWriter) Flush() {
+	if flusher, ok := irw.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+func (irw *instrumentedResponseWriter) Value(key interface{}) interface{} {
+	if keyStr, ok := key.(string); ok {
+		if keyStr == "http.response" {
+			return irw.ResponseWriter
+		}
+
+		parts := strings.Split(keyStr, ".")
+
+		if len(parts) != 3 {
+			goto fallback
+		}
+
+		irw.mu.Lock()
+		defer irw.mu.Unlock()
+
+		switch parts[2] {
+		case "written":
+			return irw.written
+		case "status":
+			if irw.status != 0 {
+				return irw.status
+			}
+		case "contenttype":
+			contentType := irw.Header().Get("Content-Type")
+			if contentType != "" {
+				return contentType
+			}
+		}
+	}
+
+fallback:
+	return irw.Context.Value(key)
+}

context/http_test.go

@@ -0,0 +1,207 @@
+package context
+
+import (
+	"net/http"
+	"reflect"
+	"testing"
+	"time"
+
+	"golang.org/x/net/context"
+)
+
+func TestWithRequest(t *testing.T) {
+	var req http.Request
+
+	start := time.Now()
+	req.Method = "GET"
+	req.Host = "example.com"
+	req.RequestURI = "/test-test"
+	req.Header = make(http.Header)
+	req.Header.Set("Referer", "foo.com/referer")
+	req.Header.Set("User-Agent", "test/0.1")
+
+	ctx := WithRequest(context.Background(), &req)
+	for _, testcase := range []struct {
+		key      string
+		expected interface{}
+	}{
+		{
+			key:      "http.request",
+			expected: &req,
+		},
+		{
+			key: "http.request.id",
+		},
+		{
+			key:      "http.request.method",
+			expected: req.Method,
+		},
+		{
+			key:      "http.request.host",
+			expected: req.Host,
+		},
+		{
+			key:      "http.request.uri",
+			expected: req.RequestURI,
+		},
+		{
+			key:      "http.request.referer",
+			expected: req.Referer(),
+		},
+		{
+			key:      "http.request.useragent",
+			expected: req.UserAgent(),
+		},
+		{
+			key:      "http.request.remoteaddr",
+			expected: req.RemoteAddr,
+		},
+		{
+			key: "http.request.startedat",
+		},
+	} {
+		v := ctx.Value(testcase.key)
+
+		if v == nil {
+			t.Fatalf("value not found for %q", testcase.key)
+		}
+
+		if testcase.expected != nil && v != testcase.expected {
+			t.Fatalf("%s: %v != %v", testcase.key, v, testcase.expected)
+		}
+
+		// Key specific checks!
+		switch testcase.key {
+		case "http.request.id":
+			if _, ok := v.(string); !ok {
+				t.Fatalf("request id not a string: %v", v)
+			}
+		case "http.request.startedat":
+			vt, ok := v.(time.Time)
+			if !ok {
+				t.Fatalf("value not a time: %v", v)
+			}
+
+			now := time.Now()
+			if vt.After(now) {
+				t.Fatalf("time generated too late: %v > %v", vt, now)
+			}
+
+			if vt.Before(start) {
+				t.Fatalf("time generated too early: %v < %v", vt, start)
+			}
+		}
+	}
+}
+
+type testResponseWriter struct {
+	flushed bool
+	status  int
+	written int64
+	header  http.Header
+}
+
+func (trw *testResponseWriter) Header() http.Header {
+	if trw.header == nil {
+		trw.header = make(http.Header)
+	}
+
+	return trw.header
+}
+
+func (trw *testResponseWriter) Write(p []byte) (n int, err error) {
+	if trw.status == 0 {
+		trw.status = http.StatusOK
+	}
+
+	n = len(p)
+	trw.written += int64(n)
+	return
+}
+
+func (trw *testResponseWriter) WriteHeader(status int) {
+	trw.status = status
+}
+
+func (trw *testResponseWriter) Flush() {
+	trw.flushed = true
+}
+
+func TestWithResponseWriter(t *testing.T) {
+	trw := testResponseWriter{}
+	ctx, rw := WithResponseWriter(context.Background(), &trw)
+
+	if ctx.Value("http.response") != &trw {
+		t.Fatalf("response not available in context: %v != %v", ctx.Value("http.response"), &trw)
+	}
+
+	if n, err := rw.Write(make([]byte, 1024)); err != nil {
+		t.Fatalf("unexpected error writing: %v", err)
+	} else if n != 1024 {
+		t.Fatalf("unexpected number of bytes written: %v != %v", n, 1024)
+	}
+
+	if ctx.Value("http.response.status") != http.StatusOK {
+		t.Fatalf("unexpected response status in context: %v != %v", ctx.Value("http.response.status"), http.StatusOK)
+	}
+
+	if ctx.Value("http.response.written") != int64(1024) {
+		t.Fatalf("unexpected number reported bytes written: %v != %v", ctx.Value("http.response.written"), 1024)
+	}
+
+	// Make sure flush propagates
+	rw.(http.Flusher).Flush()
+
+	if !trw.flushed {
+		t.Fatalf("response writer not flushed")
+	}
+
+	// Write another status and make sure context is correct. This normally
+	// wouldn't work except for in this contrived testcase.
+	rw.WriteHeader(http.StatusBadRequest)
+
+	if ctx.Value("http.response.status") != http.StatusBadRequest {
+		t.Fatalf("unexpected response status in context: %v != %v", ctx.Value("http.response.status"), http.StatusBadRequest)
+	}
+}
+
+func TestWithVars(t *testing.T) {
+	var req http.Request
+	vars := map[string]string{
+		"foo": "asdf",
+		"bar": "qwer",
+	}
+
+	getVarsFromRequest = func(r *http.Request) map[string]string {
+		if r != &req {
+			t.Fatalf("unexpected request: %v != %v", r, req)
+		}
+
+		return vars
+	}
+
+	ctx := WithVars(context.Background(), &req)
+	for _, testcase := range []struct {
+		key      string
+		expected interface{}
+	}{
+		{
+			key:      "vars",
+			expected: vars,
+		},
+		{
+			key:      "vars.foo",
+			expected: "asdf",
+		},
+		{
+			key:      "vars.bar",
+			expected: "qwer",
+		},
+	} {
+		v := ctx.Value(testcase.key)
+
+		if !reflect.DeepEqual(v, testcase.expected) {
+			t.Fatalf("%q: %v != %v", testcase.key, v, testcase.expected)
+		}
+	}
+}

context/logger.go

@@ -0,0 +1,88 @@
+package context
+
+import (
+	"fmt"
+
+	"github.com/Sirupsen/logrus"
+	"golang.org/x/net/context"
+)
+
+// Logger provides a leveled-logging interface.
+type Logger interface {
+	// standard logger methods
+	Print(args ...interface{})
+	Printf(format string, args ...interface{})
+	Println(args ...interface{})
+
+	Fatal(args ...interface{})
+	Fatalf(format string, args ...interface{})
+	Fatalln(args ...interface{})
+
+	Panic(args ...interface{})
+	Panicf(format string, args ...interface{})
+	Panicln(args ...interface{})
+
+	// Leveled methods, from logrus
+	Debug(args ...interface{})
+	Debugf(format string, args ...interface{})
+	Debugln(args ...interface{})
+
+	Error(args ...interface{})
+	Errorf(format string, args ...interface{})
+	Errorln(args ...interface{})
+
+	Info(args ...interface{})
+	Infof(format string, args ...interface{})
+	Infoln(args ...interface{})
+
+	Warn(args ...interface{})
+	Warnf(format string, args ...interface{})
+	Warnln(args ...interface{})
+}
+
+// WithLogger creates a new context with provided logger.
+func WithLogger(ctx context.Context, logger Logger) context.Context {
+	return context.WithValue(ctx, "logger", logger)
+}
+
+// GetLogger returns the logger from the current context, if present. If one
+// or more keys are provided, they will be resolved on the context and
+// included in the logger. While context.Value takes an interface, any key
+// argument passed to GetLogger will be passed to fmt.Sprint when expanded as
+// a logging key field. If context keys are integer constants, for example,
+// its recommended that a String method is implemented.
+func GetLogger(ctx context.Context, keys ...interface{}) Logger {
+	return getLogrusLogger(ctx, keys...)
+}
+
+// GetLogrusLogger returns the logrus logger for the context. If one more keys
+// are provided, they will be resolved on the context and included in the
+// logger. Only use this function if specific logrus functionality is
+// required.
+func getLogrusLogger(ctx context.Context, keys ...interface{}) *logrus.Entry {
+	var logger *logrus.Entry
+
+	// Get a logger, if it is present.
+	loggerInterface := ctx.Value("logger")
+	if loggerInterface != nil {
+		if lgr, ok := loggerInterface.(*logrus.Entry); ok {
+			logger = lgr
+		}
+	}
+
+	if logger == nil {
+		// If no logger is found, just return the standard logger.
+		logger = logrus.NewEntry(logrus.StandardLogger())
+	}
+
+	fields := logrus.Fields{}
+
+	for _, key := range keys {
+		v := ctx.Value(key)
+		if v != nil {
+			fields[fmt.Sprint(key)] = v
+		}
+	}
+
+	return logger.WithFields(fields)
+}

context/util.go

@@ -0,0 +1,34 @@
+package context
+
+import (
+	"time"
+
+	"golang.org/x/net/context"
+)
+
+// Since looks up key, which should be a time.Time, and returns the duration
+// since that time. If the key is not found, the value returned will be zero.
+// This is helpful when inferring metrics related to context execution times.
+func Since(ctx context.Context, key interface{}) time.Duration {
+	startedAtI := ctx.Value(key)
+	if startedAtI != nil {
+		if startedAt, ok := startedAtI.(time.Time); ok {
+			return time.Since(startedAt)
+		}
+	}
+
+	return 0
+}
+
+// GetStringValue returns a string value from the context. The empty string
+// will be returned if not found.
+func GetStringValue(ctx context.Context, key string) (value string) {
+	stringi := ctx.Value(key)
+	if stringi != nil {
+		if valuev, ok := stringi.(string); ok {
+			value = valuev
+		}
+	}
+
+	return value
+}