fix: track in-flight registry proxy requests for graceful drain

Adds registryInFlight atomic.Int64 counter to proxy.Service and
instruments proxyToRegistry to increment/decrement it. Adds
DrainRegistryInFlight method for use during shutdown. This allows
gracefulShutdown to wait for active manifest pushes before stopping
the registry backend.

commit: b93860a8b973cb6b6a9d52565c63cd3110b4c43c
author: brice · 7mo ago · 2026-03-01 11:06 UTC
committer: brice · 7mo ago · 2026-03-01 11:08 UTC
parent: c8b7fd0
signature: unsigned

2 files · +92 -9

M internal/usecase/proxy/service.go +28 -9


@@ -70,15 +70,16 @@ type Config struct {
70	70
71	71		// Service implements the ProxyService interface.
72	72		type Service struct {
73		-	runtime out.ContainerRuntime
74		-	containerSvc in.ContainerService
75		-	configSvc in.ConfigService
76		-	config Config
77		-	targets map[string]*domain.ProxyTarget
78		-	mu sync.RWMutex
79		-	activeConns atomic.Int64 // Atomic counter for concurrent connection limiting
80		-	inFlight map[string]int
81		-	inFlightMu sync.Mutex
	73	+	runtime out.ContainerRuntime
	74	+	containerSvc in.ContainerService
	75	+	configSvc in.ConfigService
	76	+	config Config
	77	+	targets map[string]*domain.ProxyTarget
	78	+	mu sync.RWMutex
	79	+	activeConns atomic.Int64 // Atomic counter for concurrent connection limiting
	80	+	inFlight map[string]int
	81	+	inFlightMu sync.Mutex
	82	+	registryInFlight atomic.Int64 // active registry proxy requests, for graceful drain
82	83		}
83	84
84	85		// NewService creates a new proxy service.

@@ -574,6 +575,21 @@ func (s *Service) trackInFlight(containerID string) func() {
574	575		}
575	576		}
576	577
	578	+	// DrainRegistryInFlight blocks until all in-flight registry proxy requests
	579	+	// complete or the timeout elapses. Returns true if drained cleanly, false if
	580	+	// timed out with requests still in flight. Call this before shutting down the
	581	+	// registry server.
	582	+	func (s *Service) DrainRegistryInFlight(timeout time.Duration) bool {
	583	+	deadline := time.Now().Add(timeout)
	584	+	for time.Now().Before(deadline) {
	585	+	if s.registryInFlight.Load() == 0 {
	586	+	return true
	587	+	}
	588	+	time.Sleep(10 * time.Millisecond)
	589	+	}
	590	+	return false
	591	+	}
	592	+
577	593		// proxyToRegistry forwards requests to the local registry HTTP server.
578	594		// SECURITY: This uses http://localhost:{port} which is safe because the registry
579	595		// runs on the same host and traffic never leaves the loopback interface. If Gordon

@@ -585,6 +601,9 @@ func (s Service) proxyToRegistry(w http.ResponseWriter, r http.Request) {
585	601		cfg := s.config
586	602		s.mu.RUnlock()
587	603
	604	+	s.registryInFlight.Add(1)
	605	+	defer s.registryInFlight.Add(-1)
	606	+
588	607		log := zerowrap.FromCtx(r.Context())
589	608
590	609		targetURL, err := url.Parse(fmt.Sprintf("http://localhost:%d", cfg.RegistryPort))

 // Service implements the ProxyService interface.
 type Service struct {
-	runtime      out.ContainerRuntime
+	runtime          out.ContainerRuntime
-	containerSvc in.ContainerService
+	containerSvc     in.ContainerService
-	configSvc    in.ConfigService
+	configSvc        in.ConfigService
-	config       Config
+	config           Config
-	targets      map[string]*domain.ProxyTarget
+	targets          map[string]*domain.ProxyTarget
-	mu           sync.RWMutex
+	mu               sync.RWMutex
-	activeConns  atomic.Int64 // Atomic counter for concurrent connection limiting
+	activeConns      atomic.Int64 // Atomic counter for concurrent connection limiting
-	inFlight     map[string]int
+	inFlight         map[string]int
-	inFlightMu   sync.Mutex
+	inFlightMu       sync.Mutex
+	registryInFlight atomic.Int64 // active registry proxy requests, for graceful drain
 }
 // NewService creates a new proxy service.
 	}
 }
