Repository: numaproj/numaflow
Pull request #1512 (2024-02-16 20:44:09)

fix: idle manager refactor for multi partitions
The idleManager is shared across all forwarders for map and sink vertex, but out current implementation doesn't take multiple forwarders into consideration, causing race condition when publishing control messages.

In this PR
- Refactor the current idle manager interface to include `fromBufferParitionIndex`.
- Refactor the current idle manager structure to have a map[toBufferParitionName]int64.
```go
// forwarderActivePartition is a map[toPartitionName]uint64 to record if a forwarder is sending to the toPartition.
// for each "toPartition" we have an integer represents in binary, and mark the #n bit as 1 if the #n forwarder
// is sending to the "toPartition"
// example:
//   if we have three forwarders, the initial value in binary format will be {"toPartition": 000} which is 0 in decimal
//   if forwarder0 is sending data to the toPartition, then the value will become {"toPartition": 001} which is 1 in decimal
//   if forwarder1 is sending data to the toPartition, then the value will become {"toPartition": 011} which is 3 in decimal
// when we do the ctrlMsg check, we reply on the value to decide if we need to send a ctrlMsg
// note that we use uint64, therefore there is a limit of maximum 64 partitions
forwarderActiveToPartition map[string]uint64
```


Patch:
diff --git a/pkg/reduce/data_forward.go b/pkg/reduce/data_forward.go
index 86e7968a..ff941f40 100644
--- a/pkg/reduce/data_forward.go
+++ b/pkg/reduce/data_forward.go
@@ -251,7 +251,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 			for toVertexName, toVertexBuffer := range df.toBuffers {
 				if publisher, ok := df.wmPublishers[toVertexName]; ok {
 					for _, bufferPartition := range toVertexBuffer {
-						idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, df.log, df.vertexName, df.pipelineName, dfv1.VertexTypeReduceUDF, df.vertexReplica, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+						idlehandler.PublishIdleWatermark(ctx, wmb.PARTITION_0, bufferPartition, publisher, df.idleManager, df.log, df.vertexName, df.pipelineName, dfv1.VertexTypeReduceUDF, df.vertexReplica, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
 					}
 				}
 			}
@@ -274,7 +274,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 				for toVertexName, toVertexBuffer := range df.toBuffers {
 					if publisher, ok := df.wmPublishers[toVertexName]; ok {
 						for _, bufferPartition := range toVertexBuffer {
-							idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, df.log, df.vertexName, df.pipelineName, dfv1.VertexTypeReduceUDF, df.vertexReplica, wmb.Watermark(watermark))
+							idlehandler.PublishIdleWatermark(ctx, wmb.PARTITION_0, bufferPartition, publisher, df.idleManager, df.log, df.vertexName, df.pipelineName, dfv1.VertexTypeReduceUDF, df.vertexReplica, wmb.Watermark(watermark))
 						}
 					}
 				}
@@ -425,7 +425,7 @@ func (df *DataForward) process(ctx context.Context, messages []*isb.ReadMessage)
 			for toVertexName, toVertexBuffer := range df.toBuffers {
 				if publisher, ok := df.wmPublishers[toVertexName]; ok {
 					for _, bufferPartition := range toVertexBuffer {
-						idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, df.log, df.vertexName, df.pipelineName, dfv1.VertexTypeReduceUDF, df.vertexReplica, wmb.Watermark(watermark))
+						idlehandler.PublishIdleWatermark(ctx, wmb.PARTITION_0, bufferPartition, publisher, df.idleManager, df.log, df.vertexName, df.pipelineName, dfv1.VertexTypeReduceUDF, df.vertexReplica, wmb.Watermark(watermark))
 					}
 				}
 			}
diff --git a/pkg/reduce/pnf/processandforward.go b/pkg/reduce/pnf/processandforward.go
index 7f793510..980a68ea 100644
--- a/pkg/reduce/pnf/processandforward.go
+++ b/pkg/reduce/pnf/processandforward.go
@@ -379,7 +379,7 @@ func (p *processAndForward) publishWM(ctx context.Context, endTime time.Time) {
 					publisher.PublishWatermark(wm, offsets[len(offsets)-1], int32(index))
 					activeWatermarkBuffers[toVertexName][index] = true
 					// reset because the toBuffer partition is not idling
-					p.idleManager.Reset(p.toBuffers[toVertexName][index].GetName())
+					p.idleManager.MarkActive(wmb.PARTITION_0, p.toBuffers[toVertexName][index].GetName())
 				}
 			}
 		}
@@ -392,7 +392,7 @@ func (p *processAndForward) publishWM(ctx context.Context, endTime time.Time) {
 		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, p.log, p.vertexName, p.pipelineName, dfv1.VertexTypeReduceUDF, p.vertexReplica, wm)
+					idlehandler.PublishIdleWatermark(ctx, wmb.PARTITION_0, p.toBuffers[toVertexName][index], publisher, p.idleManager, p.log, p.vertexName, p.pipelineName, dfv1.VertexTypeReduceUDF, p.vertexReplica, wm)
 				}
 			}
 		}
diff --git a/pkg/shared/idlehandler/idlehandler.go b/pkg/shared/idlehandler/idlehandler.go
index 1ef2c444..d7c328e4 100644
--- a/pkg/shared/idlehandler/idlehandler.go
+++ b/pkg/shared/idlehandler/idlehandler.go
@@ -30,10 +30,10 @@ import (
 )
 
 // PublishIdleWatermark publishes a ctrl message with isb.Kind set to WMB. We only send one ctrl message when
-func PublishIdleWatermark(ctx context.Context, toBufferPartition isb.BufferWriter, wmPublisher publish.Publisher, idleManager wmb.IdleManager, logger *zap.SugaredLogger, vertexName string, pipelineName string, vertexType dfv1.VertexType, vertexReplica int32, wm wmb.Watermark) {
-
+func PublishIdleWatermark(ctx context.Context, fromBufferPartitionIndex int32, toBufferPartition isb.BufferWriter, wmPublisher publish.Publisher, idleManager wmb.IdleManager, logger *zap.SugaredLogger, vertexName string, pipelineName string, vertexType dfv1.VertexType, vertexReplica int32, wm wmb.Watermark) {
 	var toPartitionName = toBufferPartition.GetName()
 	var toVertexPartition = toBufferPartition.GetPartitionIdx()
+	idleManager.MarkIdle(fromBufferPartitionIndex, toPartitionName)
 
 	if idleManager.NeedToSendCtrlMsg(toPartitionName) {
 		if vertexType == dfv1.VertexTypeSink {
@@ -53,7 +53,7 @@ func PublishIdleWatermark(ctx context.Context, toBufferPartition isb.BufferWrite
 
 			if len(writeOffsets) == 1 {
 				// we only write one ctrl message, so there's only one offset in the array, use index=0 to get the offset
-				idleManager.Update(toPartitionName, writeOffsets[0])
+				idleManager.Update(fromBufferPartitionIndex, toPartitionName, writeOffsets[0])
 			}
 			metrics.CtrlMessagesCount.With(map[string]string{metrics.LabelVertex: vertexName, metrics.LabelPipeline: pipelineName, metrics.LabelVertexType: string(vertexType), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(vertexReplica)), metrics.LabelPartitionName: toPartitionName}).Add(1)
 
diff --git a/pkg/sinks/forward/forward.go b/pkg/sinks/forward/forward.go
index 7d3023ca..90f5041f 100644
--- a/pkg/sinks/forward/forward.go
+++ b/pkg/sinks/forward/forward.go
@@ -198,7 +198,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 		}
 
 		// if the validation passed, we will publish the watermark to all the toBuffer partitions.
-		idlehandler.PublishIdleWatermark(ctx, df.toBuffer, df.wmPublisher, df.idleManager, df.opts.logger, df.vertexName, df.pipelineName, dfv1.VertexTypeSink, df.vertexReplica, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+		idlehandler.PublishIdleWatermark(ctx, df.toBuffer.GetPartitionIdx(), df.toBuffer, df.wmPublisher, df.idleManager, df.opts.logger, df.vertexName, df.pipelineName, dfv1.VertexTypeSink, df.vertexReplica, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
 		return
 	}
 
@@ -284,7 +284,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 	if len(writeOffsets) > 0 {
 		df.wmPublisher.PublishWatermark(processorWM, nil, int32(0))
 		// reset because the toBuffer is no longer idling
-		df.idleManager.Reset(df.toBuffer.GetName())
+		df.idleManager.MarkActive(df.fromBufferPartition.GetPartitionIdx(), df.toBuffer.GetName())
 	}
 
 	err = df.ackFromBuffer(ctx, readOffsets)
diff --git a/pkg/sinks/sink.go b/pkg/sinks/sink.go
index 71d83e88..8a8b60bf 100644
--- a/pkg/sinks/sink.go
+++ b/pkg/sinks/sink.go
@@ -133,7 +133,7 @@ func (u *SinkProcessor) Start(ctx context.Context) error {
 			// create watermark publisher using watermark stores
 			publishWatermark = jetstream.BuildPublishersFromStores(ctx, u.VertexInstance, sinkWmStores)
 			// sink vertex has only one toBuffer, so the length is 1
-			idleManager = wmb.NewIdleManager(1)
+			idleManager, _ = wmb.NewIdleManager(len(readers), 1)
 		}
 
 	default:
diff --git a/pkg/sources/forward/data_forward.go b/pkg/sources/forward/data_forward.go
index c20155a8..163f5cd1 100644
--- a/pkg/sources/forward/data_forward.go
+++ b/pkg/sources/forward/data_forward.go
@@ -236,7 +236,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 							publisher = df.createToVertexWatermarkPublisher(toVertexName, sp)
 							vertexPublishers[sp] = publisher
 						}
-						idlehandler.PublishIdleWatermark(ctx, df.toBuffers[toVertexName][index], publisher, df.idleManager, df.opts.logger, df.vertexName, df.pipelineName, dfv1.VertexTypeSource, df.vertexReplica, fetchedWm)
+						idlehandler.PublishIdleWatermark(ctx, wmb.PARTITION_0, df.toBuffers[toVertexName][index], publisher, df.idleManager, df.opts.logger, df.vertexName, df.pipelineName, dfv1.VertexTypeSource, df.vertexReplica, fetchedWm)
 					}
 				}
 			}
