diff --git a/management/internals/shared/grpc/proxy.go b/management/internals/shared/grpc/proxy.go index cd3d517ca..15c68dabe 100644 --- a/management/internals/shared/grpc/proxy.go +++ b/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. diff --git a/proxy/sync_mappings_test.go b/proxy/sync_mappings_test.go index 63a1d42f2..801587e4c 100644 --- a/proxy/sync_mappings_test.go +++ b/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) } }