eot for sink

Signed-off-by: Yashash H L <[email protected]>
numaproj · Oct 14, 2024 · d393a61 · d393a61
1 parent b7c6e55
commit d393a61
Show file tree

Hide file tree

Showing 7 changed files with 167 additions and 125 deletions.
diff --git a/pkg/apis/proto/sink/v1/sink.pb.go b/pkg/apis/proto/sink/v1/sink.pb.go
diff --git a/pkg/apis/proto/sink/v1/sink.proto b/pkg/apis/proto/sink/v1/sink.proto
@@ -28,14 +28,11 @@ message SinkRequest {
     string id = 5;
     map<string, string> headers = 6;
   }
-  message Status {
-    bool eot = 1;
-  }
   // Required field indicating the request.
   Request request = 1;
   // Required field indicating the status of the request.
   // If eot is set to true, it indicates the end of transmission.
-  Status status = 2;
+  TransmissionStatus status = 2;
   // optional field indicating the handshake message.
   optional Handshake handshake = 3;
 }
@@ -55,6 +52,13 @@ message ReadyResponse {
   bool ready = 1;
 }
 
+/**
+ * TransmissionStatus is the status of the transmission.
+ */
+message TransmissionStatus {
+  bool eot = 1;
+}
+
 /*
   * Status is the status of the response.
  */
@@ -78,4 +82,5 @@ message SinkResponse {
   }
   Result result = 1;
   optional Handshake handshake = 2;
+  optional TransmissionStatus status = 3;
 }
diff --git a/pkg/mapstreamer/examples/flatmap_stream/go.mod b/pkg/mapstreamer/examples/flatmap_stream/go.mod
@@ -8,6 +8,7 @@ require github.com/numaproj/numaflow-go v0.8.1
 
 require (
 	golang.org/x/net v0.29.0 // indirect
+	golang.org/x/sync v0.8.0 // indirect
 	golang.org/x/sys v0.25.0 // indirect
 	golang.org/x/text v0.18.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect

diff --git a/pkg/mapstreamer/examples/flatmap_stream/go.sum b/pkg/mapstreamer/examples/flatmap_stream/go.sum
@@ -8,6 +8,8 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 golang.org/x/net v0.29.0 h1:5ORfpBpCs4HzDYoodCDBbwHzdR5UrLBZ3sOnUJmFoHo=
 golang.org/x/net v0.29.0/go.mod h1:gLkgy8jTGERgjzMic6DS9+SP0ajcu6Xu3Orq/SpETg0=
+golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
+golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34=
 golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.18.0 h1:XvMDiNzPAl0jr17s6W9lcaIhGUfUORdGCNsuLmPG224=

diff --git a/pkg/mapstreamer/service.go b/pkg/mapstreamer/service.go
@@ -87,8 +87,6 @@ func (fs *Service) invokeHandler(ctx context.Context, req *mappb.MapRequest, mes
 			err = fmt.Errorf("panic inside mapStream handler: %v", r)
 			return
 		}
-		// close the message channel after the handler is done processing the request
-		close(messageCh)
 	}()
 	streamReq := req.GetRequest()
 	hd := NewHandlerDatum(streamReq.GetValue(), streamReq.GetEventTime().AsTime(), streamReq.GetWatermark().AsTime(), streamReq.GetHeaders())

diff --git a/pkg/sinker/service.go b/pkg/sinker/service.go
@@ -197,5 +197,19 @@ func (fs *Service) processData(ctx context.Context, stream sinkpb.Sink_SinkFnSer
 			return err
 		}
 	}
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+	// send the end of transmission message
+	eotResponse := &sinkpb.SinkResponse{
+		Status: &sinkpb.TransmissionStatus{Eot: true},
+	}
+	if err := stream.Send(eotResponse); err != nil {
+		log.Printf("error sending end of transmission message: %v", err)
+		return err
+	}
 	return nil
 }
diff --git a/pkg/sinker/service_test.go b/pkg/sinker/service_test.go
@@ -90,7 +90,7 @@ func TestService_SinkFn(t *testing.T) {
 				},
 				{
 					Request: &sinkpb.SinkRequest_Request{},
-					Status:  &sinkpb.SinkRequest_Status{Eot: true},
+					Status:  &sinkpb.TransmissionStatus{Eot: true},
 				},
 			},
 			sh: SinkerFunc(func(ctx context.Context, rch <-chan Datum) Responses {
@@ -129,6 +129,9 @@ func TestService_SinkFn(t *testing.T) {
 						ErrMsg: "",
 					},
 				},
+				{
+					Status: &sinkpb.TransmissionStatus{Eot: true},
+				},
 			},
 		},
 		{
@@ -171,7 +174,7 @@ func TestService_SinkFn(t *testing.T) {
 				},
 				{
 					Request: &sinkpb.SinkRequest_Request{},
-					Status:  &sinkpb.SinkRequest_Status{Eot: true},
+					Status:  &sinkpb.TransmissionStatus{Eot: true},
 				},
 			},
 			sh: SinkerFunc(func(ctx context.Context, rch <-chan Datum) Responses {
@@ -210,6 +213,9 @@ func TestService_SinkFn(t *testing.T) {
 						ErrMsg: "unknown error",
 					},
 				},
+				{
+					Status: &sinkpb.TransmissionStatus{Eot: true},
+				},
 			},
 		},
 	}