Merge pull request #461 from go-kit/issue-460
transport/http: provide more response details
Peter Bourgon authored 7 years ago
GitHub committed 7 years ago
115 | 115 | // ContextKeyRequestXRequestID is populated in the context by |
116 | 116 | // PopulateRequestContext. Its value is r.Header.Get("X-Request-Id"). |
117 | 117 | ContextKeyRequestXRequestID |
118 | ||
119 | // ContextKeyResponseHeaders is populated in the context whenever a | |
120 | // ServerFinalizerFunc is specified. Its value is of type http.Header, and | |
121 | // is captured only once the entire response has been written. | |
122 | ContextKeyResponseHeaders | |
123 | ||
124 | // ContextKeyResponseSize is populated in the context whenever a | |
125 | // ServerFinalizerFunc is specified. Its value is of type int64. | |
126 | ContextKeyResponseSize | |
118 | 127 | ) |
87 | 87 | ctx := s.ctx |
88 | 88 | |
89 | 89 | if s.finalizer != nil { |
90 | iw := &interceptingWriter{w, http.StatusOK} | |
91 | defer func() { s.finalizer(ctx, iw.code, r) }() | |
90 | iw := &interceptingWriter{w, http.StatusOK, 0} | |
91 | defer func() { | |
92 | ctx = context.WithValue(ctx, ContextKeyResponseHeaders, iw.Header()) | |
93 | ctx = context.WithValue(ctx, ContextKeyResponseSize, iw.written) | |
94 | s.finalizer(ctx, iw.code, r) | |
95 | }() | |
92 | 96 | w = iw |
93 | 97 | } |
94 | 98 | |
129 | 133 | |
130 | 134 | // ServerFinalizerFunc can be used to perform work at the end of an HTTP |
131 | 135 | // request, after the response has been written to the client. The principal |
132 | // intended use is for request logging. | |
136 | // intended use is for request logging. In addition to the response code | |
137 | // provided in the function signature, additional response parameters are | |
138 | // provided in the context under keys with the ContextKeyResponse prefix. | |
133 | 139 | type ServerFinalizerFunc func(ctx context.Context, code int, r *http.Request) |
134 | 140 | |
135 | 141 | // EncodeJSONResponse is a EncodeResponseFunc that serializes the response as a |
199 | 205 | |
200 | 206 | type interceptingWriter struct { |
201 | 207 | http.ResponseWriter |
202 | code int | |
208 | code int | |
209 | written int64 | |
203 | 210 | } |
204 | 211 | |
205 | 212 | // WriteHeader may not be explicitly called, so care must be taken to |
208 | 215 | w.code = code |
209 | 216 | w.ResponseWriter.WriteHeader(code) |
210 | 217 | } |
218 | ||
219 | func (w *interceptingWriter) Write(p []byte) (int, error) { | |
220 | n, err := w.ResponseWriter.Write(p) | |
221 | w.written += int64(n) | |
222 | return n, err | |
223 | } |
7 | 7 | "net/http/httptest" |
8 | 8 | "strings" |
9 | 9 | "testing" |
10 | "time" | |
10 | 11 | |
11 | 12 | "github.com/go-kit/kit/endpoint" |
12 | 13 | httptransport "github.com/go-kit/kit/transport/http" |
92 | 93 | } |
93 | 94 | |
94 | 95 | func TestServerFinalizer(t *testing.T) { |
95 | c := make(chan int) | |
96 | var ( | |
97 | headerKey = "X-Henlo-Lizer" | |
98 | headerVal = "Helllo you stinky lizard" | |
99 | statusCode = http.StatusTeapot | |
100 | responseBody = "go eat a fly ugly\n" | |
101 | done = make(chan struct{}) | |
102 | ) | |
96 | 103 | handler := httptransport.NewServer( |
97 | 104 | context.Background(), |
98 | 105 | endpoint.Nop, |
100 | 107 | return struct{}{}, nil |
101 | 108 | }, |
102 | 109 | func(_ context.Context, w http.ResponseWriter, _ interface{}) error { |
103 | w.WriteHeader(<-c) | |
110 | w.Header().Set(headerKey, headerVal) | |
111 | w.WriteHeader(statusCode) | |
112 | w.Write([]byte(responseBody)) | |
104 | 113 | return nil |
105 | 114 | }, |
106 | httptransport.ServerFinalizer(func(_ context.Context, code int, _ *http.Request) { | |
107 | c <- code | |
115 | httptransport.ServerFinalizer(func(ctx context.Context, code int, _ *http.Request) { | |
116 | if want, have := statusCode, code; want != have { | |
117 | t.Errorf("StatusCode: want %d, have %d", want, have) | |
118 | } | |
119 | ||
120 | responseHeader := ctx.Value(httptransport.ContextKeyResponseHeaders).(http.Header) | |
121 | if want, have := headerVal, responseHeader.Get(headerKey); want != have { | |
122 | t.Errorf("%s: want %q, have %q", headerKey, want, have) | |
123 | } | |
124 | ||
125 | responseSize := ctx.Value(httptransport.ContextKeyResponseSize).(int64) | |
126 | if want, have := int64(len(responseBody)), responseSize; want != have { | |
127 | t.Errorf("response size: want %d, have %d", want, have) | |
128 | } | |
129 | ||
130 | close(done) | |
108 | 131 | }), |
109 | 132 | ) |
110 | 133 | |
112 | 135 | defer server.Close() |
113 | 136 | go http.Get(server.URL) |
114 | 137 | |
115 | want := http.StatusTeapot | |
116 | c <- want // give status code to response encoder | |
117 | have := <-c // take status code from finalizer | |
118 | ||
119 | if want != have { | |
120 | t.Errorf("want %d, have %d", want, have) | |
138 | select { | |
139 | case <-done: | |
140 | case <-time.After(time.Second): | |
141 | t.Fatal("timeout waiting for finalizer") | |
121 | 142 | } |
122 | 143 | } |
123 | 144 |