~evanj/cms

892c0c7ea54f6c330a9c4acf52ccfc1ab091b121 — Evan J 2 months ago 7829d91
Feat(pkg/flushresp): Break out wrapped response writer into its own
package.
3 files changed, 61 insertions(+), 23 deletions(-)

M cms.go
A pkg/flushresp/flushresp.go
A pkg/flushresp/flushresp_test.go
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