Skip to content

Commit 198211e

Browse files
imduffy15imduffy15
committed
Add delayed chunk endpoint
Adds an endpoint that does chunk based encoding. The endpoint just stalls and eventually returns the stall time. Similar to the delay endpoint but in a chunked maner. Fixed up the metrics interceptor to wrap ResponseWriter correctly too.
1 parent e1ca9e2 commit 198211e

File tree

4 files changed

+346
-11
lines changed

4 files changed

+346
-11
lines changed

‎pkg/api/chunked.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package api
2+
3+
import (
4+
"math/rand"
5+
"net/http"
6+
"strconv"
7+
"time"
8+
9+
"github.com/gorilla/mux"
10+
)
11+
12+
func (s *Server) chunkedHandler(w http.ResponseWriter, r *http.Request) {
13+
vars := mux.Vars(r)
14+
15+
delay, err := strconv.Atoi(vars["wait"])
16+
if err != nil {
17+
delay = rand.Intn(int(s.config.HttpServerTimeout*time.Second)-10) + 10
18+
}
19+
20+
flusher, ok := w.(http.Flusher)
21+
if !ok {
22+
s.ErrorResponse(w, r, "Streaming unsupported!", http.StatusInternalServerError)
23+
return
24+
}
25+
26+
w.Header().Set("Connection", "Keep-Alive")
27+
w.Header().Set("Transfer-Encoding", "chunked")
28+
w.Header().Set("X-Content-Type-Options", "nosniff")
29+
30+
flusher.Flush()
31+
32+
time.Sleep(time.Duration(delay) * time.Second)
33+
s.JSONResponse(w, r, map[string]int{"delay": delay})
34+
35+
flusher.Flush()
36+
}

‎pkg/api/chunked_test.go‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
package api
2+
3+
import (
4+
"net/http"
5+
"net/http/httptest"
6+
"regexp"
7+
"testing"
8+
)
9+
10+
func TestChunkedHandler(t *testing.T) {
11+
req, err := http.NewRequest("GET", "/chunked/0", nil)
12+
if err != nil {
13+
t.Fatal(err)
14+
}
15+
16+
rr := httptest.NewRecorder()
17+
srv := NewMockServer()
18+
19+
srv.router.HandleFunc("/chunked/{wait}", srv.chunkedHandler)
20+
srv.router.ServeHTTP(rr, req)
21+
22+
// Check the status code is what we expect.
23+
if status := rr.Code; status != http.StatusOK {
24+
t.Errorf("handler returned wrong status code: got %v want %v",
25+
status, http.StatusOK)
26+
}
27+
28+
// Check the response body is what we expect.
29+
expected := ".*delay.*0.*"
30+
r := regexp.MustCompile(expected)
31+
if !r.MatchString(rr.Body.String()) {
32+
t.Fatalf("handler returned unexpected body:\ngot \n%v \nwant \n%s",
33+
rr.Body.String(), expected)
34+
}
35+
}

‎pkg/api/metrics.go‎

