Codebase list golang-github-go-kit-kit / a15f0c3
transport/http: add ServerFinalizer Peter Bourgon 7 years ago
2 changed file(s) with 65 addition(s) and 1 deletion(s). Raw diff Collapse all Expand all
+34
-1
transport/http/server.go less more
1717 before []RequestFunc
1818 after []ServerResponseFunc
1919 errorEncoder ErrorEncoder
20 finalizer ServerFinalizerFunc
2021 logger log.Logger
2122 }
2223
6869 }
6970
7071 // ServerErrorLogger is used to log non-terminal errors. By default, no errors
71 // are logged.
72 // are logged. This is intended as a diagnostic measure. Finer-grained control
73 // of error handling, including logging in more detail, should be performed in a
74 // custom ServerErrorEncoder or ServerFinalizer, both of which have access to
75 // the context.
7276 func ServerErrorLogger(logger log.Logger) ServerOption {
7377 return func(s *Server) { s.logger = logger }
78 }
79
80 // ServerFinalizer is executed at the end of every HTTP request.
81 // By default, no finalizer is registered.
82 func ServerFinalizer(f ServerFinalizerFunc) ServerOption {
83 return func(s *Server) { s.finalizer = f }
7484 }
7585
7686 // ServeHTTP implements http.Handler.
7787 func (s Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
7888 ctx := s.ctx
89
90 if s.finalizer != nil {
91 iw := &interceptingWriter{w, http.StatusOK}
92 defer func() { s.finalizer(ctx, iw.code, r) }()
93 w = iw
94 }
7995
8096 for _, f := range s.before {
8197 ctx = f(ctx, r)
115131 // for their own error types. See the example shipping/handling service.
116132 type ErrorEncoder func(ctx context.Context, err error, w http.ResponseWriter)
117133
134 // ServerFinalizerFunc can be used to perform work at the end of an HTTP
135 // request, after the response has been written to the client. The principal
136 // intended use is for request logging.
137 type ServerFinalizerFunc func(ctx context.Context, code int, r *http.Request)
138
118139 func defaultErrorEncoder(_ context.Context, err error, w http.ResponseWriter) {
119140 switch e := err.(type) {
120141 case Error:
130151 http.Error(w, err.Error(), http.StatusInternalServerError)
131152 }
132153 }
154
155 type interceptingWriter struct {
156 http.ResponseWriter
157 code int
158 }
159
160 // WriteHeader may not be explicitly called, so care must be taken to
161 // initialize w.code to its default value of http.StatusOK.
162 func (w *interceptingWriter) WriteHeader(code int) {
163 w.code = code
164 w.ResponseWriter.WriteHeader(code)
165 }
+31
-0
transport/http/server_test.go less more
88
99 "golang.org/x/net/context"
1010
11 "github.com/go-kit/kit/endpoint"
1112 httptransport "github.com/go-kit/kit/transport/http"
1213 )
1314
9091 }
9192 }
9293
94 func TestServerFinalizer(t *testing.T) {
95 c := make(chan int)
96 handler := httptransport.NewServer(
97 context.Background(),
98 endpoint.Nop,
99 func(context.Context, *http.Request) (interface{}, error) {
100 return struct{}{}, nil
101 },
102 func(_ context.Context, w http.ResponseWriter, _ interface{}) error {
103 w.WriteHeader(<-c)
104 return nil
105 },
106 httptransport.ServerFinalizer(func(_ context.Context, code int, _ *http.Request) {
107 c <- code
108 }),
109 )
110
111 server := httptest.NewServer(handler)
112 defer server.Close()
113 go http.Get(server.URL)
114
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)
121 }
122 }
123
93124 func testServer(t *testing.T) (cancel, step func(), resp <-chan *http.Response) {
94125 var (
95126 ctx, cancelfn = context.WithCancel(context.Background())