refactor to have object that could be used

Author: rybit

tracing/req_tracer.go

@@ -0,0 +1,67 @@
+package tracing
+
+import (
+	"net/http"
+	"time"
+
+	opentracing "github.com/opentracing/opentracing-go"
+	"github.com/sirupsen/logrus"
+)
+
+type RequestTracer struct {
+	*trackingWriter
+	RequestID string
+
+	request *http.Request
+	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)
+	r, span := WrapWithSpan(r, reqID, service)
+	trackWriter := &trackingWriter{
+		writer: w,
+		log:    log,
+	}
+
+	rt := &RequestTracer{
+		RequestID:      reqID,
+		span:           span,
+		trackingWriter: trackWriter,
+		request:        r,
+	}
+
+	return rt, r, rt
+}
+
+func (rt *RequestTracer) Start() {
+	rt.start = time.Now()
+	rt.Log().WithField("url", rt.request.URL.String()).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(),
+	}
+	rt.span.Finish()
+	rt.Log().WithFields(fields).Info("Completed Request")
+}
+
+func (rt *RequestTracer) Log() logrus.FieldLogger {
+	return GetLogger(rt.request)
+}
+
+func (rt *RequestTracer) SetLogField(key string, value interface{}) logrus.FieldLogger {
+	return SetLogField(rt.request, key, value)
+}
+
+func (rt *RequestTracer) SetLogFields(fields logrus.Fields) logrus.FieldLogger {
+	return SetLogFields(rt.request, fields)
+}

tracing/tracer.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"net/http"
 	"strconv"
-	"time"
 
 	opentracing "github.com/opentracing/opentracing-go"
 	ext "github.com/opentracing/opentracing-go/ext"
@@ -21,29 +20,10 @@ const (
 )
 
 func TrackRequest(w http.ResponseWriter, r *http.Request, log logrus.FieldLogger, service string, next http.Handler) {
-	reqID := RequestID(r)
-	r, log = WrapWithLogger(r, reqID, log)
-	log.WithField("url", r.URL.String()).Info("Starting Request")
-
-	r, span := WrapWithSpan(r, reqID, service)
-
-	trackWriter := &trackingWriter{
-		writer: w,
-		log:    log,
-	}
-
-	start := time.Now()
-	next.ServeHTTP(trackWriter, r)
-	dur := time.Since(start)
-
-	log = GetLogger(r) // in case people have added fields
-	log.WithFields(logrus.Fields{
-		"status_code": trackWriter.status,
-		"rsp_bytes":   trackWriter.rspBytes,
-		"dur":         dur.String(),
-		"dur_ns":      dur.Nanoseconds(),
-	}).Info("Completed Request")
-	span.Finish()
+	w, r, rt := NewTracer(w, r, log, service)
+	rt.Start()
+	next.ServeHTTP(w, r)
+	rt.Finish()
 }
 
 func RequestID(r *http.Request) string {

tracing/writer.go

@@ -17,6 +17,10 @@ type trackingWriter struct {
 	log      logrus.FieldLogger
 }
 
+func (w *trackingWriter) OriginalWriter() http.ResponseWriter {
+	return w.writer
+}
+
 func (w *trackingWriter) Write(in []byte) (int, error) {
 	w.rspBytes += len(in)
 	return w.writer.Write(in)