simplify logger to use request tracer

Author: rybit

tracing/config.go

@@ -0,0 +1,39 @@
+package tracing
+
+import (
+	"fmt"
+
+	"github.com/opentracing/opentracing-go"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/opentracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+)
+
+const (
+	HeaderNFDebugLogging = "X-NF-Debug-Logging"
+	HeaderRequestUUID    = "X-BB-CLIENT-REQUEST-UUID"
+)
+
+type Config struct {
+	Enabled bool   `default:"false"`
+	Host    string `default:"localhost"`
+	Port    string `default:"8126"`
+	Tags    map[string]string
+}
+
+func Configure(tc *Config, svcName string) {
+	var t opentracing.Tracer = opentracing.NoopTracer{}
+	if tc.Enabled {
+		tracerAddr := fmt.Sprintf("%s:%s", tc.Host, tc.Port)
+		tracerOps := []tracer.StartOption{
+			tracer.WithServiceName(svcName),
+			tracer.WithAgentAddr(tracerAddr),
+		}
+
+		for k, v := range tc.Tags {
+			tracerOps = append(tracerOps, tracer.WithGlobalTag(k, v))
+		}
+
+		t = opentracer.New(tracerOps...)
+	}
+	opentracing.SetGlobalTracer(t)
+}

tracing/context.go

@@ -0,0 +1,26 @@
+package tracing
+
+import (
+	"context"
+	"net/http"
+)
+
+type contextKey string
+
+const tracerKey = contextKey("nf-tracer-key")
+
+func WrapWithTracer(r *http.Request, rt *RequestTracer) *http.Request {
+	return r.WithContext(context.WithValue(r.Context(), tracerKey, rt))
+}
+
+func GetTracer(r *http.Request) *RequestTracer {
+	val := r.Context().Value(tracerKey)
+	if val == nil {
+		return nil
+	}
+	entry, ok := val.(*RequestTracer)
+	if ok {
+		return entry
+	}
+	return nil
+}

tracing/logging.go

