M cms.go => cms.go +4 -23
@@ 2,12 2,13 @@
package main
import (
- "bytes"
"context"
"log"
"net/http"
"strings"
"time"
+
+ "git.sr.ht/~evanj/cms/pkg/flushresp"
)
type App struct {
@@ 39,26 40,6 @@ func (a *App) ServeHTTP(w http.ResponseWriter, r *http.Request) {
a.handleWithRetry(w, r, h)
}
-type wrappedResponseWriter struct {
- w http.ResponseWriter
- h http.Header
- sc int
- b bytes.Buffer
-}
-
-func (rw *wrappedResponseWriter) Header() http.Header { return rw.h }
-func (rw *wrappedResponseWriter) Write(bytes []byte) (int, error) { return rw.b.Write(bytes) }
-func (rw *wrappedResponseWriter) WriteHeader(statusCode int) { rw.sc = statusCode }
-
-func (rw *wrappedResponseWriter) Flush() error {
- for key := range rw.h {
- rw.w.Header().Add(key, rw.h.Get(key))
- }
- rw.w.WriteHeader(rw.sc)
- _, err := rw.w.Write(rw.b.Bytes())
- return err
-}
-
// handleWithRetry will retry requests when we're approaching long request
// times in an effort to lower tail latency.
func (a *App) handleWithRetry(w http.ResponseWriter, r *http.Request, h http.Handler) {
@@ 73,7 54,7 @@ func (a *App) handleWithRetry(w http.ResponseWriter, r *http.Request, h http.Han
for i := 0; i < maxAttempts; i++ {
var (
ctx, cancel = context.WithTimeout(context.Background(), maxTime)
- wWrap = &wrappedResponseWriter{w: w, h: http.Header{}, b: bytes.Buffer{}}
+ wWrap = flushresp.New(w)
successCh = handle(wWrap, r.WithContext(ctx), h)
)
defer cancel()
@@ 82,7 63,7 @@ func (a *App) handleWithRetry(w http.ResponseWriter, r *http.Request, h http.Han
case <-ctx.Done():
err = ctx.Err()
case <-successCh:
- if err := wWrap.Flush(); err != nil {
+ if _, err := wWrap.Flush(); err != nil {
// TODO: Is there a more appropriate way to handle an error here?...
a.log.Println(err)
}
A pkg/flushresp/flushresp.go => pkg/flushresp/flushresp.go +56 -0
@@ 0,0 1,56 @@
+package flushresp
+
+import (
+ "bytes"
+ "net/http"
+)
+
+var _ http.ResponseWriter = &Response{}
+
+// Response is to be used in place of an http.ResponseWriter.
+// Example:
+// func main() {
+// http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+// w2 := New(w)
+// w2.Header().Add("Content-Type", "text/plain")
+// w2.WriteHeader(http.StatusOK)
+// w2.Write([]byte("hello world"))
+// w2.Flush()
+// })
+// }
+type Response struct {
+ w http.ResponseWriter
+ h http.Header
+ c int
+ b bytes.Buffer
+}
+
+// New will, given an http.ResponseWriter, create a flushresp.Response, which
+// is to be used as an http.ResponseWriter (as flushresp.Response implements
+// http.ResponseWriter).
+func New(w http.ResponseWriter) *Response {
+ return &Response{
+ w: w,
+ h: http.Header{},
+ b: bytes.Buffer{},
+ }
+}
+
+// Header, see http.ResponseWriter
+func (rw *Response) Header() http.Header { return rw.h }
+
+// WriteHeader, see http.ResponseWriter
+func (rw *Response) WriteHeader(statusCode int) { rw.c = statusCode }
+
+// Write, see http.ResponseWriter
+func (rw *Response) Write(bytes []byte) (int, error) { return rw.b.Write(bytes) }
+
+// Flush will complete the Response writing and write data to the underlying
+// http.ResponseWriter supplied via flushresp.New.
+func (rw *Response) Flush() (int, error) {
+ for key := range rw.h {
+ rw.w.Header().Add(key, rw.h.Get(key))
+ }
+ rw.w.WriteHeader(rw.c)
+ return rw.w.Write(rw.b.Bytes())
+}
A pkg/flushresp/flushresp_test.go => pkg/flushresp/flushresp_test.go +1 -0
@@ 0,0 1,1 @@
+package flushresp_test