Package list golang-github-go-kit-kit / 816c73e
transport/http: provide more response details ServerFinalizerFunc gets access to response headers and size. Closes #460. Peter Bourgon 4 years ago
3 changed file(s) with 57 addition(s) and 14 deletion(s). Raw diff Collapse all Expand all
115115 // ContextKeyRequestXRequestID is populated in the context by
116116 // PopulateRequestContext. Its value is r.Header.Get("X-Request-Id").
117117 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
118127 )
8787 ctx := s.ctx
8888
8989 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 }()
9296 w = iw
9397 }
9498
129133
130134 // ServerFinalizerFunc can be used to perform work at the end of an HTTP
131135 // 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.
133139 type ServerFinalizerFunc func(ctx context.Context, code int, r *http.Request)
134140
135141 // EncodeJSONResponse is a EncodeResponseFunc that serializes the response as a
199205
200206 type interceptingWriter struct {
201207 http.ResponseWriter
202 code int
208 code int
209 written int64
203210 }
204211
205212 // WriteHeader may not be explicitly called, so care must be taken to
208215 w.code = code
209216 w.ResponseWriter.WriteHeader(code)
210217 }
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 }
77 "net/http/httptest"
88 "strings"
99 "testing"
10 "time"
1011
1112 "github.com/go-kit/kit/endpoint"
1213 httptransport "github.com/go-kit/kit/transport/http"
9293 }
9394
9495 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 )
96103 handler := httptransport.NewServer(
97104 context.Background(),
98105 endpoint.Nop,
100107 return struct{}{}, nil
101108 },
102109 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))
104113 return nil
105114 },
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)
108131 }),
109132 )
110133
112135 defer server.Close()
113136 go http.Get(server.URL)
114137
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")
121142 }
122143 }
123144