Panic upon timeout

internal/timeout/responsewriter.go

 package timeout
 import (
-	"errors"
 	"net/http"
 	"time"
 )
 type responseWriter struct {
-	responseWriter      http.ResponseWriter
-	timer               StartStopper
-	status              int
-	writeHeaderTimedOut bool
+	http.ResponseWriter
+	timer  StartStopper
+	status int
 }
 type writeResult struct {
@@ -18,12 +16,14 @@ type writeResult struct {
 	err error
 }
+type timeoutPanic string
+
 // Put a timeout on each WriteHeader and Write call on the
 // ResponseWriter. The timeout gets reset for every WriteHeader / Write
-// call.
+// call. Panics when the timeout is exceeded.
 func NewResponseWriter(rw http.ResponseWriter, ti time.Duration) http.ResponseWriter {
 	return &responseWriter{
-		responseWriter: rw,
+		ResponseWriter: rw,
 		timer:          NewStartStopper(ti),
 	}
 }
@@ -33,15 +33,11 @@ func (tw *responseWriter) Write(p []byte) (int, error) {
 		tw.WriteHeader(http.StatusOK)
 	}
-	if tw.writeHeaderTimedOut {
-		return 0, errors.New("responseWriter: Write called after WriteHeader timed out")
-	}
-
 	tw.timer.Start()
-	writeChan := make(chan writeResult)
+	writeChan := make(chan writeResult, 1)
 	go func() {
-		n, err := tw.responseWriter.Write(p)
+		n, err := tw.ResponseWriter.Write(p)
 		writeChan <- writeResult{n: n, err: err}
 	}()
@@ -50,10 +46,7 @@ func (tw *responseWriter) Write(p []byte) (int, error) {
 		tw.timer.Stop()
 		return wr.n, wr.err
 	case <-tw.timer.Chan():
-		go func() {
-			<-writeChan // allow writer goroutine to finish
-		}()
-		return 0, errors.New("responseWriter: Write timeout")
+		panic(timeoutPanic("responseWriter: Write() timeout"))
 	}
 }
@@ -67,21 +60,14 @@ func (tw *responseWriter) WriteHeader(status int) {
 	writeHeaderChan := make(chan struct{})
 	go func() {
-		tw.responseWriter.WriteHeader(status)
-		writeHeaderChan <- struct{}{}
+		tw.ResponseWriter.WriteHeader(status)
+		close(writeHeaderChan)
 	}()
 	select {
 	case <-writeHeaderChan:
 		tw.timer.Stop()
 	case <-tw.timer.Chan():
-		go func() {
-			<-writeHeaderChan
-		}()
-		tw.writeHeaderTimedOut = true
+		panic(timeoutPanic("responseWriter: WriteHeader() timeout"))
 	}
 }
-
-func (tw *responseWriter) Header() http.Header {
-	return tw.responseWriter.Header()
-}

internal/timeout/responsewriter_test.go

@@ -20,29 +20,43 @@ func TestResponseWriterTimeouts(t *testing.T) {
 	}{
 		{true, []time.Duration{fast}},
 		{true, []time.Duration{fast, fast, fast}},
-		{true, []time.Duration{slow}}, // there is no way to detect a single WriteHeader failure
+		{false, []time.Duration{slow}},
 		{false, []time.Duration{slow, fast, fast}},
 		{false, []time.Duration{fast, slow, fast}},
 	}
 	for _, tc := range testCases {
 		w := NewResponseWriter(&slowWriter{delays: tc.delays}, timeout)
-		seenErrors := []error{}
-		w.WriteHeader(200)
+		seenPanics := catchPanic(nil, func() {
+			w.WriteHeader(200)
+		})
 		for range tc.delays[1:] {
-			if _, err := w.Write(writeData); err != nil {
-				seenErrors = append(seenErrors, err)
-			}
+			seenPanics = catchPanic(seenPanics, func() {
+				if _, err := w.Write(writeData); err != nil {
+					t.Fatal(err)
+				}
+			})
 		}
-		if tc.success && len(seenErrors) > 0 {
-			t.Errorf("case %v: no errors expected but received %v", tc, seenErrors[0])
+		if tc.success && len(seenPanics) > 0 {
+			t.Errorf("case %v: no panics expected but received %v", tc, seenPanics[0])
 		}
-		if !tc.success && len(seenErrors) == 0 {
-			t.Errorf("case %v: error expected, none received", tc)
+		if !tc.success && len(seenPanics) == 0 {
+			t.Errorf("case %v: panic expected, none happened", tc)
 		}
 	}
 }
+func catchPanic(seenPanics []timeoutPanic, cb func()) (result []timeoutPanic) {
+	result = append([]timeoutPanic{}, seenPanics...)
+	defer func() {
+		if p := recover(); p != nil {
+			result = append(result, p.(timeoutPanic))
+		}
+	}()
+	cb()
+	return result
+}
+
 type slowWriter struct {
 	delays []time.Duration
 	i      int