Repository: numaproj/numaflow

Issue:
Sink stops when we have multiple partitions
I have multiple partitions in my pipeline and I am consuming batches of messages.
In the numa container there is a messgae that EOT har reached it stops consuming.
Removing the partitions fixed the problem.

Numaflow: 1.4.3
Numaflow-java: 0.9.2
Jetstream: 2.10.3

Hints from issue comments:
@yhl25: @borkox Thanks for reporting the bug, I was able to reproduce the issue. I will have a fix for it in the next release.@yhl25: Fix was released in [v1.4.4](https://github.com/numaproj/numaflow/releases/tag/v1.4.4).@vigith: We have to fix this in main too.@yhl25: > We have to fix this in main too.

https://github.com/numaproj/numaflow/pull/2521

Solution patch:
diff --git a/pkg/sinks/sink.go b/pkg/sinks/sink.go
index 18702074..0800a5a7 100644
--- a/pkg/sinks/sink.go
+++ b/pkg/sinks/sink.go
@@ -65,7 +65,7 @@ func (u *SinkProcessor) Start(ctx context.Context) error {
 		fromVertexWmStores map[string]store.WatermarkStore
 		sinkWmStores       map[string]store.WatermarkStore
 		idleManager        wmb.IdleManager
-		sinkHandler        *udsink.UDSgRPCBasedUDSink
+		sinkHandlers       []*udsink.UDSgRPCBasedUDSink
 		fbSinkHandler      *udsink.UDSgRPCBasedUDSink
 		healthCheckers     = make([]metrics.HealthChecker, 0)
 		vertexName         = u.VertexInstance.Vertex.Spec.Name
@@ -156,26 +156,34 @@ func (u *SinkProcessor) Start(ctx context.Context) error {
 		}
 		metrics.SDKInfo.WithLabelValues(dfv1.ComponentVertex, fmt.Sprintf("%s-%s", pipelineName, vertexName), string(serverinfo.ContainerTypeSinker), serverInfo.Version, string(serverInfo.Language)).Set(1)
 
-		sdkClient, err := sinkclient.New(ctx, serverInfo, sdkclient.WithMaxMessageSize(maxMessageSize))
-		if err != nil {
-			return fmt.Errorf("failed to create sdk client, %w", err)
-		}
-
-		sinkHandler = udsink.NewUDSgRPCBasedUDSink(sdkClient)
-		// Readiness check
-		if err := sinkHandler.WaitUntilReady(ctx); err != nil {
-			return fmt.Errorf("failed on UDSink readiness check, %w", err)
-		}
+		// we should create a new sdk client for each partition, because each client is mapped to one bidirectional stream
+		// and the bidi stream cannot be shared.
+		for index := range u.VertexInstance.Vertex.OwnedBuffers() {
+			sdkClient, err := sinkclient.New(ctx, serverInfo, sdkclient.WithMaxMessageSize(maxMessageSize))
+			if err != nil {
+				return fmt.Errorf("failed to create sdk client, %w", err)
+			}
 
-		// add sinkHandler to healthCheckers
-		healthCheckers = append(healthCheckers, sinkHandler)
+			sinkHandler := udsink.NewUDSgRPCBasedUDSink(sdkClient)
+			// Readiness check
+			if err := sinkHandler.WaitUntilReady(ctx); err != nil {
+				return fmt.Errorf("failed on UDSink readiness check, %w", err)
+			}
 
-		defer func() {
-			err = sdkClient.CloseConn(ctx)
-			if err != nil {
-				log.Warnw("Failed to close sdk client", zap.Error(err))
+			// add sinkHandler to healthCheckers
+			if index == 0 {
+				healthCheckers = append(healthCheckers, sinkHandler)
 			}
-		}()
+
+			sinkHandlers = append(sinkHandlers, sinkHandler)
+
+			defer func() {
+				err = sdkClient.CloseConn(ctx)
+				if err != nil {
+					log.Warnw("Failed to close sdk client", zap.Error(err))
+				}
+			}()
+		}
 	}
 
 	if u.VertexInstance.Vertex.HasFallbackUDSink() {
@@ -220,7 +228,7 @@ func (u *SinkProcessor) Start(ctx context.Context) error {
 		}
 
 		// create the main sink writer
-		sinkWriter, err := u.createSinkWriter(ctx, &u.VertexInstance.Vertex.Spec.Sink.AbstractSink, sinkHandler)
+		sinkWriter, err := u.createSinkWriter(ctx, &u.VertexInstance.Vertex.Spec.Sink.AbstractSink, sinkHandlers[index])
 		if err != nil {
 			return fmt.Errorf("failed to find a sink, error: %w", err)
 		}



Test patch:
diff --git a/test/e2e/testdata/simple-pipeline.yaml b/test/e2e/testdata/simple-pipeline.yaml
index 0ddacd96..194b5237 100644
--- a/test/e2e/testdata/simple-pipeline.yaml
+++ b/test/e2e/testdata/simple-pipeline.yaml
@@ -10,10 +10,12 @@ spec:
           rpu: 5
           duration: 1s
     - name: p1
+      partitions: 2
       udf:
         builtin:
           name: cat
     - name: output
+      partitions: 3
       sink:
         log: {}
   edges: