Bump golang version, add io.StringWriter and improve performance (#33)

* feat: add support for io.StringWriter

The interface was added to the io package in go 1.12.

 Related to:
 - https://go.dev/doc/go1.12#iopkgio
 - golang/go@33d531d

* chore: support go 1.25 and up

Following the same policy as golang.
See https://go.dev/doc/devel/release#policy

* chore: ran go fmt

* chore: benchmark

* refactor: invoke middleware func once

This removes allocation for each Write call

* refactor: reduce allocation by defining types

A type is now created for each possible interface combination. Allowing Wrap to so a single allocation.

* chore: combine the if statements

* fix: capture metrics for WriteString

* chore: add httpFlushError support

---------

Co-authored-by: Felix Geisendörfer <felix@felixge.de>

Author: boekkooi-impossiblecloud

.github/workflows/main.yaml

@@ -5,9 +5,8 @@ jobs:
     strategy:
       matrix:
         go:
-          - 1.18.x
-          - 1.17.x
-          - 1.16.x
+          - 1.26.x
+          - 1.25.x
         os:
           - ubuntu-latest
     name: ${{ matrix.os }}/go${{ matrix.go }}

capture_metrics.go

@@ -69,6 +69,16 @@ func (m *Metrics) CaptureMetrics(w http.ResponseWriter, fn func(http.ResponseWri
 				}
 			},
 
+			WriteString: func(next WriteStringFunc) WriteStringFunc {
+				return func(s string) (int, error) {
+					n, err := next(s)
+
+					m.Written += int64(n)
+					headerWritten = true
+					return n, err
+				}
+			},
+
 			ReadFrom: func(next ReadFromFunc) ReadFromFunc {
 				return func(src io.Reader) (int64, error) {
 					n, err := next(src)
@@ -101,3 +111,11 @@ type deadliner interface {
 type fullDuplexEnabler interface {
 	EnableFullDuplex() error
 }
+
+// httpFlushError defines a method introduced in go 1.20. The standard
+// library seems not to provide an interface we can import, hence its definition
+// here.
+// See https://github.com/golang/go/blob/go1.20/src/net/http/responsecontroller.go#L50
+type httpFlushError interface {
+	FlushError() error
+}

bench_test.go capture_metrics_bench_test.gobench_test.go renamed to capture_metrics_bench_test.go

@@ -18,24 +18,6 @@ func BenchmarkCaptureMetricsTwice(b *testing.B) {
 	benchmark(b, 2)
 }
 
-func BenchmarkWrap(b *testing.B) {
-	b.StopTimer()
-	doneCh := make(chan struct{}, 1)
-	h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		b.StartTimer()
-		for i := 0; i < b.N; i++ {
-			Wrap(w, Hooks{})
-		}
-		doneCh <- struct{}{}
-	})
-	s := httptest.NewServer(h)
-	defer s.Close()
-	if _, err := http.Get(s.URL); err != nil {
-		b.Fatal(err)
-	}
-	<-doneCh
-}
-
 func benchmark(b *testing.B, wrappings int) {
 	dummyH := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
 	h := dummyH

capture_metrics_test.go

@@ -61,6 +61,14 @@ func TestCaptureMetrics(t *testing.T) {
 			WantWritten: 17,
 			WantCode:    http.StatusOK,
 		},
+		{
+			Name: "string writer",
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				_, _ = io.WriteString(w, "write string")
+			}),
+			WantWritten: int64(len("write string")),
+			WantCode:    http.StatusOK,
+		},
 		{
 			Name: "empty panic",
 			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -104,7 +112,7 @@ func TestCaptureMetrics(t *testing.T) {
 			}
 			if err == nil {
 				defer res.Body.Close()
- 			}
+			}
 			m := <-ch
 			if m.Code != test.WantCode {
 				t.Errorf("test %d: got=%d want=%d", i, m.Code, test.WantCode)

codegen/main.go

@@ -15,19 +15,25 @@ type Build struct {
 }
 
 func (b *Build) MustBuild() {
-	prefix := "wrap_generated_"
-	b.Implementation().MustWriteFile(prefix + b.Suffix + ".go")
-	b.Tests().MustWriteFile(prefix + b.Suffix + "_test.go")
+	prefix := "wrap_generated"
+	if b.Suffix != "" {
+		prefix += "_" + b.Suffix
+	}
+
+	b.Implementation().MustWriteFile(prefix + ".go")
+	b.Tests().MustWriteFile(prefix + "_test.go")
 }
 
 func (b *Build) writeHeader(g *Generator) {
-	g.Printf(`
-// +build %s
-// Code generated by "httpsnoop/codegen"; DO NOT EDIT.
+	if b.Tags != "" {
+		g.Printf(`// +build %s
+`, b.Tags)
+	}
+	g.buf.WriteString(`// Code generated by "httpsnoop/codegen"; DO NOT EDIT.
 
 package httpsnoop
 
-`, b.Tags)
+`)
 }
 
 func (b *Build) Implementation() *Generator {
@@ -89,64 +95,110 @@ type Hooks struct {
 // hooks can be used.
 `, strings.Join(docList, "\n"))
 	g.Printf("func Wrap(w http.ResponseWriter, hooks Hooks) http.ResponseWriter {\n")
-	g.Printf("rw := &rw{w: w, h: hooks}\n")
+	g.Printf("state := &rwState{w: w}\n")
+
+	// Precompute hook chains once per Wrap call and
+	// build a uint8 combo index so the switch compiles to a jump table.
+	g.Printf("var combo uint16\n")
+	for _, fn := range ifaces[0].Funcs {
+		g.Printf("if hooks.%s != nil {\n", fn.Name)
+		g.Printf("state.%s = hooks.%s(w.%s)\n", fieldName(fn.Name), fn.Name, fn.Name)
+		g.Printf("}\n")
+	}
+
 	for i, iface := range subIfaces {
-		g.Printf("_, i%d := w.(%s)\n", i, iface.Name)
+		g.Printf("if t%[1]d, i%[1]d := w.(%s); i%[1]d {\n", i, iface.Name)
+		bit := len(subIfaces) - i - 1
+		g.Printf("combo |= 1<<%d\n", bit)
+		for _, fn := range iface.Funcs {
+			g.Printf("if hooks.%s != nil {\n", fn.Name)
+			g.Printf("state.%s = hooks.%s(t%d.%s)\n", fieldName(fn.Name), fn.Name, i, fn.Name)
+			g.Printf("}\n")
+		}
+		g.Printf("}\n")
 	}
-	g.Printf("switch {\n")
+
+	g.Printf("switch combo {\n")
 	combinations := 1 << uint(len(subIfaces))
-	for i := 0; i < combinations; i++ {
-		conditions := make([]string, len(subIfaces))
-		fields := make([]string, 0, len(subIfaces))
-		fields = append(fields, "Unwrapper", "http.ResponseWriter")
-		for j, iface := range subIfaces {
-			ok := i&(1<<uint(len(subIfaces)-j-1)) > 0
-			if !ok {
-				conditions[j] = "!"
-			} else {
-				fields = append(fields, iface.Name)
-			}
-			conditions[j] += fmt.Sprintf("i%d", j)
-		}
-		values := make([]string, len(fields))
-		for i := range fields {
-			values[i] = "rw"
-		}
-		g.Printf("// combination %d/%d\n", i+1, combinations)
-		g.Printf("case %s:\n", strings.Join(conditions, "&&"))
-		fieldsS, valuesS := strings.Join(fields, "\n"), strings.Join(values, ",")
-		g.Printf("return struct{\n%s\n}{%s}\n", fieldsS, valuesS)
+	for c := 0; c < combinations; c++ {
+		g.Printf("case %d: return (*rw%d)(state)\n", c, c)
 	}
 	g.Printf("}\n")
 	g.Printf("panic(\"unreachable\")")
-	g.Printf("}\n")
-
-	// rw struct
-	g.Printf(`
-type rw struct {
-	w http.ResponseWriter
-	h Hooks
-}
+	g.Printf("}\n\n")
 
-func (w *rw) Unwrap() http.ResponseWriter {
-	 return w.w
-}
+	// rwState holds the underlying writer plus the precomputed hooks.
+	// All variant types are type-definitions over rwState, so a single *rwState
+	// allocation can be reinterpreted as any variant via pointer conversion.
+	g.Printf("type rwState struct {\n")
+	g.Printf("w http.ResponseWriter\n")
+	for _, iface := range ifaces {
+		for _, fn := range iface.Funcs {
+			g.Printf("%s %s\n", fieldName(fn.Name), fn.Type())
+		}
+	}
+	g.Printf("}\n\n")
 
-`)
+	// do<Name> helpers on *rwState
+	// These actual dispatch logic, defined once and called by the variant types.
 	for _, iface := range ifaces {
 		for _, fn := range iface.Funcs {
-			g.Printf("func (w *rw) %s(%s) (%s) {\n", fn.Name, fn.Args, fn.Returns)
-			g.Printf("f := w.w.(%s).%s\n", iface.Name, fn.Name)
-			g.Printf("if w.h.%s != nil {\n", fn.Name)
-			g.Printf("f = w.h.%s(f)\n", fn.Name)
-			g.Printf("}\n")
+			g.Printf("func (r *rwState) do%s(%s) (%s) {\n", fn.Name, fn.Args, fn.Returns)
+			g.Printf("if r.%s != nil {\n", fieldName(fn.Name))
 			if fn.Returns != "" {
-				g.Printf("return ")
+				g.Printf("return r.%s(%s)\n", fieldName(fn.Name), fn.Args.Names())
+			} else {
+				g.Printf("r.%s(%s)\n", fieldName(fn.Name), fn.Args.Names())
+				g.Printf("return\n")
 			}
-			g.Printf("f(%s)\n", fn.Args.Names())
 			g.Printf("}\n")
-			g.Printf("\n")
+
+			receiver := "r.w"
+			if iface.Name != "http.ResponseWriter" {
+				receiver = fmt.Sprintf("r.w.(%s)", iface.Name)
+			}
+			if fn.Returns != "" {
+				g.Printf("return %s.%s(%s)\n", receiver, fn.Name, fn.Args.Names())
+			} else {
+				g.Printf("%s.%s(%s)\n", receiver, fn.Name, fn.Args.Names())
+			}
+			g.Printf("}\n\n")
+		}
+	}
+
+	// Variant types, each is a type with the same memory layout as rwState,
+	// but exposing exactly the method set required by its combination of interfaces.
+	// This allows (*rwN)(state) to be a zero-cost pointer conversion.
+	emitVariantMethod := func(c int, fn *InterfaceFunc) {
+		g.Printf("func (w *rw%d) %s(%s) (%s) {\n", c, fn.Name, fn.Args, fn.Returns)
+		if fn.Returns != "" {
+			g.Printf("return (*rwState)(w).do%s(%s)\n", fn.Name, fn.Args.Names())
+		} else {
+			g.Printf("(*rwState)(w).do%s(%s)\n", fn.Name, fn.Args.Names())
 		}
+		g.Printf("}\n")
+	}
+	for c := 0; c < combinations; c++ {
+		supported := []string{"http.ResponseWriter"}
+		for j, iface := range subIfaces {
+			if c&(1<<uint(len(subIfaces)-j-1)) > 0 {
+				supported = append(supported, iface.Name)
+			}
+		}
+		g.Printf("// combination %d/%d: %s\n", c+1, combinations, strings.Join(supported, ", "))
+		g.Printf("type rw%d rwState\n", c)
+		g.Printf("func (w *rw%d) Unwrap() http.ResponseWriter { return w.w }\n", c)
+		for _, fn := range ifaces[0].Funcs {
+			emitVariantMethod(c, fn)
+		}
+		for j, iface := range subIfaces {
+			if c&(1<<uint(len(subIfaces)-j-1)) > 0 {
+				for _, fn := range iface.Funcs {
+					emitVariantMethod(c, fn)
+				}
+			}
+		}
+		g.Printf("\n")
 	}
 	g.Printf(`
 type Unwrapper interface {
@@ -159,9 +211,8 @@ func Unwrap(w http.ResponseWriter) http.ResponseWriter {
 	if rw, ok := w.(Unwrapper); ok {
 		// recurse until rw.Unwrap() returns a non-Unwrapper
 		return Unwrap(rw.Unwrap())
-	} else {
-		return w
 	}
+	return w
 }
 `)
 	return &g
@@ -263,12 +314,19 @@ func (fn *InterfaceFunc) Type() string {
 	return fn.Name + "Func"
 }
 
+func fieldName(s string) string {
+	if s == "" {
+		return s
+	}
+	return strings.ToLower(s[:1]) + s[1:]
+}
+
 type Generator struct {
 	buf bytes.Buffer
 }
 
 func (g *Generator) Printf(s string, args ...interface{}) {
-	fmt.Fprintf(&g.buf, s, args...)
+	_, _ = fmt.Fprintf(&g.buf, s, args...)
 }
 
 func (g *Generator) WriteFile(name string) error {
@@ -311,6 +369,12 @@ func main() {
 				{"Flush", nil, ""},
 			},
 		},
+		{
+			Name: "httpFlushError", // Introduced in Go 1.20.
+			Funcs: []*InterfaceFunc{
+				{"FlushError", nil, "error"},
+			},
+		},
 		{
 			Name: "http.CloseNotifier",
 			Funcs: []*InterfaceFunc{
@@ -342,33 +406,32 @@ func main() {
 				{"EnableFullDuplex", nil, "error"},
 			},
 		},
+		{
+			Name: "http.Pusher",
+			Funcs: []*InterfaceFunc{
+				{"Push", FuncArgs{
+					{"target", "string"},
+					{"opts", "*http.PushOptions"},
+				}, "error"},
+			},
+		}, {
+			Name: "io.StringWriter",
+			Funcs: []*InterfaceFunc{
+				{"WriteString", FuncArgs{{"s", "string"}}, "int, error"},
+			},
+		},
 	}
 	builds := []Build{
 		{
-			Suffix:     "lt_1.8",
-			Tags:       "!go1.8",
 			Interfaces: ifaces,
 		},
-		{
-			Suffix: "gteq_1.8",
-			Tags:   "go1.8",
-			Interfaces: append(ifaces, &Interface{
-				Name: "http.Pusher",
-				Funcs: []*InterfaceFunc{
-					{"Push", FuncArgs{
-						{"target", "string"},
-						{"opts", "*http.PushOptions"},
-					}, "error"},
-				},
-			}),
-		},
 	}
 	for _, build := range builds {
 		build.MustBuild()
 	}
 }
 
 func fatalf(s string, args ...interface{}) {
-	fmt.Fprintf(os.Stderr, s+"\n", args...)
+	_, _ = fmt.Fprintf(os.Stderr, s+"\n", args...)
 	os.Exit(1)
 }

go.mod

@@ -1,3 +1,3 @@
 module github.com/felixge/httpsnoop
 
-go 1.13
+go 1.25