+// DrainRegistryInFlight blocks until all in-flight registry proxy requests
+// complete or the timeout elapses. Returns true if drained cleanly, false if
+// timed out with requests still in flight. Call this before shutting down the
+// registry server.
+func (s *Service) DrainRegistryInFlight(timeout time.Duration) bool {
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		if s.registryInFlight.Load() == 0 {
+			return true
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	return false
+}
 // proxyToRegistry forwards requests to the local registry HTTP server.
 // SECURITY: This uses http://localhost:{port} which is safe because the registry
 // runs on the same host and traffic never leaves the loopback interface. If Gordon
 	cfg := s.config
 	s.mu.RUnlock()
+	s.registryInFlight.Add(1)
+	defer s.registryInFlight.Add(-1)
 	log := zerowrap.FromCtx(r.Context())
 	targetURL, err := url.Parse(fmt.Sprintf("http://localhost:%d", cfg.RegistryPort))

M internal/usecase/proxy/service_test.go +64 -0


@@ -4,6 +4,7 @@ import (
4	4		"context"
5	5		"net/http/httptest"
6	6		"testing"
	7	+	"time"
7	8
8	9		"github.com/bnema/zerowrap"
9	10		"github.com/stretchr/testify/assert"

@@ -510,3 +511,66 @@ func TestServeHTTP_MaxBodySize(t *testing.T) {
510	511		})
511	512		}
512	513		}
	514	+
	515	+	func TestRegistryInFlightTracking(t *testing.T) {
	516	+	svc := &Service{
	517	+	inFlight: make(map[string]int),
	518	+	}
	519	+
	520	+	if got := svc.registryInFlight.Load(); got != 0 {
	521	+	t.Fatalf("expected 0 in-flight, got %d", got)
	522	+	}
	523	+
	524	+	svc.registryInFlight.Add(1)
	525	+	if got := svc.registryInFlight.Load(); got != 1 {
	526	+	t.Fatalf("expected 1 in-flight after Add, got %d", got)
	527	+	}
	528	+
	529	+	svc.registryInFlight.Add(-1)
	530	+	if got := svc.registryInFlight.Load(); got != 0 {
	531	+	t.Fatalf("expected 0 in-flight after release, got %d", got)
	532	+	}
	533	+	}
	534	+
	535	+	func TestDrainRegistryInFlight(t *testing.T) {
	536	+	svc := &Service{
	537	+	inFlight: make(map[string]int),
	538	+	}
	539	+
	540	+	svc.registryInFlight.Add(2)
	541	+
	542	+	done := make(chan struct{})
	543	+	go func() {
	544	+	drained := svc.DrainRegistryInFlight(50 * time.Millisecond)
	545	+	if !drained {
	546	+	// Signal failure via done channel by leaving it open — test will time out
	547	+	return
	548	+	}
	549	+	close(done)
	550	+	}()
	551	+
	552	+	time.Sleep(5 * time.Millisecond)
	553	+	svc.registryInFlight.Add(-1)
	554	+	svc.registryInFlight.Add(-1)
	555	+
	556	+	select {
	557	+	case <-done:
	558	+	// good — drained cleanly
	559	+	case <-time.After(500 * time.Millisecond):
	560	+	t.Fatal("DrainRegistryInFlight did not return true after requests completed")
	561	+	}
	562	+	}
	563	+
	564	+	func TestDrainRegistryInFlightTimeout(t *testing.T) {
	565	+	svc := &Service{
	566	+	inFlight: make(map[string]int),
	567	+	}
	568	+
	569	+	// Add a request and never release it
	570	+	svc.registryInFlight.Add(1)
	571	+
	572	+	drained := svc.DrainRegistryInFlight(30 * time.Millisecond)
	573	+	if drained {
	574	+	t.Fatal("expected DrainRegistryInFlight to return false on timeout, got true")
	575	+	}
	576	+	}


@@ -4,6 +4,7 @@ import (
4	"context"	4		"context"
5	"net/http/httptest"	5		"net/http/httptest"
6	"testing"	6		"testing"
		7	+	"time"
7		8
8	"github.com/bnema/zerowrap"	9		"github.com/bnema/zerowrap"
9	"github.com/stretchr/testify/assert"	10		"github.com/stretchr/testify/assert"

@@ -510,3 +511,66 @@ func TestServeHTTP_MaxBodySize(t *testing.T) {
510	})	511		})
511	}	512		}
512	}	513		}
		514	+
		515	+	func TestRegistryInFlightTracking(t *testing.T) {
		516	+	svc := &Service{
		517	+	inFlight: make(map[string]int),
		518	+	}
		519	+
		520	+	if got := svc.registryInFlight.Load(); got != 0 {
		521	+	t.Fatalf("expected 0 in-flight, got %d", got)
		522	+	}
		523	+
		524	+	svc.registryInFlight.Add(1)
		525	+	if got := svc.registryInFlight.Load(); got != 1 {
		526	+	t.Fatalf("expected 1 in-flight after Add, got %d", got)
		527	+	}
		528	+
		529	+	svc.registryInFlight.Add(-1)
		530	+	if got := svc.registryInFlight.Load(); got != 0 {
		531	+	t.Fatalf("expected 0 in-flight after release, got %d", got)
		532	+	}
		533	+	}
		534	+
		535	+	func TestDrainRegistryInFlight(t *testing.T) {
		536	+	svc := &Service{
		537	+	inFlight: make(map[string]int),
		538	+	}
		539	+
		540	+	svc.registryInFlight.Add(2)
		541	+
		542	+	done := make(chan struct{})
		543	+	go func() {
		544	+	drained := svc.DrainRegistryInFlight(50 * time.Millisecond)
		545	+	if !drained {
		546	+	// Signal failure via done channel by leaving it open — test will time out
		547	+	return
		548	+	}
		549	+	close(done)
		550	+	}()
		551	+
		552	+	time.Sleep(5 * time.Millisecond)
		553	+	svc.registryInFlight.Add(-1)
		554	+	svc.registryInFlight.Add(-1)
		555	+
		556	+	select {
		557	+	case <-done:
		558	+	// good — drained cleanly
		559	+	case <-time.After(500 * time.Millisecond):
		560	+	t.Fatal("DrainRegistryInFlight did not return true after requests completed")
		561	+	}
		562	+	}
		563	+
		564	+	func TestDrainRegistryInFlightTimeout(t *testing.T) {
		565	+	svc := &Service{
		566	+	inFlight: make(map[string]int),
		567	+	}
		568	+
		569	+	// Add a request and never release it
		570	+	svc.registryInFlight.Add(1)
		571	+
		572	+	drained := svc.DrainRegistryInFlight(30 * time.Millisecond)
		573	+	if drained {
		574	+	t.Fatal("expected DrainRegistryInFlight to return false on timeout, got true")
		575	+	}
		576	+	}