fix error handling

management/internals/shared/grpc/proxy.go

@@ -9,6 +9,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"net/url"
@@ -461,9 +462,11 @@ func (s *ProxyServiceServer) SyncMappings(stream proto.ProxyService_SyncMappings
 	errChan := make(chan error, 2)
 	go s.sender(conn, errChan)
 
-	// Drain acks from the proxy in the background so the stream stays healthy.
-	// After the snapshot phase, the proxy may still send acks for incremental
-	// updates; we simply discard them.
+	// Drain acks from the proxy in the background so the stream stays
+	// healthy. The proxy sends an ack for every message it receives
+	// (including incremental updates); we discard them here.
+	// EOF or context cancellation is a normal shutdown path and is
+	// forwarded so the select below can clean up.
 	go func() {
 		for {
 			if _, err := stream.Recv(); err != nil {
@@ -503,6 +506,10 @@ func (s *ProxyServiceServer) SyncMappings(stream proto.ProxyService_SyncMappings
 
 	select {
 	case err := <-errChan:
+		if isStreamClosed(err) {
+			log.WithContext(ctx).Infof("Proxy %s stream closed", proxyID)
+			return nil
+		}
 		log.WithContext(ctx).Warnf("Failed to send update: %v", err)
 		return fmt.Errorf("send update to proxy %s: %w", proxyID, err)
 	case <-connCtx.Done():
@@ -517,6 +524,9 @@ func (s *ProxyServiceServer) sendSnapshotSync(ctx context.Context, conn *proxyCo
 	if !isProxyAddressValid(conn.address) {
 		return fmt.Errorf("proxy address is invalid")
 	}
+	if s.snapshotBatchSize <= 0 {
+		return fmt.Errorf("invalid snapshot batch size: %d", s.snapshotBatchSize)
+	}
 
 	mappings, err := s.snapshotServiceMappings(ctx, conn)
 	if err != nil {
@@ -601,6 +611,9 @@ func (s *ProxyServiceServer) sendSnapshot(ctx context.Context, conn *proxyConnec
 	if !isProxyAddressValid(conn.address) {
 		return fmt.Errorf("proxy address is invalid")
 	}
+	if s.snapshotBatchSize <= 0 {
+		return fmt.Errorf("invalid snapshot batch size: %d", s.snapshotBatchSize)
+	}
 
 	mappings, err := s.snapshotServiceMappings(ctx, conn)
 	if err != nil {
@@ -680,6 +693,18 @@ func isProxyAddressValid(addr string) bool {
 	return err == nil
 }
 
+// isStreamClosed returns true for errors that indicate normal stream
+// termination: io.EOF, context cancellation, or gRPC Canceled.
+func isStreamClosed(err error) bool {
+	if err == nil {
+		return false
+	}
+	if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) {
+		return true
+	}
+	return status.Code(err) == codes.Canceled
+}
+
 // sender handles sending messages to proxy.
 // When conn.syncStream is set the message is sent as SyncMappingsResponse;
 // otherwise the legacy GetMappingUpdateResponse stream is used.

proxy/sync_mappings_test.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -86,8 +85,8 @@ func TestIntegration_SyncMappings_BackPressure(t *testing.T) {
 	setup := setupIntegrationTest(t)
 	defer setup.cleanup()
 
-	// Add more services so we get multiple batches.
-	addServicesToStore(t, setup, 20, "test.proxy.io")
+	// Add enough services to guarantee multiple batches (default batch size 500).
+	addServicesToStore(t, setup, 600, "test.proxy.io")
 
 	conn, err := grpc.NewClient(setup.grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	require.NoError(t, err)
@@ -112,26 +111,40 @@ func TestIntegration_SyncMappings_BackPressure(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	// Record the ordering of events to verify back-pressure.
-	var mu sync.Mutex
-	var events []string
+	// Strategy: receive batch 1, then hold for a significant delay before
+	// acking. If back-pressure works, batch 2 cannot arrive until after
+	// the ack is sent — so its receive timestamp must be >= the ack
+	// timestamp. If management were fire-and-forget, all batches would
+	// already be buffered in the gRPC transport and batch 2 would arrive
+	// well before the ack time.
+	const ackDelay = 300 * time.Millisecond
+
+	type batchEvent struct {
+		recvAt time.Time
+		ackAt  time.Time
+		count  int
+	}
+	var batches []batchEvent
 	var totalMappings int
 
 	for {
 		msg, err := stream.Recv()
 		require.NoError(t, err)
 
-		mu.Lock()
-		events = append(events, "recv")
+		recvAt := time.Now()
 		totalMappings += len(msg.GetMapping())
-		mu.Unlock()
 
-		// Simulate processing delay.
-		time.Sleep(50 * time.Millisecond)
+		// Delay the ack on non-final batches to create a measurable gap.
+		if !msg.GetInitialSyncComplete() {
+			time.Sleep(ackDelay)
+		}
 
-		mu.Lock()
-		events = append(events, "ack")
-		mu.Unlock()
+		ackAt := time.Now()
+		batches = append(batches, batchEvent{
+			recvAt: recvAt,
+			ackAt:  ackAt,
+			count:  len(msg.GetMapping()),
+		})
 
 		err = stream.Send(&proto.SyncMappingsRequest{
 			Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},
@@ -143,18 +156,20 @@ func TestIntegration_SyncMappings_BackPressure(t *testing.T) {
 		}
 	}
 
-	// 2 original + 20 added = 22 services total.
-	assert.Equal(t, 22, totalMappings, "should receive all 22 mappings")
+	// 2 original + 600 added = 602 services total.
+	assert.Equal(t, 602, totalMappings, "should receive all 602 mappings")
+	require.GreaterOrEqual(t, len(batches), 2, "need at least 2 batches to verify back-pressure")
 
-	// Events should alternate recv/ack — no two recvs in a row
-	// (management waits for ack before sending next).
-	mu.Lock()
-	defer mu.Unlock()
-	for i := 0; i < len(events)-1; i += 2 {
-		assert.Equal(t, "recv", events[i], "event %d should be recv", i)
-		if i+1 < len(events) {
-			assert.Equal(t, "ack", events[i+1], "event %d should be ack", i+1)
-		}
+	// For every batch after the first, its receive time must be after the
+	// previous batch's ack time. This proves management waited for the ack
+	// before sending the next batch.
+	for i := 1; i < len(batches); i++ {
+		prevAckAt := batches[i-1].ackAt
+		thisRecvAt := batches[i].recvAt
+		assert.True(t, !thisRecvAt.Before(prevAckAt),
+			"batch %d received at %v, but batch %d was acked at %v — "+
+				"management sent the next batch before receiving the ack",
+			i, thisRecvAt, i-1, prevAckAt)
 	}
 }