Handle observer cleanup more aggressively to catch leak

proxymw/middleware_test.go

@@ -36,6 +36,7 @@ func TestMiddlewareOrder(t *testing.T) {
 		JitterDelay:  time.Second,
 
 		EnableObserver: true,
+		ClientTimeout:  time.Hour,
 	}
 
 	serveCalls := 0

proxymw/mock.go

@@ -11,11 +11,13 @@ type Mocker struct {
 	RoundTripFunc func(r *http.Request) (*http.Response, error)
 	InitFunc      func(context.Context)
 	NextFunc      func(Request) error
+	RequestFunc   func() *http.Request
 }
 
 var _ http.Handler = &Mocker{}
 var _ http.RoundTripper = &Mocker{}
 var _ ProxyClient = &Mocker{}
+var _ Request = &Mocker{}
 
 func (m *Mocker) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	m.ServeHTTPFunc(w, r)
@@ -32,3 +34,7 @@ func (m *Mocker) Init(ctx context.Context) {
 func (m *Mocker) Next(rr Request) error {
 	return m.NextFunc(rr)
 }
+
+func (m *Mocker) Request() *http.Request {
+	return m.RequestFunc()
+}

proxymw/observer.go

@@ -3,6 +3,7 @@ package proxymw
 import (
 	"context"
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -52,7 +53,7 @@ func (o *Observer) Init(ctx context.Context) {
 func (o *Observer) Next(rr Request) error {
 	o.activeGauge.Inc()
 	start := time.Now()
-	err := o.client.Next(rr)
+	err := o.executeNext(rr)
 	if err != nil {
 		var blocked *RequestBlockedError
 		if errors.As(err, &blocked) {
@@ -67,3 +68,25 @@ func (o *Observer) Next(rr Request) error {
 	o.activeGauge.Dec()
 	return err
 }
+
+// executeNext runs next in a goroutine on the off chance Next hangs so we can still run cleanup
+func (o *Observer) executeNext(rr Request) error {
+	errc := make(chan error, 1)
+	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				errc <- fmt.Errorf("panic calling Next: %v", r)
+			}
+			close(errc)
+		}()
+		errc <- o.client.Next(rr)
+	}()
+
+	ctx := rr.Request().Context()
+	select {
+	case err := <-errc:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}

proxymw/observer_test.go

@@ -3,6 +3,7 @@ package proxymw
 import (
 	"context"
 	"errors"
+	"net/http"
 	"testing"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -57,6 +58,34 @@ func TestObserverNextError(t *testing.T) {
 				require.Equal(t, float64(1), value)
 			},
 		},
+		{
+			name: "next panic",
+			observer: &Observer{
+				errCounter: prometheus.NewCounter(
+					prometheus.CounterOpts{Name: "block_test_error_count"},
+				),
+				blockCounter: prometheus.NewCounterVec(
+					prometheus.CounterOpts{Name: "block_test_block_count"}, []string{"mw_type"},
+				),
+				reqCounter: prometheus.NewCounter(
+					prometheus.CounterOpts{Name: "block_test_request_count"},
+				),
+				latencyCounter: prometheus.NewCounter(
+					prometheus.CounterOpts{Name: "block_test_request_latency_ms"},
+				),
+				activeGauge: prometheus.NewGauge(
+					prometheus.GaugeOpts{Name: "block_test_active_requests"},
+				),
+				client: &Mocker{
+					NextFunc: func(_ Request) error {
+						panic("here")
+					},
+					InitFunc: func(_ context.Context) {},
+				},
+			},
+			err:   "panic calling Next: here",
+			check: func(_ *testing.T, _ *Observer) {},
+		},
 		{
 			name: "normal error",
 			observer: &Observer{
@@ -77,7 +106,6 @@ func TestObserverNextError(t *testing.T) {
 				),
 				client: &Mocker{
 					NextFunc: func(r Request) error {
-						require.Equal(t, nil, r)
 						return errors.New("fail")
 					},
 					InitFunc: func(_ context.Context) {
@@ -141,8 +169,14 @@ func TestObserverNextError(t *testing.T) {
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
-			tt.observer.Init(context.Background())
-			err := tt.observer.Next(nil)
+			ctx := context.Background()
+			tt.observer.Init(ctx)
+			rr := &Mocker{
+				RequestFunc: func() *http.Request {
+					return (&http.Request{}).WithContext(ctx)
+				},
+			}
+			err := tt.observer.Next(rr)
 			errStr := ""
 			if err != nil {
 				errStr = err.Error()