@@ -395,7 +395,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 						activeWatermarkBuffers[toVertexName][index] = true
 						// update the watermark configs for lastTimestampSrcWMUpdated, lastFetchedSrcWatermark and lastTimestampIdleWMFound.
 						// reset because the toBuffer partition is no longer idling
-						df.idleManager.Reset(df.toBuffers[toVertexName][index].GetName())
+						df.idleManager.MarkActive(wmb.PARTITION_0, df.toBuffers[toVertexName][index].GetName())
 					}
 				}
 				// This (len(offsets) == 0) happens at conditional forwarding, there's no data written to the buffer
@@ -416,7 +416,7 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 							publisher = df.createToVertexWatermarkPublisher(toVertexName, sp)
 							vertexPublishers[sp] = publisher
 						}
-						idlehandler.PublishIdleWatermark(ctx, df.toBuffers[toVertexName][index], publisher, df.idleManager, df.opts.logger, df.vertexName, df.pipelineName, dfv1.VertexTypeSource, df.vertexReplica, processorWM)
+						idlehandler.PublishIdleWatermark(ctx, wmb.PARTITION_0, df.toBuffers[toVertexName][index], publisher, df.idleManager, df.opts.logger, df.vertexName, df.pipelineName, dfv1.VertexTypeSource, df.vertexReplica, processorWM)
 					}
 				}
 			}
diff --git a/pkg/sources/source.go b/pkg/sources/source.go
index 55ccad99..cd5358a5 100644
--- a/pkg/sources/source.go
+++ b/pkg/sources/source.go
@@ -163,7 +163,7 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 			if err != nil {
 				return err
 			}
-			idleManager = wmb.NewIdleManager(len(writersMap))
+			idleManager, _ = wmb.NewIdleManager(1, len(writersMap))
 		}
 	default:
 		return fmt.Errorf("unrecognized isb svc type %q", sp.ISBSvcType)
diff --git a/pkg/udf/forward/forward.go b/pkg/udf/forward/forward.go
index d28db484..893e7e59 100644
--- a/pkg/udf/forward/forward.go
+++ b/pkg/udf/forward/forward.go
@@ -228,7 +228,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 		for toVertexName, toVertexBuffer := range isdf.toBuffers {
 			for _, partition := range toVertexBuffer {
 				if p, ok := isdf.wmPublishers[toVertexName]; ok {
-					idlehandler.PublishIdleWatermark(ctx, partition, p, isdf.idleManager, isdf.opts.logger, isdf.vertexName, isdf.pipelineName, dfv1.VertexTypeMapUDF, isdf.vertexReplica, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+					idlehandler.PublishIdleWatermark(ctx, isdf.fromBufferPartition.GetPartitionIdx(), partition, p, isdf.idleManager, isdf.opts.logger, isdf.vertexName, isdf.pipelineName, dfv1.VertexTypeMapUDF, isdf.vertexReplica, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
 				}
 			}
 		}
@@ -356,7 +356,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 					publisher.PublishWatermark(processorWM, offsets[len(offsets)-1], int32(index))
 					activeWatermarkBuffers[toVertexName][index] = true
 					// reset because the toBuffer partition is no longer idling
-					isdf.idleManager.Reset(isdf.toBuffers[toVertexName][index].GetName())
+					isdf.idleManager.MarkActive(isdf.fromBufferPartition.GetPartitionIdx(), isdf.toBuffers[toVertexName][index].GetName())
 				}
 				// This (len(offsets) == 0) happens at conditional forwarding, there's no data written to the buffer
 			}
@@ -378,7 +378,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 					// 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, isdf.opts.logger, isdf.vertexName, isdf.pipelineName, dfv1.VertexTypeMapUDF, isdf.vertexReplica, processorWM)
+						idlehandler.PublishIdleWatermark(ctx, isdf.fromBufferPartition.GetPartitionIdx(), isdf.toBuffers[bufferName][index], p, isdf.idleManager, isdf.opts.logger, isdf.vertexName, isdf.pipelineName, dfv1.VertexTypeMapUDF, isdf.vertexReplica, processorWM)
 					}
 				}
 			}
diff --git a/pkg/udf/map_udf.go b/pkg/udf/map_udf.go
index 37c119ad..50d6651e 100644
--- a/pkg/udf/map_udf.go
+++ b/pkg/udf/map_udf.go
@@ -117,7 +117,7 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 			// create watermark publisher using watermark stores
 			publishWatermark = jetstream.BuildPublishersFromStores(ctx, u.VertexInstance, toVertexWmStores)
 
-			idleManager = wmb.NewIdleManager(len(writers))
+			idleManager, _ = wmb.NewIdleManager(len(writers), len(writers))
 		}
 	default:
 		return fmt.Errorf("unrecognized isbsvc type %q", u.ISBSvcType)
diff --git a/pkg/udf/reduce_udf.go b/pkg/udf/reduce_udf.go
index 657e3ea8..d2f44539 100644
--- a/pkg/udf/reduce_udf.go
+++ b/pkg/udf/reduce_udf.go
@@ -212,7 +212,7 @@ func (u *ReduceUDFProcessor) Start(ctx context.Context) error {
 			// create watermark publisher using watermark stores
 			publishWatermark = jetstream.BuildPublishersFromStores(ctx, u.VertexInstance, toVertexWmStores)
 
-			idleManager = wmb.NewIdleManager(len(writers))
+			idleManager, _ = wmb.NewIdleManager(1, len(writers))
 		}
 	default:
 		return fmt.Errorf("unrecognized isbsvc type %q", u.ISBSvcType)
diff --git a/pkg/watermark/wmb/idle_manager.go b/pkg/watermark/wmb/idle_manager.go
index 567f128c..1b3e7504 100644
--- a/pkg/watermark/wmb/idle_manager.go
+++ b/pkg/watermark/wmb/idle_manager.go
@@ -17,29 +17,53 @@ limitations under the License.
 package wmb
 
 import (
+	"fmt"
 	"sync"
 
 	"github.com/numaproj/numaflow/pkg/isb"
 )
 
+// PARTITION_0 is used when it is the only partition of the buffer, and this is used only for those cases
+// where we cannot have more than one partition (e.g., Reduce vertex, Source, etc.)
+const PARTITION_0 = 0
+
 // idleManager manages the idle watermark whether the control message is a duplicate and also keeps track of the idle WMB's offset.
 type idleManager struct {
+	// forwarderActivePartition is a map[toPartitionName]uint64 to record if a forwarder is sending to the toPartition.
+	// for each "toPartition" we have an integer represents in binary, and mark the #n bit as 1 if the #n forwarder
+	// is sending to the "toPartition"
+	// example:
+	//   if we have three forwarders, the initial value in binary format will be {"toPartition": 000}, which is 0 in decimal
+	//   if forwarder0 is sending data to the toPartition, then the value will become {"toPartition": 001}
+	//   if forwarder1 is sending data to the toPartition, then the value will become {"toPartition": 010}
+	//   if forwarder2 is sending data to the toPartition, then the value will become {"toPartition": 100}
+	// when we do the ctrlMsg check, we rely on the value (0 or non-0) to decide if we need to send a ctrlMsg
+	// note that we use uint64, therefore there is a limit of maximum 64 partitions
+	forwarderActiveToPartition map[string]uint64
 	// wmbOffset is a toBuffer partition name to the write offset of the idle watermark map.
 	wmbOffset map[string]isb.Offset
 	lock      sync.RWMutex
 }
 
 // NewIdleManager returns an idleManager object as the IdleManager interface type to track the watermark idle status.
