diff --git a/management/internals/shared/grpc/proxy.go b/management/internals/shared/grpc/proxy.go index 8e93f77b0..439ff1dba 100644 --- a/management/internals/shared/grpc/proxy.go +++ b/management/internals/shared/grpc/proxy.go @@ -152,6 +152,10 @@ type ProxyServiceServer struct { // initial snapshot delivery. Configurable via NB_PROXY_SNAPSHOT_BATCH_SIZE. snapshotBatchSize int + // snapshotAckTimeout bounds the wait for each snapshot batch ack. + // Defaults to defaultSnapshotAckTimeout when zero. + snapshotAckTimeout time.Duration + cancel context.CancelFunc } @@ -168,6 +172,8 @@ const defaultProxyTokenTTL = 5 * time.Minute const defaultSnapshotBatchSize = 500 +const defaultSnapshotAckTimeout = 5 * time.Minute + func snapshotBatchSizeFromEnv() int { if v := os.Getenv("NB_PROXY_SNAPSHOT_BATCH_SIZE"); v != "" { if n, err := strconv.Atoi(v); err == nil && n > 0 { @@ -185,6 +191,14 @@ func (s *ProxyServiceServer) proxyTokenTTL() time.Duration { return defaultProxyTokenTTL } +// snapshotAckWait returns the configured snapshot ack timeout or the default when unset. +func (s *ProxyServiceServer) snapshotAckWait() time.Duration { + if s.snapshotAckTimeout > 0 { + return s.snapshotAckTimeout + } + return defaultSnapshotAckTimeout +} + // proxyConnection represents a connected proxy type proxyConnection struct { proxyID string @@ -484,7 +498,7 @@ func (s *ProxyServiceServer) SyncMappings(stream proto.ProxyService_SyncMappings return err } - if err := s.sendSnapshotSync(stream.Context(), conn, stream); err != nil { + if err := s.sendSnapshotSync(conn.ctx, conn, stream); err != nil { s.cleanupFailedSnapshot(stream.Context(), conn) return fmt.Errorf("send snapshot to proxy %s: %w", params.proxyID, err) } @@ -728,7 +742,7 @@ func (s *ProxyServiceServer) sendSnapshotSync(ctx context.Context, conn *proxyCo return fmt.Errorf("send snapshot batch: %w", err) } - if err := waitForAck(stream); err != nil { + if err := waitForAck(ctx, stream, s.snapshotAckWait()); err != nil { return err } } @@ -740,7 +754,7 @@ func (s *ProxyServiceServer) sendSnapshotSync(ctx context.Context, conn *proxyCo return fmt.Errorf("send snapshot completion: %w", err) } - if err := waitForAck(stream); err != nil { + if err := waitForAck(ctx, stream, s.snapshotAckWait()); err != nil { return err } } @@ -748,15 +762,38 @@ func (s *ProxyServiceServer) sendSnapshotSync(ctx context.Context, conn *proxyCo return nil } -func waitForAck(stream proto.ProxyService_SyncMappingsServer) error { - msg, err := stream.Recv() - if err != nil { - return fmt.Errorf("receive ack: %w", err) +// waitForAck waits for the proxy to ack a snapshot batch. It gives up when ctx +// is canceled or no ack arrives within timeout, since a proxy that abandoned the +// stream without closing it would otherwise block the handler indefinitely. +func waitForAck(ctx context.Context, stream proto.ProxyService_SyncMappingsServer, timeout time.Duration) error { + type recvResult struct { + msg *proto.SyncMappingsRequest + err error } - if msg.GetAck() == nil { - return fmt.Errorf("expected ack, got %T", msg.GetMsg()) + // Recv returns once the handler exits and gRPC tears the stream down. + result := make(chan recvResult, 1) + go func() { + msg, err := stream.Recv() + result <- recvResult{msg: msg, err: err} + }() + + timer := time.NewTimer(timeout) + defer timer.Stop() + + select { + case r := <-result: + if r.err != nil { + return fmt.Errorf("receive ack: %w", r.err) + } + if r.msg.GetAck() == nil { + return fmt.Errorf("expected ack, got %T", r.msg.GetMsg()) + } + return nil + case <-ctx.Done(): + return fmt.Errorf("wait for ack: %w", ctx.Err()) + case <-timer.C: + return fmt.Errorf("no ack within %s", timeout) } - return nil } // heartbeat updates the proxy's last_seen timestamp every minute and diff --git a/management/internals/shared/grpc/sync_mappings_test.go b/management/internals/shared/grpc/sync_mappings_test.go index 6949b79e9..753142ea0 100644 --- a/management/internals/shared/grpc/sync_mappings_test.go +++ b/management/internals/shared/grpc/sync_mappings_test.go @@ -409,3 +409,75 @@ func TestConnectionSendResponse_RoutesToLegacyStream(t *testing.T) { require.Len(t, stream.messages, 1) assert.Equal(t, "svc-2", stream.messages[0].Mapping[0].Id) } + +// blockingSyncStream is a ProxyService_SyncMappingsServer whose Recv blocks +// until the test ends, like a proxy that stopped acking without closing the stream. +type blockingSyncStream struct { + grpc.ServerStream + done chan struct{} +} + +func (s *blockingSyncStream) Recv() (*proto.SyncMappingsRequest, error) { + <-s.done + return nil, context.Canceled +} + +func (s *blockingSyncStream) Send(*proto.SyncMappingsResponse) error { return nil } + +func newBlockingSyncStream(t *testing.T) *blockingSyncStream { + t.Helper() + s := &blockingSyncStream{done: make(chan struct{})} + t.Cleanup(func() { close(s.done) }) + return s +} + +func TestWaitForAck_Ack(t *testing.T) { + stream := &syncRecordingStream{recvMsgs: []*proto.SyncMappingsRequest{ackMsg()}} + + require.NoError(t, waitForAck(context.Background(), stream, time.Second)) +} + +func TestWaitForAck_UnexpectedMessage(t *testing.T) { + stream := &syncRecordingStream{recvMsgs: []*proto.SyncMappingsRequest{ + {Msg: &proto.SyncMappingsRequest_Init{Init: &proto.SyncMappingsInit{}}}, + }} + + assert.ErrorContains(t, waitForAck(context.Background(), stream, time.Second), "expected ack") +} + +func TestWaitForAck_TimesOutWhenProxyStopsAcking(t *testing.T) { + stream := newBlockingSyncStream(t) + + start := time.Now() + err := waitForAck(context.Background(), stream, 50*time.Millisecond) + + assert.ErrorContains(t, err, "no ack within") + assert.Less(t, time.Since(start), time.Second) +} + +func TestWaitForAck_ReturnsWhenConnectionCanceled(t *testing.T) { + stream := newBlockingSyncStream(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + assert.ErrorIs(t, waitForAck(ctx, stream, time.Minute), context.Canceled) +} + +func TestSendSnapshotSync_GivesUpWhenProxyStopsAcking(t *testing.T) { + const cluster = "cluster.example.com" + + ctrl := gomock.NewController(t) + mgr := rpservice.NewMockManager(ctrl) + mgr.EXPECT().GetGlobalServices(gomock.Any()).Return(makeServices(1, cluster), nil) + + s := newSnapshotTestServer(t, 10) + s.serviceManager = mgr + s.snapshotAckTimeout = 50 * time.Millisecond + + stream := newBlockingSyncStream(t) + conn := &proxyConnection{proxyID: "proxy-a", address: cluster, syncStream: stream} + + err := s.sendSnapshotSync(context.Background(), conn, stream) + + assert.ErrorContains(t, err, "no ack within") +} diff --git a/proxy/server.go b/proxy/server.go index 9ad81b50f..3bba38644 100644 --- a/proxy/server.go +++ b/proxy/server.go @@ -1303,6 +1303,10 @@ func (s *Server) proxyCapabilities() *proto.ProxyCapabilities { } func (s *Server) tryGetMappingUpdate(ctx context.Context, client proto.ProxyServiceClient, initialSyncDone *bool, onConnected func()) error { + // Canceling closes the stream on management; a bare return leaves it open. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + connectTime := time.Now() mappingClient, err := client.GetMappingUpdate(ctx, &proto.GetMappingUpdateRequest{ ProxyId: s.ID, @@ -1325,6 +1329,10 @@ func (s *Server) tryGetMappingUpdate(ctx context.Context, client proto.ProxyServ } func (s *Server) trySyncMappings(ctx context.Context, client proto.ProxyServiceClient, initialSyncDone *bool, onConnected func()) error { + // Canceling closes the stream on management; a bare return leaves it open. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + connectTime := time.Now() stream, err := client.SyncMappings(ctx) if err != nil {