From f94ac060a14c60a5bb8e11001e8c70e9e805d2de Mon Sep 17 00:00:00 2001
From: JP Robinson <john.paul.robinson@gmail.com>
Date: Mon, 23 Sep 2019 13:52:20 -0400
Subject: [PATCH] [server/kit] rearranging stack so logger is available within
 HTTPMiddleware (#232)

* rearanging stack so logger is available within HTTPMiddleware

* remove space
---
 server/kit/kitserver.go      | 31 +++++++++++++++----------------
 server/kit/kitserver_test.go |  5 ++++-
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/server/kit/kitserver.go b/server/kit/kitserver.go
index de95d0238..3d7486eaf 100644
--- a/server/kit/kitserver.go
+++ b/server/kit/kitserver.go
@@ -186,33 +186,32 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			s.errs.Flush()
 		}
 	}()
-	s.handler.ServeHTTP(w, r)
+
+	// populate context with helpful keys
+	ctx := httptransport.PopulateRequestContext(r.Context(), r)
+
+	// add google trace header to use in tracing and logging
+	ctx = context.WithValue(ctx, ContextKeyCloudTraceContext,
+		r.Header.Get("X-Cloud-Trace-Context"))
+
+	// add a request scoped logger to the context
+	ctx = SetLogger(ctx, AddLogKeyVals(ctx, s.logger))
+
+	s.handler.ServeHTTP(w, r.WithContext(ctx))
 }
 
 func (s *Server) register(svc Service) {
 	s.svc = svc
 	s.handler = s.svc.HTTPMiddleware(s.mux)
 
-	opts := []httptransport.ServerOption{
-		// populate context with helpful keys
-		httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context {
-			ctx = httptransport.PopulateRequestContext(ctx, r)
-			// add google trace header to use in tracing and logging
-			return context.WithValue(ctx, ContextKeyCloudTraceContext,
-				r.Header.Get("X-Cloud-Trace-Context"))
-		}),
-		// inject the server logger into every request context
-		httptransport.ServerBefore(func(ctx context.Context, _ *http.Request) context.Context {
-			return SetLogger(ctx, AddLogKeyVals(ctx, s.logger))
-		}),
-	}
-	opts = append(opts, svc.HTTPOptions()...)
-
 	const warmupPath = "/_ah/warmup"
 	var (
 		healthzFound bool
 		warmupFound  bool
 	)
+
+	opts := svc.HTTPOptions()
+
 	// register all endpoints with our wrappers & default decoders/encoders
 	for path, epMethods := range svc.HTTPEndpoints() {
 		for method, ep := range epMethods {
diff --git a/server/kit/kitserver_test.go b/server/kit/kitserver_test.go
index 117617fcd..0250fd61e 100644
--- a/server/kit/kitserver_test.go
+++ b/server/kit/kitserver_test.go
@@ -165,7 +165,10 @@ func (s *server) Middleware(e endpoint.Endpoint) endpoint.Endpoint {
 }
 
 func (s *server) HTTPMiddleware(h http.Handler) http.Handler {
-	return gserver.CORSHandler(h, "")
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		kit.LogDebug(r.Context(), "logging in the HTTP middleware!")
+		gserver.CORSHandler(h, "").ServeHTTP(w, r)
+	})
 }
 
 func (s *server) HTTPOptions() []httptransport.ServerOption {