3 files changed, 61 insertions(+), 23 deletions(-)

M cms.go
A pkg/flushresp/flushresp.go
A pkg/flushresp/flushresp_test.go

@@ 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)
 			}

@@ 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())
+}

@@ 0,0 1,1 @@
+package flushresp_test