Graceful shutdown blocks on running proxy handler connections (#2789)

internal/daemon/worker/proxy/proxy_conn_tracker.go

@@ -0,0 +1,40 @@
+package proxy
+
+import (
+	"context"
+	"net/http"
+	"sync/atomic"
+)
+
+// ProxyState contains the current state of proxies in this process.
+var ProxyState proxyState
+
+type proxyState struct {
+	proxyCount atomic.Int64
+}
+
+// CurrentProxiedConnections returns the current number of ongoing proxied
+// connections which are currently running the Handler's ProxyConnFn.
+func (p *proxyState) CurrentProxiedConnections() int64 {
+	return p.proxyCount.Load()
+}
+
+// HttpHandlerCounter records how many requests are currently running in the
+// wrapped Handler. This should be used for handlers that serve proxied traffic.
+func ProxyHandlerCounter(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		ProxyState.proxyCount.Add(1)
+		defer ProxyState.proxyCount.Add(-1)
+		h.ServeHTTP(rw, r)
+	})
+}
+
+// proxyConnFnCounter wraps a ProxyState and keeps the proxyCount incremented
+// while it runs.
+func proxyConnFnCounter(fn ProxyConnFn) ProxyConnFn {
+	return func(ctx context.Context) {
+		ProxyState.proxyCount.Add(1)
+		defer ProxyState.proxyCount.Add(-1)
+		fn(ctx)
+	}
+}

internal/daemon/worker/proxy/proxy_conn_tracker_test.go

@@ -0,0 +1,41 @@
+package proxy
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestProxyStateHelpers(t *testing.T) {
+	old := ProxyState
+	ProxyState = proxyState{}
+	t.Cleanup(func() {
+		ProxyState = old
+	})
+
+	t.Run("proxyConnFnCounter", func(t *testing.T) {
+		assert.EqualValues(t, 0, ProxyState.CurrentProxiedConnections())
+		proxyConnFnCounter(func(context.Context) {
+			assert.EqualValues(t, 1, ProxyState.CurrentProxiedConnections())
+		})(context.Background())
+		assert.EqualValues(t, 0, ProxyState.CurrentProxiedConnections())
+	})
+
+	t.Run("ProxyHandlerCounter", func(t *testing.T) {
+		assert.EqualValues(t, 0, ProxyState.CurrentProxiedConnections())
+		var handlerRan bool
+		h := ProxyHandlerCounter(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
+			handlerRan = true
+			assert.EqualValues(t, 1, ProxyState.CurrentProxiedConnections())
+		}))
+		assert.EqualValues(t, 0, ProxyState.CurrentProxiedConnections())
+		req, err := http.NewRequest("GET", "test/path", nil)
+		require.NoError(t, err)
+		h.ServeHTTP(httptest.NewRecorder(), req)
+		require.True(t, handlerRan)
+	})
+}