@@ -8,50 +8,61 @@ import (
 
 const (
 	logKey = contextKey("nf-log-key")
-
-	HeaderNFDebugLogging = "X-NF-Debug-Logging"
 )
 
-type structuredLoggerEntry struct {
-	Logger logrus.FieldLogger
-}
+func requestLogger(r *http.Request, log logrus.FieldLogger) logrus.FieldLogger {
+	if r.Header.Get(HeaderNFDebugLogging) != "" {
+		logger := logrus.New()
+		logger.SetLevel(logrus.DebugLevel)
 
-func getEntry(r *http.Request) *structuredLoggerEntry {
-	val := r.Context().Value(logKey)
-	if val == nil {
-		return nil
-	}
-	entry, ok := val.(*structuredLoggerEntry)
-	if ok {
-		return entry
+		if entry, ok := log.(*logrus.Entry); ok {
+			log = logger.WithFields(entry.Data)
+		}
 	}
-	return nil
+
+	reqID := RequestID(r)
+
+	log = log.WithFields(logrus.Fields{
+		"request_id": reqID,
+	})
+	return log
 }
 
 func GetLogger(r *http.Request) logrus.FieldLogger {
-	entry := getEntry(r)
+	entry := GetTracer(r)
 	if entry == nil {
 		return logrus.NewEntry(logrus.StandardLogger())
 	}
-	return entry.Logger
+	return entry.FieldLogger
 }
 
+// SetLogField will add the field to this log line and every one following
 func SetLogField(r *http.Request, key string, value interface{}) logrus.FieldLogger {
-	entry := getEntry(r)
+	entry := GetTracer(r)
 	if entry == nil {
 		return logrus.StandardLogger().WithField(key, value)
 	}
-
-	entry.Logger = entry.Logger.WithField(key, value)
-	return entry.Logger
+	return entry.SetLogField(key, value)
 }
 
+// SetLogFields will add the fields to this log line and every one following
 func SetLogFields(r *http.Request, fields logrus.Fields) logrus.FieldLogger {
-	entry := getEntry(r)
+	entry := GetTracer(r)
 	if entry == nil {
 		return logrus.StandardLogger().WithFields(fields)
 	}
 
-	entry.Logger = entry.Logger.WithFields(fields)
-	return entry.Logger
+	return entry.SetLogFields(fields)
+}
+
+// SetFinalField will add a field to the canonical line created at in Finish. It will add
+// it to this line, but not every log line in between
+func SetFinalField(r *http.Request, key string, value interface{}) logrus.FieldLogger {
+	entry := GetTracer(r)
+	if entry == nil {
+		return logrus.StandardLogger().WithField(key, value)
+	}
+
+	entry.FinalFields[key] = value
+	return entry.FieldLogger.WithField(key, value)
 }

tracing/req_tracer.go

@@ -10,58 +10,80 @@ import (
 
 type RequestTracer struct {
 	*trackingWriter
-	RequestID string
+	logrus.FieldLogger
 
-	request *http.Request
-	span    opentracing.Span
-	start   time.Time
+	RequestID   string
+	FinalFields map[string]interface{}
+
+	remoteAddr  string
+	method      string
+	originalURL string
+	referrer    string
+	span        opentracing.Span
+	start       time.Time
 }
 
 func NewTracer(w http.ResponseWriter, r *http.Request, log logrus.FieldLogger, service string) (http.ResponseWriter, *http.Request, *RequestTracer) {
 	reqID := RequestID(r)
-
-	r, log = WrapWithLogger(r, reqID, log)
+	log = requestLogger(r, log)
 	r, span := WrapWithSpan(r, reqID, service)
 	trackWriter := &trackingWriter{
 		writer: w,
 		log:    log,
 	}
 
 	rt := &RequestTracer{
+		originalURL: r.URL.String(),
+		method:      r.Method,
+		referrer:    r.Referer(),
+		remoteAddr:  r.RemoteAddr,
+
 		RequestID:      reqID,
 		span:           span,
 		trackingWriter: trackWriter,
-		request:        r,
+		FieldLogger:    log,
+		FinalFields:    make(map[string]interface{}),
 	}
+	r = WrapWithTracer(r, rt)
 
 	return rt, r, rt
 }
 
 func (rt *RequestTracer) Start() {
 	rt.start = time.Now()
-	rt.Log().WithField("url", rt.request.URL.String()).Info("Starting Request")
+	rt.WithFields(logrus.Fields{
+		"method":      rt.method,
+		"remote_addr": rt.remoteAddr,
+		"referer":     rt.referrer,
+		"url":         rt.originalURL,
+	}).Info("Starting Request")
 }
 
 func (rt *RequestTracer) Finish() {
 	dur := time.Since(rt.start)
-	fields := logrus.Fields{
-		"status_code": rt.trackingWriter.status,
-		"rsp_bytes":   rt.trackingWriter.rspBytes,
-		"dur":         dur.String(),
-		"dur_ns":      dur.Nanoseconds(),
+
+	fields := logrus.Fields{}
+	for k, v := range rt.FinalFields {
+		fields[k] = v
 	}
-	rt.span.Finish()
-	rt.Log().WithFields(fields).Info("Completed Request")
-}
 
-func (rt *RequestTracer) Log() logrus.FieldLogger {
-	return GetLogger(rt.request)
+	fields["status_code"] = rt.trackingWriter.status
+	fields["rsp_bytes"] = rt.trackingWriter.rspBytes
+	fields["url"] = rt.originalURL
+	fields["method"] = rt.method
+	fields["dur"] = dur.String()
+	fields["dur_ns"] = dur.Nanoseconds()
+
+	rt.span.Finish()
+	rt.WithFields(fields).Info("Completed Request")
 }
 
 func (rt *RequestTracer) SetLogField(key string, value interface{}) logrus.FieldLogger {
-	return SetLogField(rt.request, key, value)
+	rt.FieldLogger = rt.FieldLogger.WithField(key, value)
+	return rt.FieldLogger
 }
 
 func (rt *RequestTracer) SetLogFields(fields logrus.Fields) logrus.FieldLogger {
-	return SetLogFields(rt.request, fields)
+	rt.FieldLogger = rt.FieldLogger.WithFields(fields)
+	return rt.FieldLogger
 }

tracing/req_tracer_test.go

@@ -0,0 +1,67 @@
+package tracing
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	logtest "github.com/sirupsen/logrus/hooks/test"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTracerLogging(t *testing.T) {
+	// func NewTracer(w http.ResponseWriter, r *http.Request, log logrus.FieldLogger, service string) (http.ResponseWriter, *http.Request, *RequestTracer) {
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "http://whatever.com/something", nil)
+
+	log, hook := logtest.NewNullLogger()
+
+	_, r, rt := NewTracer(rec, req, log, t.Name())
+
+	rt.Start()
+	e := hook.LastEntry()
+	assert.Equal(t, 5, len(e.Data))
+	assert.NotEmpty(t, e.Data["request_id"])
+	assert.NotEmpty(t, e.Data["remote_addr"])
+	assert.Empty(t, e.Data["referrer"])
+	assert.NotEmpty(t, e.Data["method"])
+	assert.Equal(t, "http://whatever.com/something", e.Data["url"])
+	fmt.Println(e.Data)
+
+	_ = SetLogField(r, "first", "second")
+	SetFinalField(r, "final", "line").Info("should have the final here")
+	e = hook.LastEntry()
+	assert.Equal(t, 3, len(e.Data))
+	assert.NotEmpty(t, e.Data["request_id"])
+	assert.Equal(t, "line", e.Data["final"])
+	assert.Equal(t, "second", e.Data["first"])
+	fmt.Println(e.Data)
+
+	rt.Info("Shouldn't have the final line")
+	e = hook.LastEntry()
+	assert.Equal(t, 2, len(e.Data))
+	assert.NotEmpty(t, e.Data["request_id"])
+	assert.Equal(t, "second", e.Data["first"])
+	fmt.Println(e.Data)
+
+	rt.WriteHeader(http.StatusOK)
+	rt.Write([]byte{0, 1, 2, 3})
+	rt.Finish()
+	e = hook.LastEntry()
+
+	assert.Equal(t, 9, len(e.Data))
+
+	// the automatic fields
+	assert.NotEmpty(t, e.Data["dur"])
+	assert.NotEmpty(t, e.Data["dur_ns"])
+	assert.NotEmpty(t, e.Data["request_id"])
+	assert.Equal(t, 4, e.Data["rsp_bytes"])
+	assert.Equal(t, 200, e.Data["status_code"])
+	assert.Equal(t, "http://whatever.com/something", e.Data["url"])
+	assert.Equal(t, "GET", e.Data["method"])
+
+	// the value that we added above
+	assert.Equal(t, "second", e.Data["first"])
+	assert.Equal(t, "line", e.Data["final"])
+}

tracing/tracer.go

@@ -1,7 +1,6 @@
 package tracing
 
 import (
-	"context"
 	"net/http"
 	"strconv"
 
@@ -13,12 +12,6 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/opentracer"
 )
 
-type contextKey string
-
-const (
-	HeaderRequestUUID = "X-BB-CLIENT-REQUEST-UUID"
-)
-
 func TrackRequest(w http.ResponseWriter, r *http.Request, log logrus.FieldLogger, service string, next http.Handler) {
 	w, r, rt := NewTracer(w, r, log, service)
 	rt.Start()
@@ -35,30 +28,6 @@ func RequestID(r *http.Request) string {
 	return id
 }
 
-func WrapWithLogger(r *http.Request, reqID string, log logrus.FieldLogger) (*http.Request, logrus.FieldLogger) {
-	if r.Header.Get(HeaderNFDebugLogging) != "" {
-		logger := logrus.New()
-		logger.SetLevel(logrus.DebugLevel)
-
-		if entry, ok := log.(*logrus.Entry); ok {
-			log = logger.WithFields(entry.Data)
-		}
-	}
-
-	log = log.WithFields(logrus.Fields{
-		"method":      r.Method,
-		"path":        r.URL.Path,
-		"remote_addr": r.RemoteAddr,
-		"referer":     r.Referer(),
-		"request_id":  reqID,
-	})
-
-	entry := &structuredLoggerEntry{log}
-	r = r.WithContext(context.WithValue(r.Context(), logKey, entry))
-
-	return r, log
-}
-
 func WrapWithSpan(r *http.Request, reqID, service string) (*http.Request, opentracing.Span) {
 	span := opentracing.SpanFromContext(r.Context())
 	if span != nil {

tracing/writer.go

@@ -9,6 +9,8 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+var _ http.Hijacker = (*trackingWriter)(nil)
+
 type trackingWriter struct {
 	writer   http.ResponseWriter
 	rspBytes int