Remove arbitrary sleeps from shutdown handling

internal/king/export_test.go

@@ -1,6 +1,10 @@
 package king
 
-import "github.com/quic-go/quic-go"
+import (
+	"context"
+
+	"github.com/quic-go/quic-go"
+)
 
 // SetQUICConn sets a QUIC connection for a service ID (test helper).
 func (tunnelSrv *TunnelServer) SetQUICConn(serviceID string, conn *quic.Conn) {
@@ -35,3 +39,33 @@ func (tunnelSrv *TunnelServer) SetServiceAuth(serviceID string, auth ServiceAuth
 
 	tunnelSrv.services[serviceID] = auth
 }
+
+// RemoveQUICConn removes a QUIC connection and signals drainNotify (test helper).
+func (tunnelSrv *TunnelServer) RemoveQUICConn(serviceID string) {
+	tunnelSrv.mu.Lock()
+
+	delete(tunnelSrv.quicConns, serviceID)
+
+	empty := len(tunnelSrv.quicConns) == 0
+	tunnelSrv.mu.Unlock()
+
+	if empty {
+		select {
+		case tunnelSrv.drainNotify <- struct{}{}:
+		default:
+		}
+	}
+}
+
+// QUICConnCount returns the number of QUIC connections (test helper).
+func (tunnelSrv *TunnelServer) QUICConnCount() int {
+	tunnelSrv.mu.RLock()
+	defer tunnelSrv.mu.RUnlock()
+
+	return len(tunnelSrv.quicConns)
+}
+
+// WaitForDrain exposes the private waitForDrain method (test helper).
+func (tunnelSrv *TunnelServer) WaitForDrain(ctx context.Context) {
+	tunnelSrv.waitForDrain(ctx)
+}

internal/king/king.go

@@ -25,7 +25,7 @@ const (
 	defaultCouncilHost = "http://localhost:8080"
 	kvPairCount        = 2
 	finalSyncTimeout   = 2 * time.Second
-	shutdownGracePause = 1 * time.Second
+	quicDrainTimeout   = 5 * time.Second
 	certValidityYears  = 10 * 365 * 24 * time.Hour
 )
 
@@ -289,7 +289,24 @@ func performShutdown(
 
 	syncerInstance.sync(finalCtx)
 
-	time.Sleep(shutdownGracePause)
+	drainCtx, drainCancel := context.WithTimeout(
+		context.WithoutCancel(ctx), quicDrainTimeout,
+	)
+	defer drainCancel()
+
+	var drainWg sync.WaitGroup
+
+	for _, tunnelSrv := range tunnels {
+		drainWg.Add(1)
+
+		go func(srv *TunnelServer) {
+			defer drainWg.Done()
+
+			srv.waitForDrain(drainCtx)
+		}(tunnelSrv)
+	}
+
+	drainWg.Wait()
 
 	for _, tunnelSrv := range tunnels {
 		tunnelSrv.close()

internal/king/king_test.go

@@ -1,8 +1,10 @@
 package king_test
 
 import (
+	"context"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/firecow/burrow/internal/king"
 	"github.com/firecow/burrow/internal/state"
@@ -894,6 +896,84 @@ func TestOnStateChanged_ClearsServicesOnEmpty(t *testing.T) {
 	}
 }
 
+// --- WaitForDrain ---
+
+func TestWaitForDrain_EmptyImmediate(t *testing.T) {
+	t.Parallel()
+
+	tunnelSrv := king.NewTunnelServer(testBindPort, nil)
+
+	done := make(chan struct{})
+
+	go func() {
+		tunnelSrv.WaitForDrain(t.Context())
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("waitForDrain should return immediately when no connections")
+	}
+}
+
+func TestWaitForDrain_WaitsForRemoval(t *testing.T) {
+	t.Parallel()
+
+	tunnelSrv := king.NewTunnelServer(testBindPort, nil)
+	tunnelSrv.SetQUICConn("svc-1", nil)
+
+	done := make(chan struct{})
+
+	go func() {
+		tunnelSrv.WaitForDrain(t.Context())
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		t.Fatal("waitForDrain should not return while connections exist")
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	tunnelSrv.RemoveQUICConn("svc-1")
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("waitForDrain should return after connection removed")
+	}
+}
+
+func TestWaitForDrain_RespectsContextTimeout(t *testing.T) {
+	t.Parallel()
+
+	tunnelSrv := king.NewTunnelServer(testBindPort, nil)
+	tunnelSrv.SetQUICConn("svc-1", nil)
+
+	ctx, cancel := context.WithTimeout(
+		t.Context(), 100*time.Millisecond,
+	)
+	defer cancel()
+
+	done := make(chan struct{})
+
+	go func() {
+		tunnelSrv.WaitForDrain(ctx)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("waitForDrain should return on context cancellation")
+	}
+
+	if tunnelSrv.QUICConnCount() != 1 {
+		t.Fatal("connection should still exist after timeout")
+	}
+}
+
 // --- Syncer.TriggerSync ---
 
 func TestTriggerSync_SendsNotification(t *testing.T) {

internal/king/tunnel.go

@@ -55,6 +55,7 @@ type TunnelServer struct {
 	services        map[string]ServiceAuth
 	tcpListeners    map[int]net.Listener
 	quicConns       map[string]*quic.Conn
+	drainNotify     chan struct{}
 	onLingConnected func()
 }
 
@@ -69,6 +70,7 @@ func NewTunnelServer(
 		services:        make(map[string]ServiceAuth),
 		tcpListeners:    make(map[int]net.Listener),
 		quicConns:       make(map[string]*quic.Conn),
+		drainNotify:     make(chan struct{}, 1),
 		onLingConnected: nil,
 	}
 }
@@ -290,7 +292,35 @@ func (tunnelSrv *TunnelServer) unregisterConnections(
 		}
 	}
 
+	empty := len(tunnelSrv.quicConns) == 0
 	tunnelSrv.mu.Unlock()
+
+	if empty {
+		select {
+		case tunnelSrv.drainNotify <- struct{}{}:
+		default:
+		}
+	}
+}
+
+func (tunnelSrv *TunnelServer) waitForDrain(
+	ctx context.Context,
+) {
+	for {
+		tunnelSrv.mu.RLock()
+		empty := len(tunnelSrv.quicConns) == 0
+		tunnelSrv.mu.RUnlock()
+
+		if empty {
+			return
+		}
+
+		select {
+		case <-tunnelSrv.drainNotify:
+		case <-ctx.Done():
+			return
+		}
+	}
 }
 
 func (tunnelSrv *TunnelServer) ensureTCPListener(

internal/ling/ling.go

@@ -21,7 +21,6 @@ import (
 const (
 	keyValueParts           = 2
 	finalSyncTimeout        = 2 * time.Second
-	shutdownGracePeriod     = 750 * time.Millisecond
 	quicErrorCodeCloseClean = 0
 )
 
@@ -373,7 +372,6 @@ func performShutdown(
 	defer finalCancel()
 
 	syncerInstance.sync(finalCtx)
-	time.Sleep(shutdownGracePeriod)
 
 	tunnelCli.closeAll()