-func NewIdleManager(length int) IdleManager {
-	return &idleManager{
-		wmbOffset: make(map[string]isb.Offset, length),
+func NewIdleManager(numOfFromPartitions int, numOfToPartitions int) (IdleManager, error) {
+	if numOfFromPartitions <= 0 || numOfFromPartitions >= 64 {
+		return nil, fmt.Errorf("failed to create a new idle manager: numOfFromPartitions should be > 0 and < 64")
 	}
+	return &idleManager{
+		forwarderActiveToPartition: make(map[string]uint64, numOfToPartitions),
+		wmbOffset:                  make(map[string]isb.Offset, numOfToPartitions),
+	}, nil
 }
 
 // NeedToSendCtrlMsg returns true if the given partition hasn't got any control message and needs to create a new control message
 func (im *idleManager) NeedToSendCtrlMsg(toBufferPartitionName string) bool {
 	im.lock.RLock()
 	defer im.lock.RUnlock()
+	if im.forwarderActiveToPartition[toBufferPartitionName] != 0 {
+		// only send ctrl msg when all bits are 0 (0 means inactive)
+		return false
+	}
 	// if the given partition doesn't have a control message
 	// the map entry will be empty, return true
 	return im.wmbOffset[toBufferPartitionName] == nil
@@ -53,15 +77,48 @@ func (im *idleManager) Get(toBufferPartitionName string) isb.Offset {
 }
 
 // Update will update the existing item or add if not present for the given toBuffer partition name.
-func (im *idleManager) Update(toBufferPartitionName string, newOffset isb.Offset) {
+func (im *idleManager) Update(fromBufferPartitionIndex int32, toBufferPartitionName string, newOffset isb.Offset) {
 	im.lock.Lock()
 	defer im.lock.Unlock()
+	// clear means the fromBufferPartition forwarder is inactive as the default value is 0
+	im.forwarderActiveToPartition[toBufferPartitionName] = clearBit(im.forwarderActiveToPartition[toBufferPartitionName], uint(fromBufferPartitionIndex))
 	im.wmbOffset[toBufferPartitionName] = newOffset
 }
 
-// Reset will clear the item for the given toBuffer partition name.
-func (im *idleManager) Reset(toBufferPartitionName string) {
+// MarkActive marks active for the given toBuffer partition name.
+func (im *idleManager) MarkActive(fromBufferPartitionIndex int32, toBufferPartitionName string) {
 	im.lock.Lock()
 	defer im.lock.Unlock()
+	// set to 1 means mark the fromBufferPartition forwarder as active
+	im.forwarderActiveToPartition[toBufferPartitionName] = setBit(im.forwarderActiveToPartition[toBufferPartitionName], uint(fromBufferPartitionIndex))
 	im.wmbOffset[toBufferPartitionName] = nil
 }
+
+// MarkIdle marks idle for the given toBuffer partition name if it's active
+func (im *idleManager) MarkIdle(fromBufferPartitionIndex int32, toBufferPartitionName string) {
+	im.lock.Lock()
+	defer im.lock.Unlock()
+	if isBitSet(im.forwarderActiveToPartition[toBufferPartitionName], uint(fromBufferPartitionIndex)) {
+		// the bit value at position fromBufferPartitionIndex is 1, meaning it's marked as active, sets to idle
+		im.forwarderActiveToPartition[toBufferPartitionName] = clearBit(im.forwarderActiveToPartition[toBufferPartitionName], uint(fromBufferPartitionIndex))
+	}
+	// the bit value is 0, meaning it's already idling, no action required
+}
+
+// setBit sets the bit at pos in the integer n.
+func setBit(n uint64, pos uint) uint64 {
+	n |= uint64(1) << pos
+	return n
+}
+
+// clearBit clears the bit at pos in n.
+func clearBit(n uint64, pos uint) uint64 {
+	n &^= uint64(1) << pos
+	return n
+}
+
+// isBitSet checks if the bit at pos in n is set to 1.
+func isBitSet(n uint64, pos uint) bool {
+	conditionCheck := uint64(1) << pos
+	return (n & conditionCheck) != 0
+}
diff --git a/pkg/watermark/wmb/interface.go b/pkg/watermark/wmb/interface.go
index b3ee9139..16050fc1 100644
--- a/pkg/watermark/wmb/interface.go
+++ b/pkg/watermark/wmb/interface.go
@@ -26,7 +26,9 @@ type IdleManager interface {
 	Get(toBufferPartitionName string) isb.Offset
 	// Update updates the isb.Offset for the given partition using the new offset if the partition exists, otherwise
 	// create a new entry.
-	Update(toBufferPartitionName string, newOffset isb.Offset)
-	// Reset resets the idle status for the given partition.
-	Reset(toBufferPartitionName string)
+	Update(fromBufferPartitionIndex int32, toBufferPartitionName string, newOffset isb.Offset)
+	// MarkActive marks the active status for the given partition.
+	MarkActive(fromBufferPartitionIndex int32, toBufferPartitionName string)
+	// MarkIdle marks idle for the given toBuffer partition name if it's not idle.
+	MarkIdle(fromBufferPartitionIndex int32, toBufferPartitionName string)
 }
diff --git a/pkg/watermark/wmb/noop_idle_manager.go b/pkg/watermark/wmb/noop_idle_manager.go
index e0c68c71..8b3b9842 100644
--- a/pkg/watermark/wmb/noop_idle_manager.go
+++ b/pkg/watermark/wmb/noop_idle_manager.go
@@ -38,8 +38,11 @@ func (n *noOpIdleManager) Get(string) isb.Offset {
 	return isb.SimpleIntOffset(func() int64 { return int64(-1) })
 }
 
-func (n *noOpIdleManager) Update(string, isb.Offset) {
+func (n *noOpIdleManager) Update(int32, string, isb.Offset) {
 }
 
-func (n *noOpIdleManager) Reset(string) {
+func (n *noOpIdleManager) MarkActive(int32, string) {
+}
+
+func (n *noOpIdleManager) MarkIdle(int32, string) {
 }


Test patch:
diff --git a/pkg/isb/stores/jetstream/writer_test.go b/pkg/isb/stores/jetstream/writer_test.go
index bbc785dd..7e1f41f5 100644
--- a/pkg/isb/stores/jetstream/writer_test.go
+++ b/pkg/isb/stores/jetstream/writer_test.go
@@ -146,7 +146,9 @@ func TestForwarderJetStreamBuffer(t *testing.T) {
 			}
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := forward.NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardJetStreamTest{}, myForwardJetStreamTest{}, myForwardJetStreamTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), forward.WithReadBatchSize(tt.batchSize), forward.WithUDFStreaming(tt.streamEnabled))
+			idleManager, err := wmb.NewIdleManager(1, len(toSteps))
+			assert.NoError(t, err)
+			f, err := forward.NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardJetStreamTest{}, myForwardJetStreamTest{}, myForwardJetStreamTest{}, fetchWatermark, publishWatermark, idleManager, forward.WithReadBatchSize(tt.batchSize), forward.WithUDFStreaming(tt.streamEnabled))
 			assert.NoError(t, err)
 
 			stopped := f.Start()
diff --git a/pkg/reduce/data_forward_test.go b/pkg/reduce/data_forward_test.go
index 275aa1e9..864daa9e 100644
--- a/pkg/reduce/data_forward_test.go
+++ b/pkg/reduce/data_forward_test.go
@@ -408,7 +408,8 @@ func TestDataForward_StartWithNoOpWM(t *testing.T) {
 	// create new fixed windower of (windowTime)
 	windower := fixed.NewWindower(windowTime, keyedVertex)
 
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(child, keyedVertex, CounterReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publisher, idleManager, windower)
 
 	var reduceDataForwarder *DataForward
@@ -504,7 +505,8 @@ func TestReduceDataForward_IdleWM(t *testing.T) {
 
 	// create a fixed windower of 5s
 	windower := fixed.NewWindower(5*time.Second, keyedVertex)
-	idleManager := wmb.NewIdleManager(len(toBuffers))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffers))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, CounterReduceTest{}, toBuffers, pbqManager, CounterReduceTest{}, publisherMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -713,7 +715,8 @@ func TestReduceDataForward_Count(t *testing.T) {
 
 	// create a fixed window of 60s
 	windower := fixed.NewWindower(60*time.Second, keyedVertex)
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, CounterReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publisherMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -798,7 +801,8 @@ func TestReduceDataForward_AllowedLatencyCount(t *testing.T) {
 	// create a fixed windower of 10s
 	windower := fixed.NewWindower(5*time.Second, keyedVertex)
 
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, CounterReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publisherMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -885,7 +889,8 @@ func TestReduceDataForward_Sum(t *testing.T) {
 
 	// create a fixed window of 2 minutes
 	windower := fixed.NewWindower(2*time.Minute, keyedVertex)
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, SumReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publishersMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -970,7 +975,8 @@ func TestReduceDataForward_Max(t *testing.T) {
 
 	// create a fixed window of 5 minutes
 	windower := fixed.NewWindower(5*time.Minute, keyedVertex)
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, MaxReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publishersMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -1056,7 +1062,8 @@ func TestReduceDataForward_FixedSumWithDifferentKeys(t *testing.T) {
 	// create a fixed windower of 5 minutes
 	windower := fixed.NewWindower(5*time.Minute, keyedVertex)
 
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, SumReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publishersMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -1162,7 +1169,8 @@ func TestReduceDataForward_SumWithDifferentKeys(t *testing.T) {
 	// create a session windower with 1 minute timeout
 	windower := session.NewWindower(1*time.Minute, keyedVertex)
 
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, SessionSumReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publishersMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -1266,7 +1274,8 @@ func TestReduceDataForward_NonKeyed(t *testing.T) {
 	// create a fixed window of 5 minutes
 	windower := fixed.NewWindower(5*time.Minute, keyedVertex)
 
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, SumReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publishersMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -1356,7 +1365,8 @@ func TestDataForward_WithContextClose(t *testing.T) {
 	// create a fixed windower of 5 minutes
 	windower := fixed.NewWindower(5*time.Minute, keyedVertex)
 
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, SumReduceTest{}, toBuffer, pbqManager, CounterReduceTest{}, publishersMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
@@ -1453,7 +1463,8 @@ func TestReduceDataForward_SumMultiPartitions(t *testing.T) {
 	// create a fixed windower of 5 minutes
 	windower := fixed.NewWindower(5*time.Minute, keyedVertex)
 
-	idleManager := wmb.NewIdleManager(len(toBuffer))
+	idleManager, err := wmb.NewIdleManager(1, len(toBuffer))
+	assert.NoError(t, err)
 	op := pnf.NewPnFManager(ctx, keyedVertex, SumReduceTest{}, toBuffer, pbqManager, &myForwardTestRoundRobin{}, publishersMap, idleManager, windower)
 
 	var reduceDataForward *DataForward
diff --git a/pkg/reduce/pnf/processandforward_test.go b/pkg/reduce/pnf/processandforward_test.go
index 87e95191..c5bd235d 100644
--- a/pkg/reduce/pnf/processandforward_test.go
+++ b/pkg/reduce/pnf/processandforward_test.go
@@ -181,13 +181,15 @@ func TestPnFHandleAlignedWindowResponses(t *testing.T) {
 		latestWriteOffsets[toVertexName] = make([][]isb.Offset, len(toVertexBuffer))
 	}
 
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffersMap))
+
 	pf := &processAndForward{
 		partitionId:        id,
 		whereToDecider:     whereto,
 		windower:           windower,
 		toBuffers:          toBuffersMap,
 		wmPublishers:       wmPublishers,
-		idleManager:        wmb.NewIdleManager(len(toBuffersMap)),
+		idleManager:        idleManager,
 		pbqReader:          &pbqReader{},
 		done:               make(chan struct{}),
 		latestWriteOffsets: latestWriteOffsets,
@@ -261,6 +263,8 @@ func createProcessAndForwardAndOTStore(ctx context.Context, key string, pbqManag
 		buffers: buffers,
 	}
 
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+
 	pf := processAndForward{
 		partitionId:    testPartition,
 		UDF:            nil,
@@ -269,7 +273,7 @@ func createProcessAndForwardAndOTStore(ctx context.Context, key string, pbqManag
 		toBuffers:      toBuffers,
 		whereToDecider: whereto,
 		wmPublishers:   pw,
-		idleManager:    wmb.NewIdleManager(len(toBuffers)),
+		idleManager:    idleManager,
 	}
 
 	return pf, otStore
diff --git a/pkg/sinks/blackhole/blackhole_test.go b/pkg/sinks/blackhole/blackhole_test.go
index 07adf80d..637eefbc 100644
--- a/pkg/sinks/blackhole/blackhole_test.go
+++ b/pkg/sinks/blackhole/blackhole_test.go
@@ -51,7 +51,8 @@ func TestBlackhole_Start(t *testing.T) {
 		Replica: 0,
 	}
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferList([]string{vertex.Spec.Name})
-	s, err := NewBlackhole(vertexInstance, fromStep, fetchWatermark, publishWatermark[vertex.Spec.Name], wmb.NewIdleManager(1))
+	idleManager, _ := wmb.NewIdleManager(1, 1)
+	s, err := NewBlackhole(vertexInstance, fromStep, fetchWatermark, publishWatermark[vertex.Spec.Name], idleManager)
 	assert.NoError(t, err)
 
 	stopped := s.Start()
diff --git a/pkg/sinks/forward/forward_test.go b/pkg/sinks/forward/forward_test.go
index 84ebe9a2..7a797093 100644
--- a/pkg/sinks/forward/forward_test.go
+++ b/pkg/sinks/forward/forward_test.go
@@ -159,7 +159,8 @@ func TestNewDataForward(t *testing.T) {
 
 		_, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 		fetchWatermark := &testForwardFetcher{}
-		f, err := NewDataForward(vertexInstance, fromStep, to1, fetchWatermark, publishWatermark[testVertexName], wmb.NewIdleManager(1))
+		idleManager, _ := wmb.NewIdleManager(1, 1)
+		f, err := NewDataForward(vertexInstance, fromStep, to1, fetchWatermark, publishWatermark[testVertexName], idleManager)
 
 		assert.NoError(t, err)
 		assert.False(t, to1.IsFull())
@@ -243,7 +244,8 @@ func TestNewDataForward(t *testing.T) {
 		publishWatermark := map[string]publish.Publisher{
 			testVertexName: &testForwarderPublisher{},
 		}
-		f, err := NewDataForward(vertexInstance, fromStep, to1, fetchWatermark, publishWatermark[testVertexName], wmb.NewIdleManager(1))
+		idleManager, _ := wmb.NewIdleManager(1, 1)
+		f, err := NewDataForward(vertexInstance, fromStep, to1, fetchWatermark, publishWatermark[testVertexName], idleManager)
 
 		assert.NoError(t, err)
 		assert.False(t, to1.IsFull())
@@ -326,7 +328,8 @@ func TestWriteToBuffer(t *testing.T) {
 			publishWatermark := map[string]publish.Publisher{
 				"to1": &testForwarderPublisher{},
 			}
-			f, err := NewDataForward(vertexInstance, fromStep, buffer, fetchWatermark, publishWatermark["to1"], wmb.NewIdleManager(1))
+			idleManager, _ := wmb.NewIdleManager(1, 1)
+			f, err := NewDataForward(vertexInstance, fromStep, buffer, fetchWatermark, publishWatermark["to1"], idleManager)
 
 			assert.NoError(t, err)
 			assert.False(t, buffer.IsFull())
diff --git a/pkg/sinks/forward/shutdown_test.go b/pkg/sinks/forward/shutdown_test.go
index 4bed7159..9485b3af 100644
--- a/pkg/sinks/forward/shutdown_test.go
+++ b/pkg/sinks/forward/shutdown_test.go
@@ -72,7 +72,9 @@ func TestShutDown(t *testing.T) {
 			}
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, to1, fetchWatermark, publishWatermark["to1"], wmb.NewIdleManager(1), WithReadBatchSize(batchSize))
+
+			idleManager, _ := wmb.NewIdleManager(1, 1)
+			f, err := NewDataForward(vertexInstance, fromStep, to1, fetchWatermark, publishWatermark["to1"], idleManager, WithReadBatchSize(batchSize))
 			assert.NoError(t, err)
 			stopped := f.Start()
 			// write some data but buffer is not full even though we are not reading
@@ -110,7 +112,9 @@ func TestShutDown(t *testing.T) {
 			}
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, to1, fetchWatermark, publishWatermark["to1"], wmb.NewIdleManager(1), WithReadBatchSize(batchSize))
+
+			idleManager, _ := wmb.NewIdleManager(1, 1)
+			f, err := NewDataForward(vertexInstance, fromStep, to1, fetchWatermark, publishWatermark["to1"], idleManager, WithReadBatchSize(batchSize))
 			assert.NoError(t, err)
 			stopped := f.Start()
 			// write some data such that the fromBufferPartition can be empty, that is toBuffer gets full
diff --git a/pkg/sinks/kafka/kafka_test.go b/pkg/sinks/kafka/kafka_test.go
index 252eebf2..6575e1f6 100644
--- a/pkg/sinks/kafka/kafka_test.go
+++ b/pkg/sinks/kafka/kafka_test.go
@@ -51,7 +51,8 @@ func TestWriteSuccessToKafka(t *testing.T) {
 		Replica: 0,
 	}
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferList([]string{vertex.Spec.Name})
-	toKafka.isdf, err = sinkforward.NewDataForward(vertexInstance, fromStep, toKafka, fetchWatermark, publishWatermark["testVertex"], wmb.NewIdleManager(1))
+	idleManager, _ := wmb.NewIdleManager(1, 1)
+	toKafka.isdf, err = sinkforward.NewDataForward(vertexInstance, fromStep, toKafka, fetchWatermark, publishWatermark["testVertex"], idleManager)
 	assert.NoError(t, err)
 	toKafka.kafkaSink = vertex.Spec.Sink.Kafka
 	toKafka.name = "Test"
@@ -108,7 +109,8 @@ func TestWriteFailureToKafka(t *testing.T) {
 	}
 	toSteps := map[string][]isb.BufferWriter{vertex.Spec.Name: {toKafka}}
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-	toKafka.isdf, err = sinkforward.NewDataForward(vertexInstance, fromStep, toKafka, fetchWatermark, publishWatermark["testVertex"], wmb.NewIdleManager(1))
+	idleManager, _ := wmb.NewIdleManager(1, 1)
+	toKafka.isdf, err = sinkforward.NewDataForward(vertexInstance, fromStep, toKafka, fetchWatermark, publishWatermark["testVertex"], idleManager)
 	assert.NoError(t, err)
 	toKafka.name = "Test"
 	toKafka.topic = "topic-1"
diff --git a/pkg/sinks/logger/log_test.go b/pkg/sinks/logger/log_test.go
index 7f5d7e82..38dd6427 100644
--- a/pkg/sinks/logger/log_test.go
+++ b/pkg/sinks/logger/log_test.go
@@ -57,7 +57,8 @@ func TestToLog_Start(t *testing.T) {
 		Replica: 0,
 	}
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferList([]string{vertex.Spec.Name})
-	s, err := NewToLog(vertexInstance, fromStep, fetchWatermark, publishWatermark[vertex.Spec.Name], wmb.NewIdleManager(1))
+	idleManager, _ := wmb.NewIdleManager(1, 1)
+	s, err := NewToLog(vertexInstance, fromStep, fetchWatermark, publishWatermark[vertex.Spec.Name], idleManager)
 	assert.NoError(t, err)
 
 	stopped := s.Start()
@@ -112,7 +113,8 @@ func TestToLog_Forward(t *testing.T) {
 			}
 
 			fetchWatermark1, publishWatermark1 := generic.BuildNoOpWatermarkProgressorsFromBufferList([]string{vertex1.Spec.Name})
-			logger1, _ := NewToLog(vertexInstance1, to1, fetchWatermark1, publishWatermark1[vertex1.Spec.Name], wmb.NewIdleManager(1))
+			idleManager0, _ := wmb.NewIdleManager(1, 1)
+			logger1, _ := NewToLog(vertexInstance1, to1, fetchWatermark1, publishWatermark1[vertex1.Spec.Name], idleManager0)
 			logger1Stopped := logger1.Start()
 
 			toSteps := map[string][]isb.BufferWriter{
@@ -122,7 +124,8 @@ func TestToLog_Forward(t *testing.T) {
 			writeMessages := testutils.BuildTestWriteMessages(int64(20), testStartTime)
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := sinkforward.NewDataForward(vertexInstance1, fromStep, to1, fetchWatermark, publishWatermark["sinks.logger1"], wmb.NewIdleManager(1), sinkforward.WithReadBatchSize(batchSize))
+			idleManager1, _ := wmb.NewIdleManager(1, 1)
+			f, err := sinkforward.NewDataForward(vertexInstance1, fromStep, to1, fetchWatermark, publishWatermark["sinks.logger1"], idleManager1, sinkforward.WithReadBatchSize(batchSize))
 			assert.NoError(t, err)
 
 			stopped := f.Start()
diff --git a/pkg/sources/forward/data_forward_test.go b/pkg/sources/forward/data_forward_test.go
index c2c78940..8d01f391 100644
--- a/pkg/sources/forward/data_forward_test.go
+++ b/pkg/sources/forward/data_forward_test.go
@@ -161,7 +161,8 @@ func TestNewDataForward(t *testing.T) {
 
 			fetchWatermark, _ := generic.BuildNoOpSourceWatermarkProgressorsFromBufferMap(toSteps)
 			noOpStores := buildNoOpToVertexStores(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, noOpStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, noOpStores, idleManager, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -260,7 +261,8 @@ func TestNewDataForward(t *testing.T) {
 			fetchWatermark := &testForwardFetcher{}
 			toVertexStores := buildToVertexWatermarkStores(toSteps)
 
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, &myForwardToAllTest{}, &myForwardToAllTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, &myForwardToAllTest{}, &myForwardToAllTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -412,7 +414,8 @@ func TestNewDataForward(t *testing.T) {
 			fetchWatermark := &testForwardFetcher{}
 			toVertexStores := buildToVertexWatermarkStores(toSteps)
 
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardDropTest{}, myForwardDropTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardDropTest{}, myForwardDropTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -576,7 +579,8 @@ func TestNewDataForward(t *testing.T) {
 			fetchWatermark := &testForwardFetcher{}
 			toVertexStores := buildToVertexWatermarkStores(toSteps)
 
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -711,7 +715,9 @@ func TestNewDataForward(t *testing.T) {
 
 			fetchWatermark, _ := generic.BuildNoOpSourceWatermarkProgressorsFromBufferMap(toSteps)
 			toVertexStores := buildNoOpToVertexStores(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardApplyTransformerErrTest{}, myForwardApplyTransformerErrTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardApplyTransformerErrTest{}, myForwardApplyTransformerErrTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
 			assert.False(t, to1.IsFull())
@@ -755,7 +761,9 @@ func TestNewDataForward(t *testing.T) {
 
 			fetchWatermark, _ := generic.BuildNoOpSourceWatermarkProgressorsFromBufferMap(toSteps)
 			toVertexStores := buildNoOpToVertexStores(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardApplyWhereToErrTest{}, myForwardApplyWhereToErrTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardApplyWhereToErrTest{}, myForwardApplyWhereToErrTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
 			assert.True(t, to1.IsEmpty())
@@ -797,7 +805,9 @@ func TestNewDataForward(t *testing.T) {
 
 			fetchWatermark, _ := generic.BuildNoOpSourceWatermarkProgressorsFromBufferMap(toSteps)
 			toVertexStores := buildNoOpToVertexStores(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardInternalErrTest{}, myForwardInternalErrTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardInternalErrTest{}, myForwardInternalErrTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
 			assert.False(t, to1.IsFull())
@@ -907,7 +917,8 @@ func TestDataForwardSinglePartition(t *testing.T) {
 	fetchWatermark := &testForwardFetcher{}
 	toVertexStores := buildNoOpToVertexStores(toSteps)
 
-	f, err := NewDataForward(vertexInstance, fromStep, toSteps, mySourceForwardTest{}, mySourceForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(5))
+	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+	f, err := NewDataForward(vertexInstance, fromStep, toSteps, mySourceForwardTest{}, mySourceForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(5))
 	assert.NoError(t, err)
 	assert.False(t, to1.IsFull())
 	assert.True(t, to1.IsEmpty())
@@ -961,7 +972,8 @@ func TestDataForwardMultiplePartition(t *testing.T) {
 	fetchWatermark := &testForwardFetcher{}
 	toVertexStores := buildNoOpToVertexStores(toSteps)
 
-	f, err := NewDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, mySourceForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(5))
+	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+	f, err := NewDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, mySourceForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(5))
 	assert.NoError(t, err)
 	assert.False(t, to11.IsFull())
 	assert.False(t, to12.IsFull())
@@ -1068,7 +1080,9 @@ func TestWriteToBuffer(t *testing.T) {
 
 			fetchWatermark, _ := generic.BuildNoOpSourceWatermarkProgressorsFromBufferMap(toSteps)
 			toVertexStores := buildNoOpToVertexStores(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(value.batchSize))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexStores, idleManager, WithReadBatchSize(value.batchSize))
 			assert.NoError(t, err)
 			assert.False(t, buffer.IsFull())
 			assert.True(t, buffer.IsEmpty())
diff --git a/pkg/sources/forward/shutdown_test.go b/pkg/sources/forward/shutdown_test.go
index 21ed2f89..8e61f345 100644
--- a/pkg/sources/forward/shutdown_test.go
+++ b/pkg/sources/forward/shutdown_test.go
@@ -88,7 +88,8 @@ func TestInterStepDataForward(t *testing.T) {
 
 			fetchWatermark, _ := generic.BuildNoOpSourceWatermarkProgressorsFromBufferMap(toSteps)
 			toVertexWmStores := buildNoOpToVertexStores(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexWmStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexWmStores, idleManager, WithReadBatchSize(batchSize))
 			assert.NoError(t, err)
 			stopped := f.Start()
 			// write some data but buffer is not full even though we are not reading
@@ -126,7 +127,9 @@ func TestInterStepDataForward(t *testing.T) {
 
 			fetchWatermark, _ := generic.BuildNoOpSourceWatermarkProgressorsFromBufferMap(toSteps)
 			toVertexWmStores := buildNoOpToVertexStores(toSteps)
-			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexWmStores, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, fetchWatermark, TestSourceWatermarkPublisher{}, toVertexWmStores, idleManager, WithReadBatchSize(batchSize))
 			assert.NoError(t, err)
 			stopped := f.Start()
 			// write some data such that the reader can be empty, that is toBuffer gets full
diff --git a/pkg/sources/generator/tickgen_test.go b/pkg/sources/generator/tickgen_test.go
index f8401c1a..0c78f039 100644
--- a/pkg/sources/generator/tickgen_test.go
+++ b/pkg/sources/generator/tickgen_test.go
@@ -76,7 +76,9 @@ func TestRead(t *testing.T) {
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"writer": publishWMStore,
 	}
-	mgen, err := NewMemGen(m, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, wmb.NewIdleManager(len(toBuffers)))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	mgen, err := NewMemGen(m, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, idleManager)
 	assert.NoError(t, err)
 	_ = mgen.Start()
 
@@ -135,7 +137,9 @@ func TestStop(t *testing.T) {
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"writer": publishWMStore,
 	}
-	mgen, err := NewMemGen(m, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, wmb.NewIdleManager(len(toBuffers)))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	mgen, err := NewMemGen(m, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, idleManager)
 	assert.NoError(t, err)
 	stop := mgen.Start()
 
diff --git a/pkg/sources/http/http_test.go b/pkg/sources/http/http_test.go
index 2a28bd42..69fe76a4 100644
--- a/pkg/sources/http/http_test.go
+++ b/pkg/sources/http/http_test.go
@@ -85,7 +85,9 @@ func Test_NewHTTP(t *testing.T) {
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"test": publishWMStores,
 	}
-	h, err := New(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStores, wmb.NewIdleManager(len(toBuffers)))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	h, err := New(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStores, idleManager)
 	assert.NoError(t, err)
 	assert.False(t, h.(*httpSource).ready)
 	assert.Equal(t, v.Spec.Name, h.GetName())
diff --git a/pkg/sources/kafka/handler_test.go b/pkg/sources/kafka/handler_test.go
index 1e9772e0..9aec1df0 100644
--- a/pkg/sources/kafka/handler_test.go
+++ b/pkg/sources/kafka/handler_test.go
@@ -81,7 +81,9 @@ func TestMessageHandling(t *testing.T) {
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"test": publishWMStore,
 	}
-	ks, _ := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, wmb.NewIdleManager(len(toBuffers)), WithLogger(logging.NewLogger()),
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	ks, _ := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, idleManager, WithLogger(logging.NewLogger()),
 		WithBufferSize(100), WithReadTimeOut(100*time.Millisecond))
 
 	msg := &sarama.ConsumerMessage{
diff --git a/pkg/sources/kafka/reader_test.go b/pkg/sources/kafka/reader_test.go
index f9d22cdf..15d844e2 100644
--- a/pkg/sources/kafka/reader_test.go
+++ b/pkg/sources/kafka/reader_test.go
@@ -59,7 +59,9 @@ func TestNewKafkasource(t *testing.T) {
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"testVertex": publishWMStore,
 	}
-	ks, err := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, wmb.NewIdleManager(len(toBuffers)), WithLogger(logging.NewLogger()), WithBufferSize(100), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	ks, err := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, idleManager, WithLogger(logging.NewLogger()), WithBufferSize(100), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
 
 	// no errors if everything is good.
 	assert.Nil(t, err)
@@ -102,7 +104,9 @@ func TestGroupNameOverride(t *testing.T) {
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"testVertex": publishWMStore,
 	}
-	ks, _ := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, wmb.NewIdleManager(len(toBuffers)), WithLogger(logging.NewLogger()), WithBufferSize(100), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	ks, _ := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, idleManager, WithLogger(logging.NewLogger()), WithBufferSize(100), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
 
 	assert.Equal(t, "default", ks.(*kafkaSource).groupName)
 
@@ -135,7 +139,9 @@ func TestDefaultBufferSize(t *testing.T) {
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"testVertex": publishWMStore,
 	}
-	ks, _ := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, wmb.NewIdleManager(len(toBuffers)), WithLogger(logging.NewLogger()), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	ks, _ := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, idleManager, WithLogger(logging.NewLogger()), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
 
 	assert.Equal(t, 100, ks.(*kafkaSource).handlerBuffer)
 
@@ -168,7 +174,9 @@ func TestBufferSizeOverrides(t *testing.T) {
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"testVertex": publishWMStore,
 	}
-	ks, _ := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, wmb.NewIdleManager(len(toBuffers)), WithLogger(logging.NewLogger()), WithBufferSize(110), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	ks, _ := NewKafkaSource(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStore, idleManager, WithLogger(logging.NewLogger()), WithBufferSize(110), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
 
 	assert.Equal(t, 110, ks.(*kafkaSource).handlerBuffer)
 
diff --git a/pkg/sources/nats/nats_test.go b/pkg/sources/nats/nats_test.go
index 14ea0fbe..30463620 100644
--- a/pkg/sources/nats/nats_test.go
+++ b/pkg/sources/nats/nats_test.go
@@ -83,7 +83,9 @@ func newInstance(t *testing.T, vi *dfv1.VertexInstance) (sourcer.Sourcer, error)
 	toVertexWmStores := map[string]store.WatermarkStore{
 		"testVertex": publishWMStores,
 	}
-	return New(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStores, wmb.NewIdleManager(len(toBuffers)), WithReadTimeout(1*time.Second))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toBuffers))
+	return New(vi, toBuffers, myForwardToAllTest{}, applier.Terminal, fetchWatermark, toVertexWmStores, publishWMStores, idleManager, WithReadTimeout(1*time.Second))
 }
 
 func Test_Single(t *testing.T) {
diff --git a/pkg/udf/forward/forward_test.go b/pkg/udf/forward/forward_test.go
index 820bee2d..0adf83df 100644
--- a/pkg/udf/forward/forward_test.go
+++ b/pkg/udf/forward/forward_test.go
@@ -142,7 +142,9 @@ func TestNewInterStepDataForward(t *testing.T) {
 			writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime)
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -244,7 +246,8 @@ func TestNewInterStepDataForward(t *testing.T) {
 				}
 			}()
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &myForwardToAllTest{}, &myForwardToAllTest{}, &myForwardToAllTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &myForwardToAllTest{}, &myForwardToAllTest{}, &myForwardToAllTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -409,7 +412,8 @@ func TestNewInterStepDataForward(t *testing.T) {
 				}
 			}()
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardDropTest{}, myForwardDropTest{}, myForwardDropTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardDropTest{}, myForwardDropTest{}, myForwardDropTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -586,7 +590,8 @@ func TestNewInterStepDataForward(t *testing.T) {
 				}
 			}()
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -720,7 +725,8 @@ func TestNewInterStepDataForward(t *testing.T) {
 			writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime)
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyUDFErrTest{}, myForwardApplyUDFErrTest{}, myForwardApplyUDFErrTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyUDFErrTest{}, myForwardApplyUDFErrTest{}, myForwardApplyUDFErrTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 
 			assert.NoError(t, err)
 			assert.False(t, to1.IsFull())
@@ -763,7 +769,8 @@ func TestNewInterStepDataForward(t *testing.T) {
 			}
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyWhereToErrTest{}, myForwardApplyWhereToErrTest{}, myForwardApplyWhereToErrTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyWhereToErrTest{}, myForwardApplyWhereToErrTest{}, myForwardApplyWhereToErrTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 
 			assert.NoError(t, err)
 			assert.True(t, to1.IsEmpty())
@@ -804,7 +811,8 @@ func TestNewInterStepDataForward(t *testing.T) {
 			}
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardInternalErrTest{}, myForwardInternalErrTest{}, myForwardInternalErrTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardInternalErrTest{}, myForwardInternalErrTest{}, myForwardInternalErrTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 
 			assert.NoError(t, err)
 			assert.False(t, to1.IsFull())
@@ -940,7 +948,8 @@ func TestNewInterStepDataForwardIdleWatermark(t *testing.T) {
 		}
 	}()
 
-	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(2))
+	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(2))
 	assert.NoError(t, err)
 	assert.False(t, to1.IsFull())
 	assert.True(t, to1.IsEmpty())
@@ -1118,7 +1127,8 @@ func TestNewInterStepDataForwardIdleWatermark_Reset(t *testing.T) {
 		}
 	}()
 
-	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(2))
+	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(2))
 	assert.NoError(t, err)
 	assert.False(t, to1.IsFull())
 	assert.True(t, to1.IsEmpty())
@@ -1345,7 +1355,8 @@ func TestInterStepDataForwardSinglePartition(t *testing.T) {
 	_, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 
 	// create a forwarder
-	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, mySourceForwardTest{}, mySourceForwardTest{}, mySourceForwardTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(5))
+	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, mySourceForwardTest{}, mySourceForwardTest{}, mySourceForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(5))
 	assert.NoError(t, err)
 	assert.False(t, to1.IsFull())
 	assert.True(t, to1.IsEmpty())
@@ -1396,7 +1407,9 @@ func TestInterStepDataForwardMultiplePartition(t *testing.T) {
 	_, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 
 	// create a forwarder
-	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, mySourceForwardTest{}, mySourceForwardTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(5))
+
+	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, mySourceForwardTest{}, mySourceForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(5))
 	assert.NoError(t, err)
 	assert.False(t, to11.IsFull())
 	assert.False(t, to12.IsFull())
@@ -1497,7 +1510,8 @@ func TestWriteToBuffer(t *testing.T) {
 			}
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(value.batchSize), WithUDFStreaming(value.streamEnabled))
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(value.batchSize), WithUDFStreaming(value.streamEnabled))
 			assert.NoError(t, err)
 			assert.False(t, buffer.IsFull())
 			assert.True(t, buffer.IsEmpty())
diff --git a/pkg/udf/forward/shutdown_test.go b/pkg/udf/forward/shutdown_test.go
index bf0781bf..3190505e 100644
--- a/pkg/udf/forward/shutdown_test.go
+++ b/pkg/udf/forward/shutdown_test.go
@@ -91,7 +91,9 @@ func TestInterStepDataForward(t *testing.T) {
 			}
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, myShutdownTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, myShutdownTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 			assert.NoError(t, err)
 			stopped := f.Start()
 			// write some data but buffer is not full even though we are not reading
@@ -129,7 +131,9 @@ func TestInterStepDataForward(t *testing.T) {
 			}
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, myShutdownTest{}, fetchWatermark, publishWatermark, wmb.NewIdleManager(len(toSteps)), WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, myShutdownTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
 			assert.NoError(t, err)
 			stopped := f.Start()
 			// write some data such that the fromBufferPartition can be empty, that is toBuffer gets full
diff --git a/pkg/watermark/wmb/idle_manager_test.go b/pkg/watermark/wmb/idle_manager_test.go
index f74b8bf5..203eb574 100644
--- a/pkg/watermark/wmb/idle_manager_test.go
+++ b/pkg/watermark/wmb/idle_manager_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package wmb
 
 import (
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -31,13 +32,316 @@ func TestNewIdleManager(t *testing.T) {
 			return int64(100)
 		})
 	)
-	idleManager := NewIdleManager(10)
-	assert.NotNil(t, idleManager)
-	assert.True(t, idleManager.NeedToSendCtrlMsg(toBufferName))
-	idleManager.Update(toBufferName, o)
-	getOffset := idleManager.Get(toBufferName)
+	im, _ := NewIdleManager(1, 10)
+	assert.NotNil(t, im)
+	assert.True(t, im.NeedToSendCtrlMsg(toBufferName))
+	im.Update(PARTITION_0, toBufferName, o)
+	getOffset := im.Get(toBufferName)
 	assert.Equal(t, o.String(), getOffset.String())
-	assert.False(t, idleManager.NeedToSendCtrlMsg(toBufferName))
-	idleManager.Reset(toBufferName)
-	assert.True(t, idleManager.NeedToSendCtrlMsg(toBufferName))
+	assert.False(t, im.NeedToSendCtrlMsg(toBufferName)) // already send one ctrlMsg, skip
+	im.MarkActive(PARTITION_0, toBufferName)
+	assert.False(t, im.NeedToSendCtrlMsg(toBufferName)) // forwarder is active, skip
+	im.MarkIdle(PARTITION_0, toBufferName)
+	assert.True(t, im.NeedToSendCtrlMsg(toBufferName))
+
+	type args struct {
+		numOfFromPartitions int
+		numOfToPartitions   int
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    IdleManager
+		wantErr string
+	}{
+		{
+			name: "good",
+			args: args{
+				numOfFromPartitions: 2,
+				numOfToPartitions:   2,
+			},
+			want: &idleManager{
+				wmbOffset:                  make(map[string]isb.Offset, 2),
+				forwarderActiveToPartition: make(map[string]uint64, 2),
+				lock:                       sync.RWMutex{},
+			},
+			wantErr: "",
+		},
+		{
+			name: "bad0",
+			args: args{
+				numOfFromPartitions: 0,
+				numOfToPartitions:   2,
+			},
+			want:    nil,
+			wantErr: "failed to create a new idle manager: numOfFromPartitions should be > 0 and < 64",
+		},
+		{
+			name: "bad64",
+			args: args{
+				numOfFromPartitions: 64,
+				numOfToPartitions:   2,
+			},
+			want:    nil,
+			wantErr: "failed to create a new idle manager: numOfFromPartitions should be > 0 and < 64",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := NewIdleManager(tt.args.numOfFromPartitions, tt.args.numOfToPartitions)
+			if err != nil {
+				assert.Equal(t, tt.wantErr, err.Error())
+			} else {
+				// if err is nil, then wantErr must be empty
+				assert.Equal(t, "", tt.wantErr)
+			}
+			assert.Equalf(t, tt.want, got, "NewIdleManager(%v, %v)", tt.args.numOfFromPartitions, tt.args.numOfToPartitions)
+		})
+	}
+}
+
+func Test_idleManager_Get(t *testing.T) {
+	type fields struct {
+		wmbOffset                  map[string]isb.Offset
+		forwarderActiveToPartition map[string]uint64
+	}
+	type args struct {
+		toBufferPartitionName string
+	}
+	var (
+		o = isb.SimpleIntOffset(func() int64 {
+			return int64(100)
+		})
+		sequence, _ = o.Sequence()
+	)
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   int64
+	}{
+		{
+			name: "good",
+			fields: fields{
+				wmbOffset: map[string]isb.Offset{
+					"testToBufferPartition0": o,
+				},
+				forwarderActiveToPartition: map[string]uint64{
+					// all forwarders are inactive
+					"testToBufferPartition0": 0,
+				},
+			},
+			args: args{
+				toBufferPartitionName: "testToBufferPartition0",
+			},
+			want: sequence,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			im := &idleManager{
+				wmbOffset:                  tt.fields.wmbOffset,
+				forwarderActiveToPartition: tt.fields.forwarderActiveToPartition,
+				lock:                       sync.RWMutex{},
+			}
+			gotOffset := im.Get(tt.args.toBufferPartitionName)
+			gotSequence, err := gotOffset.Sequence()
+			assert.Nil(t, err)
+			assert.Equalf(t, tt.want, gotSequence, "Get(%v)", tt.args.toBufferPartitionName)
+		})
+	}
+}
+
+func Test_idleManager_NeedToSendCtrlMsg(t *testing.T) {
+	type fields struct {
+		wmbOffset                  map[string]isb.Offset
+		forwarderActiveToPartition map[string]uint64
+	}
+	type args struct {
+		toBufferPartitionName string
+	}
+	var o = isb.SimpleIntOffset(func() int64 {
+		return int64(100)
+	})
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   bool
+	}{
+		{
+			name: "send_new_idle",
+			fields: fields{
+				wmbOffset: map[string]isb.Offset{},
+				forwarderActiveToPartition: map[string]uint64{
+					// all forwarders are inactive
+					"testToBufferPartition0": 0,
+				},
+			},
+			args: args{
+				toBufferPartitionName: "testToBufferPartition0",
+			},
+			want: true,
+		},
+		{
+			name: "dont_send_already_idle",
+			fields: fields{
+				wmbOffset: map[string]isb.Offset{
+					"testToBufferPartition0": o,
+				},
+				forwarderActiveToPartition: map[string]uint64{
+					// all forwarders are inactive
+					"testToBufferPartition0": 0,
+				},
+			},
+			args: args{
+				toBufferPartitionName: "testToBufferPartition0",
+			},
+			want: false,
+		},
+		{
+			name: "dont_send_at_least_one_active",
+			fields: fields{
+				wmbOffset: map[string]isb.Offset{},
+				forwarderActiveToPartition: map[string]uint64{
+					// forwarder0 is inactive while forwarder1 is active -> binary 10 -> decimal 2
+					"testToBufferPartition0": 2,
+				},
+			},
+			args: args{
+				toBufferPartitionName: "testToBufferPartition0",
+			},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			im := &idleManager{
+				wmbOffset:                  tt.fields.wmbOffset,
+				forwarderActiveToPartition: tt.fields.forwarderActiveToPartition,
+				lock:                       sync.RWMutex{},
+			}
+			assert.Equalf(t, tt.want, im.NeedToSendCtrlMsg(tt.args.toBufferPartitionName), "NeedToSendCtrlMsg(%v)", tt.args.toBufferPartitionName)
+		})
+	}
+}
+
+func Test_idleManager_Reset(t *testing.T) {
+	type fields struct {
+		wmbOffset                  map[string]isb.Offset
+		forwarderActiveToPartition map[string]uint64
+	}
+	type args struct {
+		fromBufferPartitionIndex int32
+		toBufferPartitionName    string
+	}
+	var (
+		o = isb.SimpleIntOffset(func() int64 {
+			return int64(100)
+		})
+		sequence, _ = o.Sequence()
+		testMap     = map[string]uint64{
+			// forwarder0 is inactive while forwarder1 is active -> binary 10 -> decimal 2
+			"testToBufferPartition0": 2,
+		}
+	)
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+	}{
+		{
+			name: "good",
+			fields: fields{
+				wmbOffset: map[string]isb.Offset{
+					"testToBufferPartition0": o, // should be reset to nil
+				},
+				forwarderActiveToPartition: testMap,
+			},
+			args: args{
+				fromBufferPartitionIndex: 0,
+				toBufferPartitionName:    "testToBufferPartition0",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			im := &idleManager{
+				wmbOffset:                  tt.fields.wmbOffset,
+				forwarderActiveToPartition: tt.fields.forwarderActiveToPartition,
+				lock:                       sync.RWMutex{},
+			}
+			// before
+			beforeReset, err := im.Get(tt.args.toBufferPartitionName).Sequence()
+			assert.Nil(t, err)
+			assert.Equal(t, sequence, beforeReset)
+			assert.Equal(t, uint64(2), testMap[tt.args.toBufferPartitionName])
+			// reset
+			im.MarkActive(tt.args.fromBufferPartitionIndex, tt.args.toBufferPartitionName)
+			// after
+			assert.Nil(t, im.Get(tt.args.toBufferPartitionName))
+			// now both bits should be 1 -> 11 -> 3
+			assert.Equal(t, uint64(3), testMap[tt.args.toBufferPartitionName])
+		})
+	}
+}
+
+func Test_idleManager_Update(t *testing.T) {
+	type fields struct {
+		wmbOffset                  map[string]isb.Offset
+		forwarderActiveToPartition map[string]uint64
+	}
+	type args struct {
+		fromBufferPartitionIndex int32
+		toBufferPartitionName    string
+		newOffset                isb.Offset
+	}
+	var (
+		o = isb.SimpleIntOffset(func() int64 {
+			return int64(100)
+		})
+		sequence, _ = o.Sequence()
+		testMap     = map[string]uint64{
+			// forwarder0 is active while forwarder1 is inactive -> binary 01 -> decimal 1
+			"testToBufferPartition0": 1,
+		}
+	)
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+	}{
+		{
+			name: "good",
+			fields: fields{
+				wmbOffset:                  map[string]isb.Offset{},
+				forwarderActiveToPartition: testMap,
+			},
+			args: args{
+				fromBufferPartitionIndex: 0,
+				toBufferPartitionName:    "testToBufferPartition0",
+				newOffset:                o,
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			im := &idleManager{
+				wmbOffset:                  tt.fields.wmbOffset,
+				forwarderActiveToPartition: tt.fields.forwarderActiveToPartition,
+				lock:                       sync.RWMutex{},
+			}
+			// before
+			assert.Nil(t, im.Get(tt.args.toBufferPartitionName))
+			assert.Equal(t, uint64(1), testMap[tt.args.toBufferPartitionName])
+			// update
+			im.Update(tt.args.fromBufferPartitionIndex, tt.args.toBufferPartitionName, tt.args.newOffset)
+			// after
+			afterUpdate, err := im.Get(tt.args.toBufferPartitionName).Sequence()
+			assert.Nil(t, err)
+			assert.Equal(t, sequence, afterUpdate)
+			// now the bits should be 0 -> 00 -> 0
+			assert.Equal(t, uint64(0), testMap[tt.args.toBufferPartitionName])
+		})
+	}
 }