set request id in the context

middleware/logger.go

@@ -1,11 +1,21 @@
 package middleware
 
 import (
+	"context"
+	"fmt"
 	"log"
 	"net/http"
 	"time"
 )
 
+// ContextKey avoids collision in Context.
+// Read more on https://blog.golang.org/context#TOC_3.2.
+type ContextKey string
+
+func (c ContextKey) String() string {
+	return string(c)
+}
+
 // RequestIDHeader is the name of the HTTP Header which contains the request id.
 // Exported so that it can be changed by developers
 var RequestIDHeader = "X-Request-Id"
@@ -38,17 +48,27 @@ func getReqID(r *http.Request) string {
 // TODO: logger might set the common log object in the ctx
 func Logger(next http.Handler) http.Handler {
 	fn := func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+
+		reqID := r.Header.Get(RequestIDHeader)
+		if reqID == "" {
+			reqID = getReqID(r)
+			if reqID == "" {
+				reqID = fmt.Sprintf("%d", time.Now().UnixNano())
+				ctx = context.WithValue(ctx, ContextKey(RequestIDHeader), reqID)
+			}
+		}
 		start := time.Now()
 		rw := &responseWriter{
 			ResponseWriter: w,
 		}
 
 		// call the next handler
-		next.ServeHTTP(rw, r)
+		next.ServeHTTP(rw, r.WithContext(ctx))
 
 		log.Printf(
 			"[%s] %s [%d] %s?%s %v",
-			getReqID(r), r.Method, rw.code, r.URL.Path, r.URL.RawQuery, time.Since(start),
+			reqID, r.Method, rw.code, r.URL.Path, r.URL.RawQuery, time.Since(start),
 		)
 	}