Lines changed: 267 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package api
33
import (
44
"bufio"
55
"fmt"
6+
"io"
67
"net"
78
"net/http"
89
"regexp"
@@ -51,7 +52,7 @@ func (p *PrometheusMiddleware) Handler(next http.Handler) http.Handler {
5152
begin := time.Now()
5253
interceptor := &interceptor{ResponseWriter: w, statusCode: http.StatusOK}
5354
path := p.getRouteName(r)
54-
next.ServeHTTP(interceptor, r)
55+
next.ServeHTTP(interceptor.wrappedResponseWriter(), r)
5556
var (
5657
status = strconv.Itoa(interceptor.statusCode)
5758
took = time.Since(begin)
@@ -94,6 +95,14 @@ type interceptor struct {
9495
recorded bool
9596
}
9697

98+
func (i *interceptor) Hijack() (net.Conn, *bufio.ReadWriter, error) {
99+
hj, ok := i.ResponseWriter.(http.Hijacker)
100+
if !ok {
101+
return nil, nil, fmt.Errorf("interceptor: can't cast parent ResponseWriter to Hijacker")
102+
}
103+
return hj.Hijack()
104+
}
105+
97106
func (i *interceptor) WriteHeader(code int) {
98107
if !i.recorded {
99108
i.statusCode = code
@@ -102,10 +111,262 @@ func (i *interceptor) WriteHeader(code int) {
102111
i.ResponseWriter.WriteHeader(code)
103112
}
104113

105-
func (i *interceptor) Hijack() (net.Conn, *bufio.ReadWriter, error) {
106-
hj, ok := i.ResponseWriter.(http.Hijacker)
107-
if !ok {
108-
return nil, nil, fmt.Errorf("interceptor: can't cast parent ResponseWriter to Hijacker")
114+
// Returns a wrapped http.ResponseWriter that implements the same optional interfaces
115+
// that the underlying ResponseWriter has.
116+
// Handle every possible combination so that code that checks for the existence of each
117+
// optional interface functions properly.
118+
// Based on https://github.com/felixge/httpsnoop/blob/eadd4fad6aac69ae62379194fe0219f3dbc80fd3/wrap_generated_gteq_1.8.go#L66
119+
func (i *interceptor) wrappedResponseWriter() http.ResponseWriter {
120+
closeNotifier, isCloseNotifier := i.ResponseWriter.(http.CloseNotifier)
121+
flush, isFlusher := i.ResponseWriter.(http.Flusher)
122+
hijack, isHijacker := i.ResponseWriter.(http.Hijacker)
123+
push, isPusher := i.ResponseWriter.(http.Pusher)
124+
readFrom, isReaderFrom := i.ResponseWriter.(io.ReaderFrom)
125+
126+
switch {
127+
case !isCloseNotifier && !isFlusher && !isHijacker && !isPusher && !isReaderFrom:
128+
return struct {
129+
http.ResponseWriter
130+
}{i}
131+
132+
case isCloseNotifier && !isFlusher && !isHijacker && !isPusher && !isReaderFrom:
133+
return struct {
134+
http.ResponseWriter
135+
http.CloseNotifier
136+
}{i, closeNotifier}
137+
138+
case !isCloseNotifier && isFlusher && !isHijacker && !isPusher && !isReaderFrom:
139+
return struct {
140+
http.ResponseWriter
141+
http.Flusher
142+
}{i, flush}
143+
144+
case !isCloseNotifier && !isFlusher && isHijacker && !isPusher && !isReaderFrom:
145+
return struct {
146+
http.ResponseWriter
147+
http.Hijacker
148+
}{i, hijack}
149+
150+
case !isCloseNotifier && !isFlusher && !isHijacker && isPusher && !isReaderFrom:
151+
return struct {
152+
http.ResponseWriter
153+
http.Pusher
154+
}{i, push}
155+
156+
case !isCloseNotifier && !isFlusher && !isHijacker && !isPusher && isReaderFrom:
157+
return struct {
158+
http.ResponseWriter
159+
io.ReaderFrom
160+
}{i, readFrom}
161+
162+
case isCloseNotifier && isFlusher && !isHijacker && !isPusher && !isReaderFrom:
163+
return struct {
164+
http.ResponseWriter
165+
http.CloseNotifier
166+
http.Flusher
167+
}{i, closeNotifier, flush}
168+
169+
case isCloseNotifier && !isFlusher && isHijacker && !isPusher && !isReaderFrom:
170+
return struct {
171+
http.ResponseWriter
172+
http.CloseNotifier
173+
http.Hijacker
174+
}{i, closeNotifier, hijack}
175+
176+
case isCloseNotifier && !isFlusher && !isHijacker && isPusher && !isReaderFrom:
177+
return struct {
178+
http.ResponseWriter
179+
http.CloseNotifier
180+
http.Pusher
181+
}{i, closeNotifier, push}
182+
183+
case isCloseNotifier && !isFlusher && !isHijacker && !isPusher && isReaderFrom:
184+
return struct {
185+
http.ResponseWriter
186+
http.CloseNotifier
187+
io.ReaderFrom
188+
}{i, closeNotifier, readFrom}
189+
190+
case !isCloseNotifier && isFlusher && isHijacker && !isPusher && !isReaderFrom:
191+
return struct {
192+
http.ResponseWriter
193+
http.Flusher
194+
http.Hijacker
195+
}{i, flush, hijack}
196+
197+
case !isCloseNotifier && isFlusher && !isHijacker && isPusher && !isReaderFrom:
198+
return struct {
199+
http.ResponseWriter
200+
http.Flusher
201+
http.Pusher
202+
}{i, flush, push}
203+
204+
case !isCloseNotifier && isFlusher && !isHijacker && !isPusher && isReaderFrom:
205+
return struct {
206+
http.ResponseWriter
207+
http.Flusher
208+
io.ReaderFrom
209+
}{i, flush, readFrom}
210+
211+
case !isCloseNotifier && !isFlusher && isHijacker && isPusher && !isReaderFrom:
212+
return struct {
213+
http.ResponseWriter
214+
http.Hijacker
215+
http.Pusher
216+
}{i, hijack, push}
217+
218+
case !isCloseNotifier && !isFlusher && isHijacker && !isPusher && isReaderFrom:
219+
return struct {
220+
http.ResponseWriter
221+
http.Hijacker
222+
io.ReaderFrom
223+
}{i, hijack, readFrom}
224+
225+
case !isCloseNotifier && !isFlusher && !isHijacker && isPusher && isReaderFrom:
226+
return struct {
227+
http.ResponseWriter
228+
http.Pusher
229+
io.ReaderFrom
230+
}{i, push, readFrom}
231+
232+
case isCloseNotifier && isFlusher && isHijacker && !isPusher && !isReaderFrom:
233+
return struct {
234+
http.ResponseWriter
235+
http.CloseNotifier
236+
http.Flusher
237+
http.Hijacker
238+
}{i, closeNotifier, flush, hijack}
239+
240+
case isCloseNotifier && isFlusher && !isHijacker && isPusher && !isReaderFrom:
241+
return struct {
242+
http.ResponseWriter
243+
http.CloseNotifier
244+
http.Flusher
245+
http.Pusher
246+
}{i, closeNotifier, flush, push}
247+
248+
case isCloseNotifier && isFlusher && !isHijacker && !isPusher && isReaderFrom:
249+
return struct {
250+
http.ResponseWriter
251+
http.CloseNotifier
252+
http.Flusher
253+
io.ReaderFrom
254+
}{i, closeNotifier, flush, readFrom}
255+
256+
case isCloseNotifier && !isFlusher && isHijacker && isPusher && !isReaderFrom:
257+
return struct {
258+
http.ResponseWriter
259+
http.CloseNotifier
260+
http.Hijacker
261+
http.Pusher
262+
}{i, closeNotifier, hijack, push}
263+
264+
case isCloseNotifier && !isFlusher && isHijacker && !isPusher && isReaderFrom:
265+
return struct {
266+
http.ResponseWriter
267+
http.CloseNotifier
268+
http.Hijacker
269+
io.ReaderFrom
270+
}{i, closeNotifier, hijack, readFrom}
271+
272+
case isCloseNotifier && !isFlusher && !isHijacker && isPusher && isReaderFrom:
273+
return struct {
274+
http.ResponseWriter
275+
http.CloseNotifier
276+
http.Pusher
277+
io.ReaderFrom
278+
}{i, closeNotifier, push, readFrom}
279+
280+
case !isCloseNotifier && isFlusher && isHijacker && isPusher && !isReaderFrom:
281+
return struct {
282+
http.ResponseWriter
283+
http.Flusher
284+
http.Hijacker
285+
http.Pusher
286+
}{i, flush, hijack, push}
287+
288+
case !isCloseNotifier && isFlusher && isHijacker && !isPusher && isReaderFrom:
289+
return struct {
290+
http.ResponseWriter
291+
http.Flusher
292+
http.Hijacker
293+
io.ReaderFrom
294+
}{i, flush, hijack, readFrom}
295+
296+
case !isCloseNotifier && isFlusher && !isHijacker && isPusher && isReaderFrom:
297+
return struct {
298+
http.ResponseWriter
299+
http.Flusher
300+
http.Pusher
301+
io.ReaderFrom
302+
}{i, flush, push, readFrom}
303+
304+
case !isCloseNotifier && !isFlusher && isHijacker && isPusher && isReaderFrom:
305+
return struct {
306+
http.ResponseWriter
307+
http.Hijacker
308+
http.Pusher
309+
io.ReaderFrom
310+
}{i, hijack, push, readFrom}
311+
312+
case isCloseNotifier && isFlusher && isHijacker && isPusher && !isReaderFrom:
313+
return struct {
314+
http.ResponseWriter
315+
http.CloseNotifier
316+
http.Flusher
317+
http.Hijacker
318+
http.Pusher
319+
}{i, closeNotifier, flush, hijack, push}
320+
321+
case isCloseNotifier && isFlusher && isHijacker && !isPusher && isReaderFrom:
322+
return struct {
323+
http.ResponseWriter
324+
http.CloseNotifier
325+
http.Flusher
326+
http.Hijacker
327+
io.ReaderFrom
328+
}{i, closeNotifier, flush, hijack, readFrom}
329+
330+
case isCloseNotifier && isFlusher && !isHijacker && isPusher && isReaderFrom:
331+
return struct {
332+
http.ResponseWriter
333+
http.CloseNotifier
334+
http.Flusher
335+
http.Pusher
336+
io.ReaderFrom
337+
}{i, closeNotifier, flush, push, readFrom}
338+
339+
case isCloseNotifier && !isFlusher && isHijacker && isPusher && isReaderFrom:
340+
return struct {
341+
http.ResponseWriter
342+
http.CloseNotifier
343+
http.Hijacker
344+
http.Pusher
345+
io.ReaderFrom
346+
}{i, closeNotifier, hijack, push, readFrom}
347+
348+
case !isCloseNotifier && isFlusher && isHijacker && isPusher && isReaderFrom:
349+
return struct {
350+
http.ResponseWriter
351+
http.Flusher
352+
http.Hijacker
353+
http.Pusher
354+
io.ReaderFrom
355+
}{i, flush, hijack, push, readFrom}
356+
357+
case isCloseNotifier && isFlusher && isHijacker && isPusher && isReaderFrom:
358+
return struct {
359+
http.ResponseWriter
360+
http.CloseNotifier
361+
http.Flusher
362+
http.Hijacker
363+
http.Pusher
364+
io.ReaderFrom
365+
}{i, closeNotifier, flush, hijack, push, readFrom}
366+
367+
default:
368+
return struct {
369+
http.ResponseWriter
370+
}{i}
109371
}
110-
return hj.Hijack()
111372
}

0 commit comments

Comments
 (0)