[management] fix proxy conn lifecycle during acks (#8220)

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 {
+// waitForAck waits for the proxy to ack a snapshot batch. It gives up when ctx
-	msg, err := stream.Recv()
+// is canceled or no ack arrives within timeout, since a proxy that abandoned the
-	if err != nil {
+// stream without closing it would otherwise block the handler indefinitely.
-		return fmt.Errorf("receive ack: %w", err)
+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 {
+	// Recv returns once the handler exits and gRPC tears the stream down.
-		return fmt.Errorf("expected ack, got %T", msg.GetMsg())
+	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

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")
+}

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 {