Repository: numaproj/numaflow
Pull request #762 (2023-05-31 23:43:22)

chore: fix idle watermark bug
![Screenshot 2023-06-01 at 2 08 00 PM](https://github.com/numaproj/numaflow/assets/109710325/ada9d7ca-2e2b-41da-812e-11a2200a9500)

fixes https://github.com/numaproj/numaflow/issues/748#issuecomment-1570702145


Patch:
diff --git a/pkg/forward/forward.go b/pkg/forward/forward.go
index d06dfadf..d5d02f9e 100644
--- a/pkg/forward/forward.go
+++ b/pkg/forward/forward.go
@@ -52,14 +52,15 @@ type InterStepDataForward struct {
 	cancelFn   context.CancelFunc
 	fromBuffer isb.BufferReader
 	// toBuffers is a map of toVertex name to the toVertex's owned buffers.
-	toBuffers        map[string][]isb.BufferWriter
-	FSD              ToWhichStepDecider
-	UDF              applier.MapApplier
-	fetchWatermark   fetch.Fetcher
-	publishWatermark map[string]publish.Publisher
-	opts             options
-	vertexName       string
-	pipelineName     string
+	toBuffers map[string][]isb.BufferWriter
+	FSD       ToWhichStepDecider
+	UDF       applier.MapApplier
+	wmFetcher fetch.Fetcher
+	// wmPublishers stores the vertex to publisher mapping
+	wmPublishers map[string]publish.Publisher
+	opts         options
+	vertexName   string
+	pipelineName string
 	// idleManager manages the idle watermark status.
 	idleManager *wmb.IdleManager
 	// wmbChecker checks if the idle watermark is valid.
@@ -87,14 +88,14 @@ func NewInterStepDataForward(vertex *dfv1.Vertex,
 	ctx, cancel := context.WithCancel(context.Background())
 
 	var isdf = InterStepDataForward{
-		ctx:              ctx,
-		cancelFn:         cancel,
-		fromBuffer:       fromStep,
-		toBuffers:        toSteps,
-		FSD:              fsd,
-		UDF:              applyUDF,
-		fetchWatermark:   fetchWatermark,
-		publishWatermark: publishWatermark,
+		ctx:          ctx,
+		cancelFn:     cancel,
+		fromBuffer:   fromStep,
+		toBuffers:    toSteps,
+		FSD:          fsd,
+		UDF:          applyUDF,
+		wmFetcher:    fetchWatermark,
+		wmPublishers: publishWatermark,
 		// should we do a check here for the values not being null?
 		vertexName:   vertex.Spec.Name,
 		pipelineName: vertex.Spec.PipelineName,
@@ -175,12 +176,12 @@ func (isdf *InterStepDataForward) Start() <-chan struct{} {
 		}
 
 		// stop watermark fetcher
-		if err := isdf.fetchWatermark.Close(); err != nil {
+		if err := isdf.wmFetcher.Close(); err != nil {
 			log.Errorw("Failed to close watermark fetcher", zap.Error(err))
 		}
 
 		// stop watermark publisher
-		for _, publisher := range isdf.publishWatermark {
+		for _, publisher := range isdf.wmPublishers {
 			if err := publisher.Close(); err != nil {
 				log.Errorw("Failed to close watermark publisher", zap.Error(err))
 			}
@@ -224,7 +225,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 		// We also publish a control message if this is the first time we get this idle situation.
 
 		// we use the HeadWMB as the watermark for the idle
-		var processorWMB = isdf.fetchWatermark.GetHeadWMB()
+		var processorWMB = isdf.wmFetcher.GetHeadWMB()
 		if !isdf.wmbChecker.ValidateHeadWMB(processorWMB) {
 			// validation failed, skip publishing
 			isdf.opts.logger.Debugw("skip publishing idle watermark",
@@ -237,8 +238,10 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 
 		// TODO(multi-partition) handle for multiple partitions
 		for toVertexName, toVertexBuffer := range isdf.toBuffers {
-			if p, ok := isdf.publishWatermark[toVertexName]; ok {
-				idlehandler.PublishIdleWatermark(ctx, toVertexBuffer[0], p, isdf.idleManager, 0, isdf.opts.logger, isdf.opts.vertexType, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+			for index, partition := range toVertexBuffer {
+				if p, ok := isdf.wmPublishers[toVertexName]; ok {
+					idlehandler.PublishIdleWatermark(ctx, partition, p, isdf.idleManager, int32(index), isdf.opts.logger, isdf.opts.vertexType, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+				}
 			}
 		}
 		return
@@ -266,7 +269,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 		// TODO: make it async (concurrent and wait later)
 		// let's track only the first element's watermark. This is important because we reassign the watermark we fetch
 		// to all the elements in the batch. If we were to assign last element's watermark, we will wrongly mark on-time data as late.
-		processorWM = isdf.fetchWatermark.GetWatermark(readMessages[0].ReadOffset)
+		processorWM = isdf.wmFetcher.GetWatermark(readMessages[0].ReadOffset)
 	}
 
 	var writeOffsets map[string][][]isb.Offset
@@ -334,7 +337,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 			isdf.opts.srcWatermarkPublisher.PublishSourceWatermarks(transformedReadMessages)
 			// fetch the source watermark again, we might not get the latest watermark because of publishing delay,
 			// but ideally we should use the latest to determine the IsLate attribute.
-			processorWM = isdf.fetchWatermark.GetWatermark(readMessages[0].ReadOffset)
+			processorWM = isdf.wmFetcher.GetWatermark(readMessages[0].ReadOffset)
 			// assign isLate
 			for _, m := range writeMessages {
 				if processorWM.After(m.EventTime) { // Set late data at source level
@@ -385,24 +388,26 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 	// activeWatermarkBuffers records the buffers that the publisher has published
 	// a watermark in this batch processing cycle.
 	// it's used to determine which buffers should receive an idle watermark.
-	var activeWatermarkBuffers = make(map[string]bool)
+	// It is created as a slice because it tracks per partition activity info.
+	var activeWatermarkBuffers = make(map[string][]bool)
 	// forward the highest watermark to all the edges to avoid idle edge problem
 	// TODO: sort and get the highest value
 	for toVertexName, toVertexBufferOffsets := range writeOffsets {
-		if publisher, ok := isdf.publishWatermark[toVertexName]; ok {
+		activeWatermarkBuffers[toVertexName] = make([]bool, len(toVertexBufferOffsets))
+		if publisher, ok := isdf.wmPublishers[toVertexName]; ok {
 			for index, offsets := range toVertexBufferOffsets {
 				if isdf.opts.vertexType == dfv1.VertexTypeSource || isdf.opts.vertexType == dfv1.VertexTypeMapUDF ||
 					isdf.opts.vertexType == dfv1.VertexTypeReduceUDF {
 					if len(offsets) > 0 {
 						publisher.PublishWatermark(processorWM, offsets[len(offsets)-1], int32(index))
-						activeWatermarkBuffers[toVertexName] = true
+						activeWatermarkBuffers[toVertexName][index] = true
 						// reset because the toBuffer is no longer idling
 						isdf.idleManager.Reset(toVertexName)
 					}
 					// This (len(offsets) == 0) happens at conditional forwarding, there's no data written to the buffer
 				} else { // For Sink vertex, and it does not care about the offset during watermark publishing
 					publisher.PublishWatermark(processorWM, nil, int32(index))
-					activeWatermarkBuffers[toVertexName] = true
+					activeWatermarkBuffers[toVertexName][index] = true
 					// reset because the toBuffer is no longer idling
 					isdf.idleManager.Reset(toVertexName)
 				}
@@ -410,20 +415,24 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 		}
 	}
 	// TODO(multi-partition): handle idle watermark publishing for multi partitioned buffer
-	if len(dataMessages) > 0 && len(activeWatermarkBuffers) < len(isdf.publishWatermark) {
-		// - condition1 "len(dataMessages) > 0" :
-		//   Meaning, we do have some data messages, but not all out buffers get written.
-		//   Tt could be all data messages are dropped, or conditional forwarding to part of the out buffers.
-		//   If we don't have this condition check, when dataMessages is zero but ctrlMessages > 0, we will
-		//   wrongly publish an idle watermark without the ctrl message and the ctrl message tracking map.
-		// - condition 2 "len(activeWatermarkBuffers) < len(isdf.publishWatermark)" :
-		//   send idle watermark only if we have idle out buffers
-		for bufferName := range isdf.publishWatermark {
-			if !activeWatermarkBuffers[bufferName] {
-				// use the watermark of the current read batch for the idle watermark
-				// same as read len==0 because there's no event published to the buffer
-				if p, ok := isdf.publishWatermark[bufferName]; ok {
-					idlehandler.PublishIdleWatermark(ctx, isdf.toBuffers[bufferName][0], p, isdf.idleManager, 0, isdf.opts.logger, isdf.opts.vertexType, processorWM)
+	// - condition1 "len(dataMessages) > 0" :
+	//   Meaning, we do have some data messages, but we may not have written to all out buffers or its partitions.
+	//   It could be all data messages are dropped, or conditional forwarding to part of the out buffers.
+	//   If we don't have this condition check, when dataMessages is zero but ctrlMessages > 0, we will
+	//   wrongly publish an idle watermark without the ctrl message and the ctrl message tracking map.
+	// - condition 2 "len(activeWatermarkBuffers) < len(isdf.wmPublishers)" :
+	//   send idle watermark only if we have idle out buffers
+	// Note: When the len(dataMessages) is 0, meaning all the readMessages are control messages, we choose not to do extra steps
+	// This is because, if the idle continues, we will eventually handle the idle watermark when we read the next batch where the len(readMessages) will be zero
+	if len(dataMessages) > 0 {
+		for bufferName := range isdf.wmPublishers {
+			for index, activePartition := range activeWatermarkBuffers[bufferName] {
+				if !activePartition {
+					// use the watermark of the current read batch for the idle watermark
+					// same as read len==0 because there's no event published to the buffer
+					if p, ok := isdf.wmPublishers[bufferName]; ok {
+						idlehandler.PublishIdleWatermark(ctx, isdf.toBuffers[bufferName][index], p, isdf.idleManager, int32(index), isdf.opts.logger, isdf.opts.vertexType, processorWM)
+					}
 				}
 			}
 		}
diff --git a/pkg/reduce/data_forward.go b/pkg/reduce/data_forward.go
index a9ef2b10..303d79a5 100644
--- a/pkg/reduce/data_forward.go
+++ b/pkg/reduce/data_forward.go
@@ -61,8 +61,8 @@ type DataForward struct {
 	vertexReplica         int32
 	fromBuffer            isb.BufferReader
 	toBuffers             map[string][]isb.BufferWriter
-	watermarkFetcher      fetch.Fetcher
-	watermarkPublishers   map[string]publish.Publisher
+	wmFetcher             fetch.Fetcher
+	wmPublishers          map[string]publish.Publisher
 	windower              window.Windower
 	keyed                 bool
 	idleManager           *wmb.IdleManager
@@ -104,8 +104,8 @@ func NewDataForward(ctx context.Context,
 		vertexReplica:         vertexInstance.Replica,
 		fromBuffer:            fromBuffer,
 		toBuffers:             toBuffers,
-		watermarkFetcher:      fw,
-		watermarkPublishers:   watermarkPublishers,
+		wmFetcher:             fw,
+		wmPublishers:          watermarkPublishers,
 		windower:              windowingStrategy,
 		keyed:                 vertexInstance.Vertex.Spec.UDF.GroupBy.Keyed,
 		idleManager:           idleManager,
@@ -192,7 +192,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 
 	if len(readMessages) == 0 {
 		// we use the HeadWMB as the watermark for the idle
-		var processorWMB = df.watermarkFetcher.GetHeadWMB()
+		var processorWMB = df.wmFetcher.GetHeadWMB()
 		if !df.wmbChecker.ValidateHeadWMB(processorWMB) {
 			// validation failed, skip publishing
 			df.log.Debugw("skip publishing idle watermark",
@@ -210,8 +210,10 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 			// in this case, send idle watermark to all the toBuffers
 			// TODO(multi-partition): support multi partitioned buffer
 			for toVertexName, toVertexBuffer := range df.toBuffers {
-				if publisher, ok := df.watermarkPublishers[toVertexName]; ok {
-					idlehandler.PublishIdleWatermark(ctx, toVertexBuffer[0], publisher, df.idleManager, 0, df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+				for index, bufferPartition := range toVertexBuffer {
+					if publisher, ok := df.wmPublishers[toVertexName]; ok {
+						idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, int32(index), df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+					}
 				}
 			}
 		} else {
@@ -229,8 +231,10 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 				// so in this case, we publish an idle watermark
 				// TODO(multi-partition): support multi partitioned edges
 				for toVertexName, toVertexBuffer := range df.toBuffers {
-					if publisher, ok := df.watermarkPublishers[toVertexName]; ok {
-						idlehandler.PublishIdleWatermark(ctx, toVertexBuffer[0], publisher, df.idleManager, 0, df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(watermark))
+					for index, bufferPartition := range toVertexBuffer {
+						if publisher, ok := df.wmPublishers[toVertexName]; ok {
+							idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, int32(index), df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(watermark))
+						}
 					}
 				}
 			}
@@ -245,7 +249,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 	}).Add(float64(len(readMessages)))
 	// fetch watermark using the first element's watermark, because we assign the watermark to all other
 	// elements in the batch based on the watermark we fetch from 0th offset.
-	processorWM := df.watermarkFetcher.GetWatermark(readMessages[0].ReadOffset)
+	processorWM := df.wmFetcher.GetWatermark(readMessages[0].ReadOffset)
 	for _, m := range readMessages {
 		if !df.keyed {
 			m.Keys = []string{dfv1.DefaultKeyForNonKeyedData}
@@ -358,8 +362,10 @@ func (df *DataForward) Process(ctx context.Context, messages []*isb.ReadMessage)
 			// this is to minimize watermark latency
 			//TODO(multi-partition): support multi partitioned edges
 			for toVertex, toVertexBuffer := range df.toBuffers {
-				if publisher, ok := df.watermarkPublishers[toVertex]; ok {
-					idlehandler.PublishIdleWatermark(ctx, toVertexBuffer[0], publisher, df.idleManager, 0, df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(watermark))
+				for index, bufferPartition := range toVertexBuffer {
+					if publisher, ok := df.wmPublishers[toVertex]; ok {
+						idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, int32(index), df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(watermark))
+					}
 				}
 			}
 		}
@@ -568,12 +574,12 @@ func (df *DataForward) ShutDown(ctx context.Context) {
 	df.pbqManager.ShutDown(ctx)
 
 	// stop watermark fetcher
-	if err := df.watermarkFetcher.Close(); err != nil {
+	if err := df.wmFetcher.Close(); err != nil {
 		df.log.Errorw("Failed to close watermark fetcher", zap.Error(err))
 	}
 
 	// stop watermark publisher
-	for _, publisher := range df.watermarkPublishers {
+	for _, publisher := range df.wmPublishers {
 		if err := publisher.Close(); err != nil {
 			df.log.Errorw("Failed to close watermark publisher", zap.Error(err))
 		}
diff --git a/pkg/reduce/pnf/metrics.go b/pkg/reduce/pnf/metrics.go
index 3277ca63..1a18ec36 100644
--- a/pkg/reduce/pnf/metrics.go
+++ b/pkg/reduce/pnf/metrics.go
@@ -73,7 +73,7 @@ var reduceProcessTime = promauto.NewHistogramVec(prometheus.HistogramOpts{
 	Buckets:   prometheus.ExponentialBucketsRange(1, 1200000, 5),
 }, []string{metrics.LabelVertex, metrics.LabelPipeline, metrics.LabelVertexReplicaIndex})
 
-// reduceForwardTime is used to indicate the time it took to forward the result
+// reduceForwardTime is used to indicate the time it took to forward the writeMessages
 var reduceForwardTime = promauto.NewHistogramVec(prometheus.HistogramOpts{
 	Subsystem: "reduce_pnf",
 	Name:      "forward_time",
diff --git a/pkg/reduce/pnf/ordered.go b/pkg/reduce/pnf/ordered.go
index cbf09688..f75f9c86 100644
--- a/pkg/reduce/pnf/ordered.go
+++ b/pkg/reduce/pnf/ordered.go
@@ -47,7 +47,7 @@ type ForwardTask struct {
 	pf     *processAndForward
 }
 
-// OrderedProcessor orders the forwarding of the result of the execution of the tasks, even though the tasks itself are
+// OrderedProcessor orders the forwarding of the writeMessages of the execution of the tasks, even though the tasks itself are
 // run concurrently in an out of ordered fashion.
 type OrderedProcessor struct {
 	vertexName    string
@@ -127,7 +127,7 @@ func (op *OrderedProcessor) SchedulePnF(
 }
 
 // reduceOp invokes the reduce function. The reducer is a long-running function since we stream in the data and it has
-// to wait for the close-of-book on the PBQ to materialize the result.
+// to wait for the close-of-book on the PBQ to materialize the writeMessages.
 func (op *OrderedProcessor) reduceOp(ctx context.Context, t *ForwardTask) {
 	start := time.Now()
 	for {
@@ -161,7 +161,7 @@ func (op *OrderedProcessor) reduceOp(ctx context.Context, t *ForwardTask) {
 	}
 }
 
-// forward monitors the ForwardTask queue, as soon as the ForwardTask at the head of the queue has been completed, the result is
+// forward monitors the ForwardTask queue, as soon as the ForwardTask at the head of the queue has been completed, the writeMessages is
 // forwarded to the next ISB. It keeps doing this for forever or until ctx.Done() happens.
 func (op *OrderedProcessor) forward(ctx context.Context) {
 	var currElement *list.Element
diff --git a/pkg/reduce/pnf/processandforward.go b/pkg/reduce/pnf/processandforward.go
index 272a788d..50bf78af 100644
--- a/pkg/reduce/pnf/processandforward.go
+++ b/pkg/reduce/pnf/processandforward.go
@@ -47,18 +47,18 @@ import (
 // processAndForward reads messages from pbq, invokes udf using grpc, forwards the results to ISB, and then publishes
 // the watermark for that partition.
 type processAndForward struct {
-	vertexName       string
-	pipelineName     string
-	vertexReplica    int32
-	PartitionID      partition.ID
-	UDF              applier.ReduceApplier
-	result           []*isb.WriteMessage
-	pbqReader        pbq.Reader
-	log              *zap.SugaredLogger
-	toBuffers        map[string][]isb.BufferWriter
-	whereToDecider   forward.ToWhichStepDecider
-	publishWatermark map[string]publish.Publisher
-	idleManager      *wmb.IdleManager
+	vertexName     string
+	pipelineName   string
+	vertexReplica  int32
+	PartitionID    partition.ID
+	UDF            applier.ReduceApplier
+	writeMessages  []*isb.WriteMessage
+	pbqReader      pbq.Reader
+	log            *zap.SugaredLogger
+	toBuffers      map[string][]isb.BufferWriter
+	whereToDecider forward.ToWhichStepDecider
+	wmPublishers   map[string]publish.Publisher
+	idleManager    *wmb.IdleManager
 }
 
 // newProcessAndForward will return a new processAndForward instance
@@ -75,21 +75,21 @@ func newProcessAndForward(ctx context.Context,
 	idleManager *wmb.IdleManager) *processAndForward {
 
 	return &processAndForward{
-		vertexName:       vertexName,
-		pipelineName:     pipelineName,
-		vertexReplica:    vr,
-		PartitionID:      partitionID,
-		UDF:              udf,
-		pbqReader:        pbqReader,
-		log:              logging.FromContext(ctx),
-		toBuffers:        toBuffers,
-		whereToDecider:   whereToDecider,
-		publishWatermark: pw,
-		idleManager:      idleManager,
+		vertexName:     vertexName,
+		pipelineName:   pipelineName,
+		vertexReplica:  vr,
+		PartitionID:    partitionID,
+		UDF:            udf,
+		pbqReader:      pbqReader,
+		log:            logging.FromContext(ctx),
+		toBuffers:      toBuffers,
+		whereToDecider: whereToDecider,
+		wmPublishers:   pw,
+		idleManager:    idleManager,
 	}
 }
 
-// Process method reads messages from the supplied PBQ, invokes UDF to reduce the result.
+// Process method reads messages from the supplied PBQ, invokes UDF to reduce the writeMessages.
 func (p *processAndForward) Process(ctx context.Context) error {
 	var err error
 	startTime := time.Now()
@@ -99,8 +99,8 @@ func (p *processAndForward) Process(ctx context.Context) error {
 		metrics.LabelVertexReplicaIndex: strconv.Itoa(int(p.vertexReplica)),
 	}).Observe(float64(time.Since(startTime).Milliseconds()))
 
-	// blocking call, only returns the result after it has read all the messages from pbq
-	p.result, err = p.UDF.ApplyReduce(ctx, &p.PartitionID, p.pbqReader.ReadCh())
+	// blocking call, only returns the writeMessages after it has read all the messages from pbq
+	p.writeMessages, err = p.UDF.ApplyReduce(ctx, &p.PartitionID, p.pbqReader.ReadCh())
 	return err
 }
 
@@ -174,7 +174,7 @@ func (p *processAndForward) whereToStep() map[string][][]isb.Message {
 
 	var to []forward.VertexBuffer
 	var err error
-	for _, msg := range p.result {
+	for _, msg := range p.writeMessages {
 		to, err = p.whereToDecider.WhereTo(msg.Keys, msg.Tags)
 		if err != nil {
 			platformError.With(map[string]string{
@@ -284,26 +284,29 @@ func (p *processAndForward) publishWM(ctx context.Context, wm wmb.Watermark, wri
 	// activeWatermarkBuffers records the buffers that the publisher has published
 	// a watermark in this batch processing cycle.
 	// it's used to determine which buffers should receive an idle watermark.
-	var activeWatermarkBuffers = make(map[string]bool)
-	for edgeName, edgeOffsets := range writeOffsets {
-		for index, offsets := range edgeOffsets {
-			if publisher, ok := p.publishWatermark[edgeName]; ok {
+	// Created as a slice since it tracks per partition of the buffer.
+	var activeWatermarkBuffers = make(map[string][]bool)
+	for toVertexName, bufferOffsets := range writeOffsets {
+		activeWatermarkBuffers[toVertexName] = make([]bool, len(bufferOffsets))
+		for index, offsets := range bufferOffsets {
+			if publisher, ok := p.wmPublishers[toVertexName]; ok {
 				if len(offsets) > 0 {
 					publisher.PublishWatermark(wm, offsets[len(offsets)-1], int32(index))
-					activeWatermarkBuffers[edgeName] = true
+					activeWatermarkBuffers[toVertexName][index] = true
 					// reset because the toBuffer is not idling
-					p.idleManager.Reset(edgeName)
+					p.idleManager.Reset(toVertexName)
 				}
 			}
 		}
 	}
-	if len(activeWatermarkBuffers) < len(p.publishWatermark) {
-		// if there's any buffers that haven't received any watermark during this
-		// batch processing cycle, send an idle watermark
-		for toVertexName := range p.publishWatermark {
-			if !activeWatermarkBuffers[toVertexName] {
-				if publisher, ok := p.publishWatermark[toVertexName]; ok {
-					idlehandler.PublishIdleWatermark(ctx, p.toBuffers[toVertexName][0], publisher, p.idleManager, 0, p.log, dfv1.VertexTypeReduceUDF, wm)
+
+	// if there's any buffers that haven't received any watermark during this
+	// batch processing cycle, send an idle watermark
+	for toVertexName := range p.wmPublishers {
+		for index, activePartition := range activeWatermarkBuffers[toVertexName] {
+			if !activePartition {
+				if publisher, ok := p.wmPublishers[toVertexName]; ok {
+					idlehandler.PublishIdleWatermark(ctx, p.toBuffers[toVertexName][index], publisher, p.idleManager, int32(index), p.log, dfv1.VertexTypeReduceUDF, wm)
 				}
 			}
 		}
diff --git a/pkg/sources/generator/tickgen.go b/pkg/sources/generator/tickgen.go
index 49742c3c..78091a2a 100644
--- a/pkg/sources/generator/tickgen.go
+++ b/pkg/sources/generator/tickgen.go
@@ -233,7 +233,8 @@ func (mg *memgen) buildSourceWatermarkPublisher(publishWMStores store.WatermarkS
 	// for tickgen, it can be the name of the replica
 	entityName := fmt.Sprintf("%s-%d", mg.vertexInstance.Vertex.Name, mg.vertexInstance.Replica)
 	processorEntity := processor.NewProcessorEntity(entityName)
-	return publish.NewPublish(mg.lifecycleCtx, processorEntity, publishWMStores, publish.IsSource(), publish.WithDelay(mg.vertexInstance.Vertex.Spec.Watermark.GetMaxDelay()))
+	// source publisher toVertexPartitionCount will be 1, because we publish watermarks within source itself.
+	return publish.NewPublish(mg.lifecycleCtx, processorEntity, publishWMStores, 1, publish.IsSource(), publish.WithDelay(mg.vertexInstance.Vertex.Spec.Watermark.GetMaxDelay()))
 }
 
 func (mg *memgen) GetName() string {
@@ -269,6 +270,7 @@ func (mg *memgen) PublishSourceWatermarks(msgs []*isb.ReadMessage) {
 		return
 	}
 	// use the first event time of the message as watermark to make it conservative
+	// toVertexPartitionCount is 1 because we publish watermarks within source itself.
 	mg.sourcePublishWM.PublishWatermark(wmb.Watermark(msgs[0].EventTime), nil, 0) // Source publisher does not care about the offset
 }
 
diff --git a/pkg/sources/http/http.go b/pkg/sources/http/http.go
index ff2cbf8b..1e180d11 100644
--- a/pkg/sources/http/http.go
+++ b/pkg/sources/http/http.go
@@ -207,7 +207,8 @@ func New(
 	h.cancelFunc = cancel
 	entityName := fmt.Sprintf("%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica)
 	processorEntity := processor.NewProcessorEntity(entityName)
-	h.sourcePublishWM = publish.NewPublish(ctx, processorEntity, publishWMStores, publish.IsSource(), publish.WithDelay(vertexInstance.Vertex.Spec.Watermark.GetMaxDelay()))
+	// source publisher toVertexPartitionCount will be 1, because we publish watermarks within source itself.
+	h.sourcePublishWM = publish.NewPublish(ctx, processorEntity, publishWMStores, 1, publish.IsSource(), publish.WithDelay(vertexInstance.Vertex.Spec.Watermark.GetMaxDelay()))
 	return h, nil
 }
 
@@ -242,6 +243,7 @@ func (h *httpSource) PublishSourceWatermarks(msgs []*isb.ReadMessage) {
 	}
 	if len(msgs) > 0 && !oldest.IsZero() {
 		h.logger.Debugf("Publishing watermark %v to source", oldest)
+		// toVertexPartitionIdx is 0, because we publish watermarks within source itself.
 		h.sourcePublishWM.PublishWatermark(wmb.Watermark(oldest), nil, 0) // Source publisher does not care about the offset
 	}
 }
diff --git a/pkg/sources/kafka/reader.go b/pkg/sources/kafka/reader.go
index db21c805..6aeb0fe6 100644
--- a/pkg/sources/kafka/reader.go
+++ b/pkg/sources/kafka/reader.go
@@ -152,6 +152,7 @@ func (r *KafkaSource) PublishSourceWatermarks(msgs []*isb.ReadMessage) {
 	}
 	for p, t := range oldestTimestamps {
 		publisher := r.loadSourceWatermarkPublisher(p)
+		// toVertexPartitionIdx is 0 because we publish watermarks within source itself.
 		publisher.PublishWatermark(wmb.Watermark(t), nil, 0) // Source publisher does not care about the offset
 	}
 }
@@ -165,7 +166,8 @@ func (r *KafkaSource) loadSourceWatermarkPublisher(partitionID int32) publish.Pu
 	}
 	entityName := fmt.Sprintf("%s-%s-%d", r.pipelineName, r.name, partitionID)
 	processorEntity := processor.NewProcessorEntity(entityName)
-	sourcePublishWM := publish.NewPublish(r.lifecyclectx, processorEntity, r.srcPublishWMStores, publish.IsSource(), publish.WithDelay(r.watermarkMaxDelay))
+	// toVertexPartitionCount is 1 because we publish watermarks within source itself.
+	sourcePublishWM := publish.NewPublish(r.lifecyclectx, processorEntity, r.srcPublishWMStores, 1, publish.IsSource(), publish.WithDelay(r.watermarkMaxDelay))
 	r.sourcePublishWMs[partitionID] = sourcePublishWM
 	return sourcePublishWM
 }
diff --git a/pkg/sources/nats/nats.go b/pkg/sources/nats/nats.go
index 226b908c..f05ec283 100644
--- a/pkg/sources/nats/nats.go
+++ b/pkg/sources/nats/nats.go
@@ -100,7 +100,8 @@ func New(
 	n.cancelfn = cancel
 	entityName := fmt.Sprintf("%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica)
 	processorEntity := processor.NewProcessorEntity(entityName)
-	n.sourcePublishWM = publish.NewPublish(ctx, processorEntity, publishWMStores, publish.IsSource(), publish.WithDelay(vertexInstance.Vertex.Spec.Watermark.GetMaxDelay()))
+	// toVertexPartitionCount is 1 because we publish watermarks within source itself.
+	n.sourcePublishWM = publish.NewPublish(ctx, processorEntity, publishWMStores, 1, publish.IsSource(), publish.WithDelay(vertexInstance.Vertex.Spec.Watermark.GetMaxDelay()))
 
 	source := vertexInstance.Vertex.Spec.Source.Nats
 	opt := []natslib.Option{
@@ -240,6 +241,7 @@ func (ns *natsSource) PublishSourceWatermarks(msgs []*isb.ReadMessage) {
 		}
 	}
 	if len(msgs) > 0 && !oldest.IsZero() {
+		// toVertexPartitionIdx is 0 because we publish watermarks within source itself.
 		ns.sourcePublishWM.PublishWatermark(wmb.Watermark(oldest), nil, 0) // Source publisher does not care about the offset
 	}
 }
diff --git a/pkg/sources/redisstreams/redisstream.go b/pkg/sources/redisstreams/redisstream.go
index e5281478..b7f8a32c 100644
--- a/pkg/sources/redisstreams/redisstream.go
+++ b/pkg/sources/redisstreams/redisstream.go
@@ -155,7 +155,8 @@ func New(
 	redisStreamsSource.cancelfn = cancel
 	entityName := fmt.Sprintf("%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica)
 	processorEntity := processor.NewProcessorEntity(entityName)
-	redisStreamsSource.sourcePublishWM = publish.NewPublish(ctx, processorEntity, publishWMStores, publish.IsSource(), publish.WithDelay(vertexInstance.Vertex.Spec.Watermark.GetMaxDelay()))
+	// toVertexPartitionCount is 1 because we publish watermarks within source itself.
+	redisStreamsSource.sourcePublishWM = publish.NewPublish(ctx, processorEntity, publishWMStores, 1, publish.IsSource(), publish.WithDelay(vertexInstance.Vertex.Spec.Watermark.GetMaxDelay()))
 
 	// create the ConsumerGroup here if not already created
 	err = redisStreamsSource.createConsumerGroup(ctx, redisSpec)
@@ -293,6 +294,7 @@ func (rsSource *redisStreamsSource) PublishSourceWatermarks(msgs []*isb.ReadMess
 		}
 	}
 	if len(msgs) > 0 && !oldest.IsZero() {
+		// toVertexPartitionIdx is 0 because we publish watermarks within source itself.
 		rsSource.sourcePublishWM.PublishWatermark(wmb.Watermark(oldest), nil, 0) // Source publisher does not care about the offset
 	}
 }
diff --git a/pkg/watermark/fetch/edge_fetcher.go b/pkg/watermark/fetch/edge_fetcher.go
index 2af99cdc..1113be96 100644
--- a/pkg/watermark/fetch/edge_fetcher.go
+++ b/pkg/watermark/fetch/edge_fetcher.go
@@ -65,7 +65,7 @@ func (e *edgeFetcher) GetWatermark(inputOffset isb.Offset) wmb.Watermark {
 	var offset, err = inputOffset.Sequence()
 	if err != nil {
 		e.log.Errorw("Unable to get offset from isb.Offset.Sequence()", zap.Error(err))
-		return wmb.Watermark(time.Unix(-1, 0))
+		return wmb.InitialWatermark
 	}
 	var debugString strings.Builder
 	var epoch int64 = math.MaxInt64
@@ -120,7 +120,7 @@ func (e *edgeFetcher) GetHeadWatermark() wmb.Watermark {
 	e.log.Debugf("GetHeadWatermark: %s", debugString.String())
 	if headWatermark == math.MaxInt64 {
 		// Use -1 as default watermark value to indicate there is no valid watermark yet.
-		return wmb.Watermark(time.UnixMilli(-1))
+		return wmb.InitialWatermark
 	}
 	return wmb.Watermark(time.UnixMilli(headWatermark))
 }
diff --git a/pkg/watermark/fetch/source_fetcher.go b/pkg/watermark/fetch/source_fetcher.go
index 8105c25b..a6f5c334 100644
--- a/pkg/watermark/fetch/source_fetcher.go
+++ b/pkg/watermark/fetch/source_fetcher.go
@@ -72,7 +72,7 @@ func (e *sourceFetcher) GetWatermark(_ isb.Offset) wmb.Watermark {
 		}
 	}
 	if epoch == math.MaxInt64 {
-		epoch = -1
+		epoch = wmb.InitialWatermark.UnixMilli()
 	}
 	e.log.Debugf("%s get watermark for offset : %+v", debugString.String(), epoch)
 	return wmb.Watermark(time.UnixMilli(epoch))
@@ -91,7 +91,7 @@ func (e *sourceFetcher) GetHeadWatermark() wmb.Watermark {
 	}
 	if epoch == math.MinInt64 {
 		// Use -1 as default watermark value to indicate there is no valid watermark yet.
-		return wmb.Watermark(time.UnixMilli(-1))
+		return wmb.InitialWatermark
 	}
 	return wmb.Watermark(time.UnixMilli(epoch))
 }
diff --git a/pkg/watermark/generic/jetstream/generic.go b/pkg/watermark/generic/jetstream/generic.go
index bbe8625e..f8379cdb 100644
--- a/pkg/watermark/generic/jetstream/generic.go
+++ b/pkg/watermark/generic/jetstream/generic.go
@@ -95,7 +95,7 @@ func BuildWatermarkProgressors(ctx context.Context, vertexInstance *v1alpha1.Ver
 			return nil, nil, fmt.Errorf("failed at new OT Publish JetStreamKVStore, OTBucket: %s, %w", otStoreBucketName, err)
 		}
 		// For sink vertex, we use the vertex name as the to buffer name, which is the key for the publisher map.
-		publishWatermark[vertexInstance.Vertex.Spec.Name] = publish.NewPublish(ctx, publishEntity, store.BuildWatermarkStore(hbStore, otStore), publish.IsSink())
+		publishWatermark[vertexInstance.Vertex.Spec.Name] = publish.NewPublish(ctx, publishEntity, store.BuildWatermarkStore(hbStore, otStore), 1, publish.IsSink())
 	} else {
 		for _, e := range vertexInstance.Vertex.Spec.ToEdges {
 			toBucket := v1alpha1.GenerateEdgeBucketName(vertexInstance.Vertex.Namespace, pipelineName, e.From, e.To)
@@ -110,7 +110,7 @@ func BuildWatermarkProgressors(ctx context.Context, vertexInstance *v1alpha1.Ver
 			if err != nil {
 				return nil, nil, fmt.Errorf("failed at new OT Publish JetStreamKVStore, OTBucket: %s, %w", otStoreBucketName, err)
 			}
-			publishWatermark[e.To] = publish.NewPublish(ctx, publishEntity, store.BuildWatermarkStore(hbStore, otStore))
+			publishWatermark[e.To] = publish.NewPublish(ctx, publishEntity, store.BuildWatermarkStore(hbStore, otStore), int32(e.GetToVertexPartitions()))
 		}
 	}
 	return fetchWatermark, publishWatermark, nil
diff --git a/pkg/watermark/publish/publisher.go b/pkg/watermark/publish/publisher.go
index 667ad756..199a60cf 100644
--- a/pkg/watermark/publish/publisher.go
+++ b/pkg/watermark/publish/publisher.go
@@ -55,15 +55,16 @@ type publish struct {
 	// heartbeatStore uses second as the time unit for the value
 	heartbeatStore store.WatermarkKVStorer
 	// osStore uses millisecond as the time unit for the value
-	otStore       store.WatermarkKVStorer
-	log           *zap.SugaredLogger
-	headWatermark wmb.Watermark
-	headWMLock    sync.RWMutex
-	opts          *publishOptions
+	otStore                store.WatermarkKVStorer
+	log                    *zap.SugaredLogger
+	headWatermarks         []wmb.Watermark
+	headWMLock             sync.RWMutex
+	toVertexPartitionCount int32
+	opts                   *publishOptions
 }
 
 // NewPublish returns `Publish`.
-func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier, watermarkStores store.WatermarkStorer, inputOpts ...PublishOption) Publisher {
+func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier, watermarkStores store.WatermarkStorer, toVertexPartitionCount int32, inputOpts ...PublishOption) Publisher {
 	log := logging.FromContext(ctx).With("entityID", processorEntity.GetName()).
 		With("otStore", watermarkStores.OffsetTimelineStore().GetStoreName()).
 		With("hbStore", watermarkStores.HeartbeatStore().GetStoreName())
@@ -79,12 +80,13 @@ func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier
 	}
 
 	p := &publish{
-		ctx:            ctx,
-		entity:         processorEntity,
-		heartbeatStore: watermarkStores.HeartbeatStore(),
-		otStore:        watermarkStores.OffsetTimelineStore(),
-		log:            log,
-		opts:           opts,
+		ctx:                    ctx,
+		entity:                 processorEntity,
+		heartbeatStore:         watermarkStores.HeartbeatStore(),
+		otStore:                watermarkStores.OffsetTimelineStore(),
+		toVertexPartitionCount: toVertexPartitionCount,
+		log:                    log,
+		opts:                   opts,
 	}
 
 	p.initialSetup()
@@ -95,29 +97,35 @@ func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier
 	return p
 }
 
-// GetHeadWM gets the p.headWatermark field.
-func (p *publish) GetHeadWM() wmb.Watermark {
+// GetHeadWM gets the headWatermark for the given partition.
+func (p *publish) GetHeadWM(toVertexPartitionIdx int32) wmb.Watermark {
 	p.headWMLock.RLock()
 	defer p.headWMLock.RUnlock()
-	return p.headWatermark
+	return p.headWatermarks[toVertexPartitionIdx]
 }
 
-// SetHeadWM sets the p.headWatermark field using the given wm.
-func (p *publish) SetHeadWM(wm wmb.Watermark) {
+// SetHeadWM sets the headWatermark using the given wm for the given partition.
+func (p *publish) SetHeadWM(wm wmb.Watermark, toVertexPartitionIdx int32) {
 	p.headWMLock.Lock()
 	defer p.headWMLock.Unlock()
-	p.headWatermark = wm
+	p.headWatermarks[toVertexPartitionIdx] = wm
 }
 
 // initialSetup inserts the default values as the ProcessorEntitier starts emitting watermarks.
+// We will be initializing all to -1
+// TODO: we could ideally resume from where we left off, but this will introduce a new key.
 func (p *publish) initialSetup() {
-	p.SetHeadWM(p.loadLatestFromStore())
+	var headWms []wmb.Watermark
+	for i := 0; i < int(p.toVertexPartitionCount); i++ {
+		headWms = append(headWms, wmb.InitialWatermark)
+	}
+	p.headWatermarks = headWms
 }
 
 // PublishWatermark publishes watermark and will retry until it can succeed. It will not publish if the new-watermark
 // is less than the current head watermark.
 func (p *publish) PublishWatermark(wm wmb.Watermark, offset isb.Offset, toVertexPartitionIdx int32) {
-	validWM, skipWM := p.validateWatermark(wm)
+	validWM, skipWM := p.validateWatermark(wm, toVertexPartitionIdx)
 	if skipWM {
 		return
 	}
@@ -150,7 +158,7 @@ func (p *publish) PublishWatermark(wm wmb.Watermark, offset isb.Offset, toVertex
 			// TODO: better exponential backoff
 			time.Sleep(time.Millisecond * 250)
 		} else {
-			p.log.Debugw("New watermark published with offset", zap.Int64("head", p.GetHeadWM().UnixMilli()), zap.Int64("new", validWM.UnixMilli()), zap.Int64("offset", seq))
+			p.log.Debugw("New watermark published with offset", zap.Int64("head", p.GetHeadWM(toVertexPartitionIdx).UnixMilli()), zap.Int64("new", validWM.UnixMilli()), zap.Int64("offset", seq))
 			break
 		}
 	}
@@ -158,15 +166,15 @@ func (p *publish) PublishWatermark(wm wmb.Watermark, offset isb.Offset, toVertex
 
 // validateWatermark checks if the new watermark is greater than the head watermark, return true if yes,
 // otherwise, return false
-func (p *publish) validateWatermark(wm wmb.Watermark) (wmb.Watermark, bool) {
+func (p *publish) validateWatermark(wm wmb.Watermark, toVertexPartitionIdx int32) (wmb.Watermark, bool) {
 	if p.opts.isSource && p.opts.delay.Nanoseconds() > 0 && !time.Time(wm).IsZero() {
 		wm = wmb.Watermark(time.Time(wm).Add(-p.opts.delay))
 	}
-	// update p.headWatermark only if wm > p.headWatermark
-	headWM := p.GetHeadWM()
+	// update p.headWatermarks only if wm > p.headWatermarks
+	headWM := p.GetHeadWM(toVertexPartitionIdx)
 	if wm.After(time.Time(headWM)) {
 		p.log.Debugw("New watermark is updated for the head watermark", zap.String("head", headWM.String()), zap.String("new", wm.String()))
-		p.SetHeadWM(wm)
+		p.SetHeadWM(wm, toVertexPartitionIdx)
 	} else if wm.Before(time.Time(headWM)) {
 		p.log.Warnw("Skip publishing the new watermark because it's older than the current watermark", zap.String("entity", p.entity.GetName()), zap.Int64("head", headWM.UnixMilli()), zap.Int64("new", wm.UnixMilli()))
 		return wmb.Watermark{}, true
@@ -181,7 +189,7 @@ func (p *publish) validateWatermark(wm wmb.Watermark) (wmb.Watermark, bool) {
 // TODO: merge with PublishWatermark
 func (p *publish) PublishIdleWatermark(wm wmb.Watermark, offset isb.Offset, toVertexPartitionIdx int32) {
 	var key = p.entity.GetName()
-	validWM, skipWM := p.validateWatermark(wm)
+	validWM, skipWM := p.validateWatermark(wm, toVertexPartitionIdx)
 	if skipWM {
 		return
 	}
@@ -222,27 +230,33 @@ func (p *publish) PublishIdleWatermark(wm wmb.Watermark, offset isb.Offset, toVe
 // TODO: how to repopulate if the processing unit is down for a really long time?
 func (p *publish) loadLatestFromStore() wmb.Watermark {
 	var (
-		timeWatermark = time.UnixMilli(-1)
+		timeWatermark = wmb.InitialWatermark
 		key           = p.entity.GetName()
 	)
 	byteValue, err := p.otStore.GetValue(p.ctx, key)
 	// could happen during boot up
 	if err != nil {
 		p.log.Warnw("Unable to load latest watermark from wmb store (failed to get value from wmb store)", zap.String("OT", p.otStore.GetStoreName()), zap.String("processorEntity", p.entity.GetName()), zap.Error(err))
-		return wmb.Watermark(timeWatermark)
+		return timeWatermark
 	}
 	otValue, err := wmb.DecodeToWMB(byteValue)
 	if err != nil {
 		p.log.Errorw("Unable to load latest watermark from wmb store (failed to decode wmb value)", zap.String("OT", p.otStore.GetStoreName()), zap.String("processorEntity", p.entity.GetName()), zap.Error(err))
-		return wmb.Watermark(timeWatermark)
+		return timeWatermark
 	}
-	timeWatermark = time.UnixMilli(otValue.Watermark)
-	return wmb.Watermark(timeWatermark)
+	timeWatermark = wmb.Watermark(time.UnixMilli(otValue.Watermark))
+	return timeWatermark
 }
 
 // GetLatestWatermark returns the latest watermark for that processor.
 func (p *publish) GetLatestWatermark() wmb.Watermark {
-	return p.GetHeadWM()
+	var latestWatermark = wmb.InitialWatermark
+	for _, wm := range p.headWatermarks {
+		if wm.After(time.Time(latestWatermark)) {
+			latestWatermark = wm
+		}
+	}
+	return latestWatermark
 }
 
 func (p *publish) publishHeartbeat() {
diff --git a/pkg/watermark/wmb/watermark.go b/pkg/watermark/wmb/watermark.go
index 457fa63d..f26119b4 100644
--- a/pkg/watermark/wmb/watermark.go
+++ b/pkg/watermark/wmb/watermark.go
@@ -23,6 +23,8 @@ import "time"
 // NOTE: today we support only second progression of watermark, we need to support millisecond too.
 type Watermark time.Time
 
+var InitialWatermark = Watermark(time.UnixMilli(-1))
+
 func (w Watermark) String() string {
 	var location, _ = time.LoadLocation("UTC")
 	var t = time.Time(w).In(location)


Test patch:
diff --git a/pkg/forward/forward_test.go b/pkg/forward/forward_test.go
index 3c0741b8..3f178d6a 100644
--- a/pkg/forward/forward_test.go
+++ b/pkg/forward/forward_test.go
@@ -758,7 +758,7 @@ func TestNewInterStepDataForwardIdleWatermark(t *testing.T) {
 	// 5th & 6th batch: again idling but got diff head WMB
 	// so the head is still the same active watermark
 	// and no new ctrl message to the next vertex
-	f.fetchWatermark.(*testWMBFetcher).RevertBoolValue()
+	f.wmFetcher.(*testWMBFetcher).RevertBoolValue()
 	time.Sleep(2 * time.Second) // default read timeout is 1s
 	otKeys1, _ = otStores["to1"].GetAllKeys(ctx)
 	otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0])
@@ -1325,7 +1325,7 @@ func buildPublisherMapAndOTStore(toBuffers map[string][]isb.BufferWriter) (map[s
 		heartbeatKV, _, _ := inmem.NewKVInMemKVStore(ctx, testPipelineName, fmt.Sprintf(publisherHBKeyspace, key))
 		otKV, _, _ := inmem.NewKVInMemKVStore(ctx, testPipelineName, fmt.Sprintf(publisherOTKeyspace, key))
 		otStores[key] = otKV
-		p := publish.NewPublish(ctx, processorEntity, wmstore.BuildWatermarkStore(heartbeatKV, otKV), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
+		p := publish.NewPublish(ctx, processorEntity, wmstore.BuildWatermarkStore(heartbeatKV, otKV), 1, publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
 		publishers[key] = p
 	}
 	return publishers, otStores
diff --git a/pkg/reduce/data_forward_test.go b/pkg/reduce/data_forward_test.go
index eacdf763..179d3ed6 100644
--- a/pkg/reduce/data_forward_test.go
+++ b/pkg/reduce/data_forward_test.go
@@ -1110,7 +1110,7 @@ func fetcherAndPublisher(ctx context.Context, fromBuffer *simplebuffer.InMemoryB
 	ot, otWatcherCh, _ := inmem.NewKVInMemKVStore(ctx, pipelineName, otBucketName)
 
 	// publisher for source
-	sourcePublisher := publish.NewPublish(ctx, sourcePublishEntity, wmstore.BuildWatermarkStore(hb, ot), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
+	sourcePublisher := publish.NewPublish(ctx, sourcePublishEntity, wmstore.BuildWatermarkStore(hb, ot), 1, publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
 
 	// publish heartbeat for the processors
 	go func() {
@@ -1144,7 +1144,7 @@ func buildPublisherMapAndOTStore(ctx context.Context, toBuffers map[string][]isb
 		hb, hbKVEntry, _ := inmem.NewKVInMemKVStore(ctx, pipelineName, key+"_PROCESSORS")
 		ot, otKVEntry, _ := inmem.NewKVInMemKVStore(ctx, pipelineName, key+"_OT")
 		otStores[key] = ot
-		p := publish.NewPublish(ctx, publishEntity, wmstore.BuildWatermarkStore(hb, ot), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
+		p := publish.NewPublish(ctx, publishEntity, wmstore.BuildWatermarkStore(hb, ot), 1, publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
 		publishers[key] = p
 
 		go func() {
diff --git a/pkg/reduce/pnf/processandforward_test.go b/pkg/reduce/pnf/processandforward_test.go
index c1d60e0c..db18859d 100644
--- a/pkg/reduce/pnf/processandforward_test.go
+++ b/pkg/reduce/pnf/processandforward_test.go
@@ -93,7 +93,7 @@ func TestProcessAndForward_Process(t *testing.T) {
 	// 1. create a pbq which has to be passed to the process method
 	// 2. pass the pbqReader interface and create a new processAndForward instance
 	// 3. mock the grpc client methods
-	// 4. assert to check if the process method is returning the result
+	// 4. assert to check if the process method is returning the writeMessages
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
@@ -174,7 +174,7 @@ func TestProcessAndForward_Process(t *testing.T) {
 
 	err = pf.Process(ctx)
 	assert.NoError(t, err)
-	assert.Len(t, pf.result, 1)
+	assert.Len(t, pf.writeMessages, 1)
 }
 
 func TestProcessAndForward_Forward(t *testing.T) {
@@ -309,7 +309,7 @@ func TestProcessAndForward_Forward(t *testing.T) {
 			assert.Equal(t, value.expected[1], msgs1[0].Header.Kind == isb.WMB)
 			// pbq entry from the manager will be removed after forwarding
 			assert.Equal(t, nil, pbqManager.GetPBQ(value.id))
-			for bufferName := range value.pf.publishWatermark {
+			for bufferName := range value.pf.wmPublishers {
 				// NOTE: in this test we only have one processor to publish
 				// so len(otKeys) should always be 1
 				otKeys, _ := value.otStores[bufferName].GetAllKeys(ctx)
@@ -407,15 +407,15 @@ func createProcessAndForwardAndOTStore(ctx context.Context, key string, pbqManag
 	}
 
 	pf := processAndForward{
-		PartitionID:      testPartition,
-		UDF:              nil,
-		result:           result,
-		pbqReader:        simplePbq,
-		log:              logging.FromContext(ctx),
-		toBuffers:        toBuffers,
-		whereToDecider:   whereto,
-		publishWatermark: pw,
-		idleManager:      wmb.NewIdleManager(len(toBuffers)),
+		PartitionID:    testPartition,
+		UDF:            nil,
+		writeMessages:  result,
+		pbqReader:      simplePbq,
+		log:            logging.FromContext(ctx),
+		toBuffers:      toBuffers,
+		whereToDecider: whereto,
+		wmPublishers:   pw,
+		idleManager:    wmb.NewIdleManager(len(toBuffers)),
 	}
 
 	return pf, otStore
@@ -431,7 +431,7 @@ func buildPublisherMapAndOTStore(toBuffers map[string][]isb.BufferWriter) (map[s
 		heartbeatKV, _, _ := inmem.NewKVInMemKVStore(ctx, testPipelineName, fmt.Sprintf(publisherHBKeyspace, key))
 		otKV, _, _ := inmem.NewKVInMemKVStore(ctx, testPipelineName, fmt.Sprintf(publisherOTKeyspace, key))
 		otStores[key] = otKV
-		p := publish.NewPublish(ctx, processorEntity, wmstore.BuildWatermarkStore(heartbeatKV, otKV), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
+		p := publish.NewPublish(ctx, processorEntity, wmstore.BuildWatermarkStore(heartbeatKV, otKV), 1, publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
 		publishers[key] = p
 	}
 	return publishers, otStores
diff --git a/pkg/watermark/publish/publisher_inmem_test.go b/pkg/watermark/publish/publisher_inmem_test.go
index c34fa43e..2cb8a310 100644
--- a/pkg/watermark/publish/publisher_inmem_test.go
+++ b/pkg/watermark/publish/publisher_inmem_test.go
@@ -48,7 +48,7 @@ func TestPublisherWithSharedOTBuckets_InMem(t *testing.T) {
 
 	publishEntity := processor.NewProcessorEntity("publisherTestPod1")
 
-	p := NewPublish(ctx, publishEntity, store.BuildWatermarkStore(heartbeatKV, otKV), WithAutoRefreshHeartbeatDisabled(), WithPodHeartbeatRate(1)).(*publish)
+	p := NewPublish(ctx, publishEntity, store.BuildWatermarkStore(heartbeatKV, otKV), 1, WithAutoRefreshHeartbeatDisabled(), WithPodHeartbeatRate(1)).(*publish)
 
 	var epoch int64 = 1651161600000
 	var location, _ = time.LoadLocation("UTC")
diff --git a/pkg/watermark/publish/publisher_test.go b/pkg/watermark/publish/publisher_test.go
index 19fdcb27..707cf7b3 100644
--- a/pkg/watermark/publish/publisher_test.go
+++ b/pkg/watermark/publish/publisher_test.go
@@ -75,7 +75,7 @@ func TestPublisherWithSharedOTBucket(t *testing.T) {
 	otKV, err := jetstream.NewKVJetStreamKVStore(ctx, "testPublisher", keyspace+"_OT", defaultJetStreamClient)
 	assert.NoError(t, err)
 
-	p := NewPublish(ctx, publishEntity, store.BuildWatermarkStore(heartbeatKV, otKV), WithAutoRefreshHeartbeatDisabled(), WithPodHeartbeatRate(1)).(*publish)
+	p := NewPublish(ctx, publishEntity, store.BuildWatermarkStore(heartbeatKV, otKV), 1, WithAutoRefreshHeartbeatDisabled(), WithPodHeartbeatRate(1)).(*publish)
 
 	var epoch int64 = 1651161600000
 	var location, _ = time.LoadLocation("UTC")