Repository: numaproj/numaflow
Pull request #170 (2022-09-14 02:15:39)

chore: publish watermark to sink bucket and refactor.
Explain what this PR does.

<!--

Before you push your changes:

* Run `make pre-push -B` to fix codegen and lint problems.

Then, you MUST:

* Sign-off your commit (otherwise the DCO check will fail).
* Use [a conventional commit message](https://www.conventionalcommits.org/en/v1.0.0/) (otherwise the commit message check will fail).

If you did not do this, reset all your commit and replace them with a single commit:

```
git reset HEAD~1 ;# change 1 to how many commits you made
git commit --signoff -m 'feat: my feat. Fixes #1234'
```

When creating your PR: 

* Make sure that "Fixes #" or "Closes #" is in both the PR title (for release notes) and description (to automatically link and close the issue).
* Say how you tested your changes. If you changed the UI, attach screenshots.
* Set your PR as a draft initially.
* Your PR needs to pass the required checks before it can be approved. 
* Once required tests have passed, mark your PR "Ready for review".

If changes were requested, once you've made them, you MUST dismiss the review to get it reviewed again.

-->



Patch:
diff --git a/pkg/daemon/server/service/pipeline_watermark_query.go b/pkg/daemon/server/service/pipeline_watermark_query.go
index d9072491..e46d999b 100644
--- a/pkg/daemon/server/service/pipeline_watermark_query.go
+++ b/pkg/daemon/server/service/pipeline_watermark_query.go
@@ -13,6 +13,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
 )
 
@@ -85,8 +86,7 @@ func createWatermarkFetcher(ctx context.Context, pipelineName string, fromBuffer
 	if err != nil {
 		return nil, err
 	}
-	var fetchWmWatchers = generic.BuildFetchWMWatchers(hbWatch, otWatch)
-	fetchWatermark := generic.NewGenericEdgeFetch(ctx, vertexName, fetchWmWatchers)
+	fetchWatermark := generic.NewGenericEdgeFetch(ctx, vertexName, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
 	return fetchWatermark, nil
 }
 
diff --git a/pkg/isb/forward/forward.go b/pkg/isb/forward/forward.go
index c042fe96..1d91f290 100644
--- a/pkg/isb/forward/forward.go
+++ b/pkg/isb/forward/forward.go
@@ -252,10 +252,12 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 	// forward the highest watermark to all the edges to avoid idle edge problem
 	// TODO: sort and get the highest value
 	// TODO: Should also publish to those edges without writing (fall out of conditional forwarding)?
-	for edgeName, offsets := range writeOffsets {
-		if len(offsets) > 0 {
-			if publisher, ok := isdf.publishWatermark[edgeName]; ok {
+	for bufferName, offsets := range writeOffsets {
+		if publisher, ok := isdf.publishWatermark[bufferName]; ok {
+			if len(offsets) > 0 {
 				publisher.PublishWatermark(processorWM, offsets[len(offsets)-1])
+			} else { // This only happens on sink vertex, and it does not care about the offset during watermark publishing
+				publisher.PublishWatermark(processorWM, nil)
 			}
 		}
 	}
diff --git a/pkg/sinks/sink.go b/pkg/sinks/sink.go
index a786fdf1..00c17dea 100644
--- a/pkg/sinks/sink.go
+++ b/pkg/sinks/sink.go
@@ -60,7 +60,7 @@ func (u *SinkProcessor) Start(ctx context.Context) error {
 			readOptions = append(readOptions, jetstreamisb.WithReadTimeOut(x.ReadTimeout.Duration))
 		}
 		// build watermark progressors
-		fetchWatermark, publishWatermark, err = jetstream.BuildJetStreamWatermarkProgressors(ctx, u.VertexInstance)
+		fetchWatermark, publishWatermark, err = jetstream.BuildWatermarkProgressors(ctx, u.VertexInstance)
 		if err != nil {
 			return err
 		}
@@ -71,7 +71,7 @@ func (u *SinkProcessor) Start(ctx context.Context) error {
 			return err
 		}
 	default:
-		return fmt.Errorf("unrecognized isbs type %q", u.ISBSvcType)
+		return fmt.Errorf("unrecognized isb svc type %q", u.ISBSvcType)
 	}
 
 	sinker, err := u.getSinker(reader, log, fetchWatermark, publishWatermark)
@@ -79,7 +79,7 @@ func (u *SinkProcessor) Start(ctx context.Context) error {
 		return fmt.Errorf("failed to find a sink, errpr: %w", err)
 	}
 
-	log.Infow("Start processing sink messages", zap.String("isbs", string(u.ISBSvcType)), zap.String("from", fromBufferName))
+	log.Infow("Start processing sink messages", zap.String("isbsvc", string(u.ISBSvcType)), zap.String("from", fromBufferName))
 	stopped := sinker.Start()
 	wg := &sync.WaitGroup{}
 	wg.Add(1)
diff --git a/pkg/sinks/udsink/sink.go b/pkg/sinks/udsink/sink.go
index 86ff79ce..c3bfd2fd 100644
--- a/pkg/sinks/udsink/sink.go
+++ b/pkg/sinks/udsink/sink.go
@@ -2,6 +2,8 @@ package udsink
 
 import (
 	"context"
+	"time"
+
 	sinksdk "github.com/numaproj/numaflow-go/sink"
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
@@ -12,7 +14,6 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 	"go.uber.org/zap"
-	"time"
 )
 
 type userDefinedSink struct {
@@ -55,7 +56,7 @@ func NewUserDefinedSink(vertex *dfv1.Vertex, fromBuffer isb.BufferReader, fetchW
 	}
 	contentType := sharedutil.LookupEnvStringOr(dfv1.EnvUDSinkContentType, string(dfv1.MsgPackType))
 	s.udsink = NewUDSHTTPBasedUDSink(dfv1.PathVarRun+"/udsink.sock", withTimeout(20*time.Second), withContentType(dfv1.ContentType(contentType)))
-	isdf, err := forward.NewInterStepDataForward(vertex, fromBuffer, map[string]isb.BufferWriter{name: s}, forward.All, applier.Terminal, fetchWatermark, publishWatermark, forwardOpts...)
+	isdf, err := forward.NewInterStepDataForward(vertex, fromBuffer, map[string]isb.BufferWriter{vertex.GetToBuffers()[0].Name: s}, forward.All, applier.Terminal, fetchWatermark, publishWatermark, forwardOpts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/sources/generator/tickgen.go b/pkg/sources/generator/tickgen.go
index 4e0483b8..33df4b29 100644
--- a/pkg/sources/generator/tickgen.go
+++ b/pkg/sources/generator/tickgen.go
@@ -19,9 +19,9 @@ import (
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
-	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
 var log = logging.NewLogger()
@@ -120,7 +120,7 @@ func NewMemGen(vertexInstance *dfv1.VertexInstance,
 	msgSize int32,
 	timeunit time.Duration,
 	writers []isb.BufferWriter,
-	fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores *generic.PublishWMStores, // watermarks
+	fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores store.WatermarkStorer, // watermarks
 	opts ...Option) (*memgen, error) {
 	gensrc := &memgen{
 		rpu:            rpu,
diff --git a/pkg/sources/generator/watermark.go b/pkg/sources/generator/watermark.go
index a639b9e7..74e66cd6 100644
--- a/pkg/sources/generator/watermark.go
+++ b/pkg/sources/generator/watermark.go
@@ -4,14 +4,14 @@ import (
 	"fmt"
 
 	sharedutil "github.com/numaproj/numaflow/pkg/shared/util"
-	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
-func (mg *memgen) buildSourceWatermarkPublisher(publishWMStores *generic.PublishWMStores) publish.Publisher {
+func (mg *memgen) buildSourceWatermarkPublisher(publishWMStores store.WatermarkStorer) publish.Publisher {
 	// 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.HBStore, publishWMStores.OTStore, publish.IsSource(), publish.WithDelay(sharedutil.GetWatermarkMaxDelay()))
+	return publish.NewPublish(mg.lifecycleCtx, processorEntity, publishWMStores, publish.IsSource(), publish.WithDelay(sharedutil.GetWatermarkMaxDelay()))
 }
diff --git a/pkg/sources/http/http.go b/pkg/sources/http/http.go
index 8854aaf5..c4868548 100644
--- a/pkg/sources/http/http.go
+++ b/pkg/sources/http/http.go
@@ -21,9 +21,9 @@ import (
 	sharedutil "github.com/numaproj/numaflow/pkg/shared/util"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
-	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
 type httpSource struct {
@@ -68,7 +68,7 @@ func WithBufferSize(s int) Option {
 	}
 }
 
-func New(vertexInstance *dfv1.VertexInstance, writers []isb.BufferWriter, fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores *generic.PublishWMStores, opts ...Option) (*httpSource, error) {
+func New(vertexInstance *dfv1.VertexInstance, writers []isb.BufferWriter, fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores store.WatermarkStorer, opts ...Option) (*httpSource, error) {
 	h := &httpSource{
 		name:         vertexInstance.Vertex.Spec.Name,
 		pipelineName: vertexInstance.Vertex.Spec.PipelineName,
@@ -191,7 +191,7 @@ func New(vertexInstance *dfv1.VertexInstance, writers []isb.BufferWriter, fetchW
 	h.cancelfn = cancel
 	entityName := fmt.Sprintf("%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica)
 	processorEntity := processor.NewProcessorEntity(entityName)
-	h.sourcePublishWM = publish.NewPublish(ctx, processorEntity, publishWMStores.HBStore, publishWMStores.OTStore, publish.IsSource(), publish.WithDelay(sharedutil.GetWatermarkMaxDelay()))
+	h.sourcePublishWM = publish.NewPublish(ctx, processorEntity, publishWMStores, publish.IsSource(), publish.WithDelay(sharedutil.GetWatermarkMaxDelay()))
 	return h, nil
 }
 
diff --git a/pkg/sources/kafka/reader.go b/pkg/sources/kafka/reader.go
index 530bbc92..923d8238 100644
--- a/pkg/sources/kafka/reader.go
+++ b/pkg/sources/kafka/reader.go
@@ -20,9 +20,9 @@ import (
 	sharedutil "github.com/numaproj/numaflow/pkg/shared/util"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
-	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
 type KafkaSource struct {
@@ -61,7 +61,7 @@ type KafkaSource struct {
 	// source watermark publishers for different partitions
 	sourcePublishWMs map[int32]publish.Publisher
 	// source watermark publisher stores
-	srcPublishWMStores *generic.PublishWMStores
+	srcPublishWMStores store.WatermarkStorer
 	lock               *sync.RWMutex
 }
 
@@ -146,7 +146,7 @@ func (r *KafkaSource) loadSourceWartermarkPublisher(partitionID int32) publish.P
 	}
 	entityName := fmt.Sprintf("%s-%s-%d", r.pipelineName, r.name, partitionID)
 	processorEntity := processor.NewProcessorEntity(entityName)
-	sourcePublishWM := publish.NewPublish(r.lifecyclectx, processorEntity, r.srcPublishWMStores.HBStore, r.srcPublishWMStores.OTStore, publish.IsSource(), publish.WithDelay(sharedutil.GetWatermarkMaxDelay()))
+	sourcePublishWM := publish.NewPublish(r.lifecyclectx, processorEntity, r.srcPublishWMStores, publish.IsSource(), publish.WithDelay(sharedutil.GetWatermarkMaxDelay()))
 	r.sourcePublishWMs[partitionID] = sourcePublishWM
 	return sourcePublishWM
 }
@@ -239,7 +239,7 @@ func (r *KafkaSource) Pending(ctx context.Context) (int64, error) {
 }
 
 // NewKafkaSource returns a KafkaSource reader based on Kafka Consumer Group .
-func NewKafkaSource(vertexInstance *dfv1.VertexInstance, writers []isb.BufferWriter, fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores *generic.PublishWMStores, opts ...Option) (*KafkaSource, error) {
+func NewKafkaSource(vertexInstance *dfv1.VertexInstance, writers []isb.BufferWriter, fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores store.WatermarkStorer, opts ...Option) (*KafkaSource, error) {
 	source := vertexInstance.Vertex.Spec.Source.Kafka
 	kafkasource := &KafkaSource{
 		name:               vertexInstance.Vertex.Spec.Name,
diff --git a/pkg/sources/source.go b/pkg/sources/source.go
index ac3c0f19..0e66a3cc 100644
--- a/pkg/sources/source.go
+++ b/pkg/sources/source.go
@@ -22,6 +22,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/generic/jetstream"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 )
 
@@ -39,7 +40,7 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 	// watermark variables no-op initialization
 	// publishWatermark is a map representing a progressor per edge, we are initializing them to a no-op progressor
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromEdgeList(generic.GetBufferNameList(sp.VertexInstance.Vertex.GetToBuffers()))
-	var publishWMStore = generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	var sourcePublisherStores = store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	var err error
 
 	switch sp.ISBSvcType {
@@ -59,13 +60,15 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 			writers = append(writers, writer)
 		}
 	case dfv1.ISBSvcTypeJetStream:
-		// build the right watermark progessor if watermark is enabled
 		// build watermark progressors
-		fetchWatermark, publishWatermark, publishWMStore, err = jetstream.BuildJetStreamWatermarkProgressorsForSource(ctx, sp.VertexInstance)
+		fetchWatermark, publishWatermark, err = jetstream.BuildWatermarkProgressors(ctx, sp.VertexInstance)
+		if err != nil {
+			return err
+		}
+		sourcePublisherStores, err = jetstream.BuildSourcePublisherStores(ctx, sp.VertexInstance)
 		if err != nil {
 			return err
 		}
-
 		for _, e := range sp.VertexInstance.Vertex.Spec.ToEdges {
 			writeOpts := []jetstreamisb.WriteOption{
 				jetstreamisb.WithUsingWriteInfoAsRate(true),
@@ -89,7 +92,7 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 		return fmt.Errorf("unrecognized isb svc type %q", sp.ISBSvcType)
 	}
 
-	sourcer, err := sp.getSourcer(writers, fetchWatermark, publishWatermark, publishWMStore, log)
+	sourcer, err := sp.getSourcer(writers, fetchWatermark, publishWatermark, sourcePublisherStores, log)
 	if err != nil {
 		return fmt.Errorf("failed to find a sourcer, error: %w", err)
 	}
@@ -129,7 +132,7 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 }
 
 // getSourcer is used to send the sourcer information
-func (sp *SourceProcessor) getSourcer(writers []isb.BufferWriter, fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores *generic.PublishWMStores, logger *zap.SugaredLogger) (Sourcer, error) {
+func (sp *SourceProcessor) getSourcer(writers []isb.BufferWriter, fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores store.WatermarkStorer, logger *zap.SugaredLogger) (Sourcer, error) {
 	src := sp.VertexInstance.Vertex.Spec.Source
 
 	if x := src.Generator; x != nil {
diff --git a/pkg/udf/udf.go b/pkg/udf/udf.go
index 674aa1f7..7dca6ab8 100644
--- a/pkg/udf/udf.go
+++ b/pkg/udf/udf.go
@@ -81,7 +81,7 @@ func (u *UDFProcessor) Start(ctx context.Context) error {
 		}
 
 		// build watermark progressors
-		fetchWatermark, publishWatermark, err = jetstream.BuildJetStreamWatermarkProgressors(ctx, u.VertexInstance)
+		fetchWatermark, publishWatermark, err = jetstream.BuildWatermarkProgressors(ctx, u.VertexInstance)
 		if err != nil {
 			return err
 		}
@@ -138,7 +138,7 @@ func (u *UDFProcessor) Start(ctx context.Context) error {
 			log.Warnw("Failed to close gRPC client conn", zap.Error(err))
 		}
 	}()
-	log.Infow("Start processing udf messages", zap.String("isbs", string(u.ISBSvcType)), zap.String("from", fromBufferName), zap.Any("to", toBuffers))
+	log.Infow("Start processing udf messages", zap.String("isbsvc", string(u.ISBSvcType)), zap.String("from", fromBufferName), zap.Any("to", toBuffers))
 
 	opts := []forward.Option{forward.WithLogger(log)}
 	if x := u.VertexInstance.Vertex.Spec.Limits; x != nil {
diff --git a/pkg/watermark/fetch/processor_manager.go b/pkg/watermark/fetch/processor_manager.go
index 53e38531..e52dbe45 100644
--- a/pkg/watermark/fetch/processor_manager.go
+++ b/pkg/watermark/fetch/processor_manager.go
@@ -32,7 +32,7 @@ type ProcessorManager struct {
 }
 
 // NewProcessorManager returns a new ProcessorManager instance
-func NewProcessorManager(ctx context.Context, hbWatcher store.WatermarkKVWatcher, otWatcher store.WatermarkKVWatcher, inputOpts ...ProcessorManagerOption) *ProcessorManager {
+func NewProcessorManager(ctx context.Context, watermarkStoreWatcher store.WatermarkStoreWatcher, inputOpts ...ProcessorManagerOption) *ProcessorManager {
 	opts := &processorManagerOptions{
 		podHeartbeatRate:         5,
 		refreshingProcessorsRate: 5,
@@ -44,8 +44,8 @@ func NewProcessorManager(ctx context.Context, hbWatcher store.WatermarkKVWatcher
 
 	v := &ProcessorManager{
 		ctx:        ctx,
-		hbWatcher:  hbWatcher,
-		otWatcher:  otWatcher,
+		hbWatcher:  watermarkStoreWatcher.HeartbeatWatcher(),
+		otWatcher:  watermarkStoreWatcher.OffsetTimelineWatcher(),
 		heartbeat:  NewProcessorHeartbeat(),
 		processors: make(map[string]*ProcessorToFetch),
 		log:        logging.FromContext(ctx),
diff --git a/pkg/watermark/generic/fetcher.go b/pkg/watermark/generic/generic.go
similarity index 52%
rename from pkg/watermark/generic/fetcher.go
rename to pkg/watermark/generic/generic.go
index 607175fa..30bbd936 100644
--- a/pkg/watermark/generic/fetcher.go
+++ b/pkg/watermark/generic/generic.go
@@ -7,30 +7,16 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
-// FetchWMWatchers has the watcher information required for fetching watermarks.
-type FetchWMWatchers struct {
-	HBWatch store.WatermarkKVWatcher
-	OTWatch store.WatermarkKVWatcher
-}
-
-// BuildFetchWMWatchers builds the FetchWMWatchers
-func BuildFetchWMWatchers(hbWatch store.WatermarkKVWatcher, otWatch store.WatermarkKVWatcher) FetchWMWatchers {
-	return FetchWMWatchers{
-		HBWatch: hbWatch,
-		OTWatch: otWatch,
-	}
-}
-
 // NewGenericEdgeFetch returns a Fetcher, where bufferName is the from buffer of the vertex that is currently processing.
 // fetchWM is a struct for retrieving both the heartbeat
 // and the offset watermark timeline (Vn-1 vertex).
-func NewGenericEdgeFetch(ctx context.Context, bufferName string, fetchWM FetchWMWatchers) fetch.Fetcher {
-	processorManager := fetch.NewProcessorManager(ctx, fetchWM.HBWatch, fetchWM.OTWatch)
+func NewGenericEdgeFetch(ctx context.Context, bufferName string, storeWatcher store.WatermarkStoreWatcher) fetch.Fetcher {
+	processorManager := fetch.NewProcessorManager(ctx, storeWatcher)
 	return fetch.NewEdgeFetcher(ctx, bufferName, processorManager)
 }
 
 // NewGenericSourceFetch returns Fetcher, where sourceBufferName is the source buffer of the source vertex.
-func NewGenericSourceFetch(ctx context.Context, sourceBufferName string, fetchWM FetchWMWatchers) fetch.Fetcher {
-	processorManager := fetch.NewProcessorManager(ctx, fetchWM.HBWatch, fetchWM.OTWatch)
+func NewGenericSourceFetch(ctx context.Context, sourceBufferName string, storeWatcher store.WatermarkStoreWatcher) fetch.Fetcher {
+	processorManager := fetch.NewProcessorManager(ctx, storeWatcher)
 	return fetch.NewSourceFetcher(ctx, sourceBufferName, processorManager)
 }
diff --git a/pkg/watermark/generic/jetstream/generic.go b/pkg/watermark/generic/jetstream/generic.go
index 65737871..200a4254 100644
--- a/pkg/watermark/generic/jetstream/generic.go
+++ b/pkg/watermark/generic/jetstream/generic.go
@@ -7,6 +7,8 @@ import (
 	"fmt"
 
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/processor"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 
 	"github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isbsvc"
@@ -15,15 +17,14 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
-	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 )
 
-// BuildJetStreamWatermarkProgressors is used to populate fetchWatermark, and a map of publishWatermark with edge name as the key.
+// BuildWatermarkProgressors is used to populate fetchWatermark, and a map of publishWatermark with edge name as the key.
 // These are used as watermark progressors in the pipeline, and is attached to each edge of the vertex.
 // Fetcher has one-to-one relationship , whereas we have multiple publishers as the vertex can read only from one edge,
 // and it can write to many.
 // The function is used only when watermarking is enabled on the pipeline.
-func BuildJetStreamWatermarkProgressors(ctx context.Context, vertexInstance *v1alpha1.VertexInstance) (fetch.Fetcher, map[string]publish.Publisher, error) {
+func BuildWatermarkProgressors(ctx context.Context, vertexInstance *v1alpha1.VertexInstance) (fetch.Fetcher, map[string]publish.Publisher, error) {
 	// if watermark is not enabled, use no-op.
 	if !sharedutil.IsWatermarkEnabled() {
 		fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromEdgeList(generic.GetBufferNameList(vertexInstance.Vertex.GetToBuffers()))
@@ -46,16 +47,14 @@ func BuildJetStreamWatermarkProgressors(ctx context.Context, vertexInstance *v1a
 		return nil, nil, fmt.Errorf("failed at new OT KVJetStreamKVWatch, OTBucket: %s, %w", otBucket, err)
 	}
 
-	var fetchWmWatchers = generic.BuildFetchWMWatchers(hbWatch, otWatch)
 	var fetchWatermark fetch.Fetcher
 	if fromBuffer.Type == v1alpha1.SourceBuffer {
-		fetchWatermark = generic.NewGenericSourceFetch(ctx, fromBuffer.Name, fetchWmWatchers)
+		fetchWatermark = generic.NewGenericSourceFetch(ctx, fromBuffer.Name, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
 	} else {
-		fetchWatermark = generic.NewGenericEdgeFetch(ctx, fromBuffer.Name, fetchWmWatchers)
+		fetchWatermark = generic.NewGenericEdgeFetch(ctx, fromBuffer.Name, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
 	}
 
 	// Publisher map creation, we need a publisher per edge.
-
 	for _, buffer := range vertexInstance.Vertex.GetToBuffers() {
 		hbPublisherBucket := isbsvc.JetStreamProcessorBucket(pipelineName, buffer.Name)
 		// We create a separate Heartbeat bucket for each edge though it can be reused. We can reuse because heartbeat is at
@@ -71,52 +70,38 @@ func BuildJetStreamWatermarkProgressors(ctx context.Context, vertexInstance *v1a
 			return nil, nil, fmt.Errorf("failed at new OT Publish JetStreamKVStore, OTBucket: %s, %w", otStoreBucket, err)
 		}
 
-		var publishStores = generic.BuildPublishWMStores(hbStore, otStore)
 		var processorName = fmt.Sprintf("%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica)
-		publishWatermark[buffer.Name] = generic.NewGenericPublish(ctx, processorName, *publishStores)
+		publishEntity := processor.NewProcessorEntity(processorName)
+		opts := []publish.PublishOption{}
+		if buffer.Type == v1alpha1.SinkBuffer {
+			opts = append(opts, publish.IsSink())
+		}
+		publishWatermark[buffer.Name] = publish.NewPublish(ctx, publishEntity, store.BuildWatermarkStore(hbStore, otStore), opts...)
 	}
 
 	return fetchWatermark, publishWatermark, nil
 }
 
-// BuildJetStreamWatermarkProgressorsForSource is an extension of BuildJetStreamWatermarkProgressors to also return the publish stores. This is
-// for letting source implement as many publishers that it requires to progress the watermark monotonically for each individual processing entity.
-// Eg, watermark progresses independently and monotonically for each partition in a Kafka topic.
-func BuildJetStreamWatermarkProgressorsForSource(ctx context.Context, vertexInstance *v1alpha1.VertexInstance) (fetch.Fetcher, map[string]publish.Publisher, *generic.PublishWMStores, error) {
-	if vertexInstance == nil || vertexInstance.Vertex == nil {
-		return nil, nil, nil, fmt.Errorf("nil vertex instance")
-	}
+// BuildSourcePublisherStores builds the watermark stores for source publisher.
+func BuildSourcePublisherStores(ctx context.Context, vertexInstance *v1alpha1.VertexInstance) (store.WatermarkStorer, error) {
 	if !vertexInstance.Vertex.IsASource() {
-		return nil, nil, nil, fmt.Errorf("invalid vertex: not a source")
-	}
-	fetchWatermark, publishWatermark, err := BuildJetStreamWatermarkProgressors(ctx, vertexInstance)
-	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed to build JetStreamWatermarkProgressors, %w", err)
+		return nil, fmt.Errorf("not a source vertex")
 	}
-	// return no-ops if not enabled!f
-	if !sharedutil.IsWatermarkEnabled() {
-		publishWM := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
-		return fetchWatermark, publishWatermark, publishWM, nil
-	}
-
 	pipelineName := vertexInstance.Vertex.Spec.PipelineName
-
 	sourceBufferName := vertexInstance.Vertex.GetFromBuffers()[0].Name
 	// heartbeat
 	hbBucket := isbsvc.JetStreamProcessorBucket(pipelineName, sourceBufferName)
 	hbKVStore, err := jetstream.NewKVJetStreamKVStore(ctx, pipelineName, hbBucket, jsclient.NewInClusterJetStreamClient())
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed at new HB KVJetStreamKVStore for source, HeartbeatBucket: %s, %w", hbBucket, err)
+		return nil, fmt.Errorf("failed at new HB KVJetStreamKVStore for source, HeartbeatBucket: %s, %w", hbBucket, err)
 	}
 
 	// OT
 	otStoreBucket := isbsvc.JetStreamOTBucket(pipelineName, sourceBufferName)
 	otKVStore, err := jetstream.NewKVJetStreamKVStore(ctx, pipelineName, otStoreBucket, jsclient.NewInClusterJetStreamClient())
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed at new OT KVJetStreamKVStore for source, OTBucket: %s, %w", otStoreBucket, err)
+		return nil, fmt.Errorf("failed at new OT KVJetStreamKVStore for source, OTBucket: %s, %w", otStoreBucket, err)
 	}
-
-	// interface for publisher store (HB and OT)
-	publishWM := generic.BuildPublishWMStores(hbKVStore, otKVStore)
-	return fetchWatermark, publishWatermark, publishWM, nil
+	sourcePublishStores := store.BuildWatermarkStore(hbKVStore, otKVStore)
+	return sourcePublishStores, nil
 }
diff --git a/pkg/watermark/generic/publisher.go b/pkg/watermark/generic/publisher.go
deleted file mode 100644
index 9fb359bb..00000000
--- a/pkg/watermark/generic/publisher.go
+++ /dev/null
@@ -1,30 +0,0 @@
-package generic
-
-import (
-	"context"
-
-	"github.com/numaproj/numaflow/pkg/watermark/processor"
-	"github.com/numaproj/numaflow/pkg/watermark/publish"
-	"github.com/numaproj/numaflow/pkg/watermark/store"
-)
-
-// PublishWMStores stores the store information for publishing the watermark.
-type PublishWMStores struct {
-	HBStore store.WatermarkKVStorer
-	OTStore store.WatermarkKVStorer
-}
-
-// BuildPublishWMStores builds the PublishWMStores.
-func BuildPublishWMStores(hbStore store.WatermarkKVStorer, otStore store.WatermarkKVStorer) *PublishWMStores {
-	return &PublishWMStores{
-		HBStore: hbStore,
-		OTStore: otStore,
-	}
-}
-
-// NewGenericPublish returns GenericPublish, where processorName is the unique processor (pod) that is running on this
-// vertex, publishWM is a struct for storing both the heartbeat and the offset watermark timeline stores for the Vn vertex.
-func NewGenericPublish(ctx context.Context, processorName string, publishWM PublishWMStores) publish.Publisher {
-	publishEntity := processor.NewProcessorEntity(processorName)
-	return publish.NewPublish(ctx, publishEntity, publishWM.HBStore, publishWM.OTStore)
-}
diff --git a/pkg/watermark/publish/options.go b/pkg/watermark/publish/options.go
index 41fd62ea..89adb48a 100644
--- a/pkg/watermark/publish/options.go
+++ b/pkg/watermark/publish/options.go
@@ -13,7 +13,11 @@ type publishOptions struct {
 	// It should only be used in a source publisher.
 	delay time.Duration
 	// Whether it is source publisher or not
+	// isSource and isSink should not be both true
 	isSource bool
+	// Whether it is sink publisher or not
+	// isSource and isSink should not be both true
+	isSink bool
 }
 
 type PublishOption func(*publishOptions)
@@ -43,3 +47,10 @@ func IsSource() PublishOption {
 		opts.isSource = true
 	}
 }
+
+// IsSink indicates it's a sink publisher
+func IsSink() PublishOption {
+	return func(opts *publishOptions) {
+		opts.isSink = true
+	}
+}
diff --git a/pkg/watermark/publish/publisher.go b/pkg/watermark/publish/publisher.go
index 8dd4fed7..8271b3df 100644
--- a/pkg/watermark/publish/publisher.go
+++ b/pkg/watermark/publish/publisher.go
@@ -38,7 +38,7 @@ type publish struct {
 }
 
 // NewPublish returns `Publish`.
-func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier, hbStore store.WatermarkKVStorer, otStore store.WatermarkKVStorer, inputOpts ...PublishOption) Publisher {
+func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier, watermarkStores store.WatermarkStorer, inputOpts ...PublishOption) Publisher {
 
 	log := logging.FromContext(ctx)
 
@@ -55,15 +55,14 @@ func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier
 	p := &publish{
 		ctx:            ctx,
 		entity:         processorEntity,
-		heartbeatStore: hbStore,
-		otStore:        otStore,
+		heartbeatStore: watermarkStores.HeartbeatStore(),
+		otStore:        watermarkStores.OffsetTimelineStore(),
 		log:            log,
 		opts:           opts,
 	}
 
 	p.initialSetup()
 
-	// TODO: i do not think we need this autoRefreshHeartbeat flag anymore. Remove it?
 	if opts.autoRefreshHeartbeat {
 		go p.publishHeartbeat()
 	}
@@ -96,7 +95,7 @@ func (p *publish) PublishWatermark(wm processor.Watermark, offset isb.Offset) {
 	// build value (offset)
 	value := make([]byte, 8)
 	var seq int64
-	if p.opts.isSource { // For source publisher, we dont' care about the offset, also the sequence of the offset might not be integer.
+	if p.opts.isSource || p.opts.isSink { // For source and sink publisher, we dont' care about the offset, also the sequence of the offset might not be integer.
 		seq = time.Now().UnixNano()
 	} else {
 		seq, _ = offset.Sequence()
diff --git a/pkg/watermark/store/interface.go b/pkg/watermark/store/interface.go
index 5fc0ed25..5c55ecb8 100644
--- a/pkg/watermark/store/interface.go
+++ b/pkg/watermark/store/interface.go
@@ -4,10 +4,10 @@ import (
 	"context"
 )
 
-// WMStorer is the watermark store implementation.
-type WMStorer interface {
-	WatermarkKVStorer
-	WatermarkKVWatcher
+// WatermarkStorer defines a pair of heartbeat KV store and offset timeline KV store
+type WatermarkStorer interface {
+	HeartbeatStore() WatermarkKVStorer
+	OffsetTimelineStore() WatermarkKVStorer
 }
 
 // WatermarkKVStorer defines the storage for publishing the watermark.
@@ -67,3 +67,9 @@ type WatermarkKVWatcher interface {
 	GetKVName() string
 	Close()
 }
+
+// WatermarkStoreWatcher defines a pair of heartbeat KV watcher and offset timeline KV watcher
+type WatermarkStoreWatcher interface {
+	HeartbeatWatcher() WatermarkKVWatcher
+	OffsetTimelineWatcher() WatermarkKVWatcher
+}
diff --git a/pkg/watermark/store/watermark_store.go b/pkg/watermark/store/watermark_store.go
new file mode 100644
index 00000000..7e168c8f
--- /dev/null
+++ b/pkg/watermark/store/watermark_store.go
@@ -0,0 +1,26 @@
+package store
+
+// watermarkStore wraps a pair of heartbeatStore and offsetTimelineStore,
+// it implements interface WatermarkStorers.
+type watermarkStore struct {
+	heartbeatStore      WatermarkKVStorer
+	offsetTimelineStore WatermarkKVStorer
+}
+
+var _ WatermarkStorer = (*watermarkStore)(nil)
+
+// BuildWatermarkStore returns a WatermarkStorer instance
+func BuildWatermarkStore(hbStore, otStore WatermarkKVStorer) WatermarkStorer {
+	return &watermarkStore{
+		heartbeatStore:      hbStore,
+		offsetTimelineStore: otStore,
+	}
+}
+
+func (ws *watermarkStore) HeartbeatStore() WatermarkKVStorer {
+	return ws.heartbeatStore
+}
+
+func (ws *watermarkStore) OffsetTimelineStore() WatermarkKVStorer {
+	return ws.offsetTimelineStore
+}
diff --git a/pkg/watermark/store/watermark_store_watcher.go b/pkg/watermark/store/watermark_store_watcher.go
new file mode 100644
index 00000000..7cadb97e
--- /dev/null
+++ b/pkg/watermark/store/watermark_store_watcher.go
@@ -0,0 +1,25 @@
+package store
+
+// watermarkStoreWatcher defines a pair of heartbeatStoreWatcher and offsetTimelineStoreWatcher,
+// it implements interface WatermarkStoreWatcher
+type watermarkStoreWatcher struct {
+	heartbeatStoreWatcher      WatermarkKVWatcher
+	offsetTimelineStoreWatcher WatermarkKVWatcher
+}
+
+var _ WatermarkStoreWatcher = (*watermarkStoreWatcher)(nil)
+
+func (w *watermarkStoreWatcher) HeartbeatWatcher() WatermarkKVWatcher {
+	return w.heartbeatStoreWatcher
+}
+func (w *watermarkStoreWatcher) OffsetTimelineWatcher() WatermarkKVWatcher {
+	return w.offsetTimelineStoreWatcher
+}
+
+// BuildWatermarkStoreWatcher returns a WatermarkStoreWatcher instance
+func BuildWatermarkStoreWatcher(hbStoreWatcher, otStoreWatcher WatermarkKVWatcher) WatermarkStoreWatcher {
+	return &watermarkStoreWatcher{
+		heartbeatStoreWatcher:      hbStoreWatcher,
+		offsetTimelineStoreWatcher: otStoreWatcher,
+	}
+}


Test patch:
diff --git a/pkg/sources/generator/tickgen_test.go b/pkg/sources/generator/tickgen_test.go
index 85cfadb5..89b88cca 100644
--- a/pkg/sources/generator/tickgen_test.go
+++ b/pkg/sources/generator/tickgen_test.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
@@ -28,7 +29,7 @@ func TestRead(t *testing.T) {
 		Replica:  0,
 	}
 
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStore := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	toSteps := map[string]isb.BufferWriter{
 		"writer": dest,
 	}
@@ -61,7 +62,7 @@ func TestStop(t *testing.T) {
 		Hostname: "TestRead",
 		Replica:  0,
 	}
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStore := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	toSteps := map[string]isb.BufferWriter{
 		"writer": dest,
 	}
diff --git a/pkg/sources/generator/watermark_test.go b/pkg/sources/generator/watermark_test.go
index c8395745..37a8371b 100644
--- a/pkg/sources/generator/watermark_test.go
+++ b/pkg/sources/generator/watermark_test.go
@@ -7,7 +7,7 @@ import (
 	"time"
 
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
-	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
@@ -34,7 +34,7 @@ func TestWatermark(t *testing.T) {
 		Hostname: "TestRead",
 		Replica:  0,
 	}
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStore := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	mgen, err := NewMemGen(m, 1, 8, time.Millisecond, []isb.BufferWriter{dest}, nil, nil, publishWMStore)
 	assert.NoError(t, err)
 	stop := mgen.Start()
diff --git a/pkg/sources/http/http_test.go b/pkg/sources/http/http_test.go
index d11aa0ba..7a128bf1 100644
--- a/pkg/sources/http/http_test.go
+++ b/pkg/sources/http/http_test.go
@@ -8,6 +8,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 	"github.com/stretchr/testify/assert"
 )
@@ -47,9 +48,9 @@ func Test_NewHTTP(t *testing.T) {
 		Replica:  0,
 	}
 	dest := []isb.BufferWriter{simplebuffer.NewInMemoryBuffer("test", 100)}
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStores := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(map[string]isb.BufferWriter{})
-	h, err := New(vi, dest, fetchWatermark, publishWatermark, publishWMStore)
+	h, err := New(vi, dest, fetchWatermark, publishWatermark, publishWMStores)
 	assert.NoError(t, err)
 	assert.False(t, h.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 181acef8..93f6ce71 100644
--- a/pkg/sources/kafka/handler_test.go
+++ b/pkg/sources/kafka/handler_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 	"github.com/stretchr/testify/assert"
 )
@@ -44,7 +45,7 @@ func TestMessageHandling(t *testing.T) {
 		Hostname: "test-host",
 		Replica:  0,
 	}
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStore := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(map[string]isb.BufferWriter{})
 	ks, _ := NewKafkaSource(vi, dest, fetchWatermark, publishWatermark, publishWMStore, WithLogger(logging.NewLogger()),
 		WithBufferSize(100), WithReadTimeOut(100*time.Millisecond))
diff --git a/pkg/sources/kafka/reader_test.go b/pkg/sources/kafka/reader_test.go
index 824ad33f..49fc9d65 100644
--- a/pkg/sources/kafka/reader_test.go
+++ b/pkg/sources/kafka/reader_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 	"github.com/stretchr/testify/assert"
 )
@@ -32,7 +33,7 @@ func TestNewKafkasource(t *testing.T) {
 		Hostname: "test-host",
 		Replica:  0,
 	}
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStore := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(map[string]isb.BufferWriter{})
 	ks, err := NewKafkaSource(vi, dest, fetchWatermark, publishWatermark, publishWMStore, WithLogger(logging.NewLogger()), WithBufferSize(100), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
 
@@ -68,7 +69,7 @@ func TestGroupNameOverride(t *testing.T) {
 		Hostname: "test-host",
 		Replica:  0,
 	}
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStore := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(map[string]isb.BufferWriter{})
 	ks, _ := NewKafkaSource(vi, dest, fetchWatermark, publishWatermark, publishWMStore, WithLogger(logging.NewLogger()), WithBufferSize(100), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
 
@@ -94,7 +95,7 @@ func TestDefaultBufferSize(t *testing.T) {
 		Hostname: "test-host",
 		Replica:  0,
 	}
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStore := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(map[string]isb.BufferWriter{})
 	ks, _ := NewKafkaSource(vi, dest, fetchWatermark, publishWatermark, publishWMStore, WithLogger(logging.NewLogger()), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
 
@@ -120,7 +121,7 @@ func TestBufferSizeOverrides(t *testing.T) {
 		Hostname: "test-host",
 		Replica:  0,
 	}
-	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	publishWMStore := store.BuildWatermarkStore(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(map[string]isb.BufferWriter{})
 	ks, _ := NewKafkaSource(vi, dest, fetchWatermark, publishWatermark, publishWMStore, WithLogger(logging.NewLogger()), WithBufferSize(110), WithReadTimeOut(100*time.Millisecond), WithGroupName("default"))
 
diff --git a/pkg/watermark/fetch/edge_fetcher_test.go b/pkg/watermark/fetch/edge_fetcher_test.go
index 4917dc19..4cea1155 100644
--- a/pkg/watermark/fetch/edge_fetcher_test.go
+++ b/pkg/watermark/fetch/edge_fetcher_test.go
@@ -15,6 +15,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb"
 	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
 )
 
@@ -50,7 +51,7 @@ func TestBuffer_GetWatermark(t *testing.T) {
 
 	hbWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", publisherHBBucketName, defaultJetStreamClient)
 	otWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", publisherOTBucketName, defaultJetStreamClient)
-	processorManager := NewProcessorManager(ctx, hbWatcher, otWatcher)
+	processorManager := NewProcessorManager(ctx, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher))
 	var (
 		// TODO: watcher should not be nil
 		testPod0     = NewProcessorToFetch(ctx, processor.NewProcessorEntity("testPod1"), 5, otWatcher)
diff --git a/pkg/watermark/fetch/processor_manager_test.go b/pkg/watermark/fetch/processor_manager_test.go
index ad491e6c..e4e94e60 100644
--- a/pkg/watermark/fetch/processor_manager_test.go
+++ b/pkg/watermark/fetch/processor_manager_test.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
 )
 
@@ -73,7 +74,7 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 
 	hbWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", keyspace+"_PROCESSORS", defaultJetStreamClient)
 	otWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", keyspace+"_OT", defaultJetStreamClient)
-	var testVertex = NewProcessorManager(ctx, hbWatcher, otWatcher, WithPodHeartbeatRate(1), WithRefreshingProcessorsRate(1), WithSeparateOTBuckets(false))
+	var testVertex = NewProcessorManager(ctx, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher), WithPodHeartbeatRate(1), WithRefreshingProcessorsRate(1), WithSeparateOTBuckets(false))
 	var testBuffer = NewEdgeFetcher(ctx, "testBuffer", testVertex).(*edgeFetcher)
 
 	go func() {
@@ -308,7 +309,7 @@ func TestFetcherWithSeparateOTBucket(t *testing.T) {
 
 	hbWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", keyspace+"_PROCESSORS", defaultJetStreamClient)
 	otWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", keyspace+"_OT", defaultJetStreamClient)
-	var testVertex = NewProcessorManager(ctx, hbWatcher, otWatcher, WithPodHeartbeatRate(1), WithRefreshingProcessorsRate(1), WithSeparateOTBuckets(true))
+	var testVertex = NewProcessorManager(ctx, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher), WithPodHeartbeatRate(1), WithRefreshingProcessorsRate(1), WithSeparateOTBuckets(true))
 	var testBuffer = NewEdgeFetcher(ctx, "testBuffer", testVertex).(*edgeFetcher)
 
 	//var location, _ = time.LoadLocation("UTC")
diff --git a/pkg/watermark/publish/publisher_test.go b/pkg/watermark/publish/publisher_test.go
index ab848d6c..18c2b8fa 100644
--- a/pkg/watermark/publish/publisher_test.go
+++ b/pkg/watermark/publish/publisher_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb"
 	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
 )
 
@@ -56,7 +57,7 @@ func TestPublisherWithSeparateOTBuckets(t *testing.T) {
 
 	publishEntity := processor.NewProcessorEntity("publisherTestPod1")
 
-	p := NewPublish(ctx, publishEntity, heartbeatKV, otKV, WithAutoRefreshHeartbeatDisabled(), WithPodHeartbeatRate(1)).(*publish)
+	p := NewPublish(ctx, publishEntity, store.BuildWatermarkStore(heartbeatKV, otKV), WithAutoRefreshHeartbeatDisabled(), WithPodHeartbeatRate(1)).(*publish)
 
 	var epoch int64 = 1651161600
 	var location, _ = time.LoadLocation("UTC")
@@ -108,7 +109,7 @@ func TestPublisherWithSharedOTBucket(t *testing.T) {
 	otKV, err := jetstream.NewKVJetStreamKVStore(ctx, "testPublisher", keyspace+"_OT", defaultJetStreamClient)
 	assert.NoError(t, err)
 
-	p := NewPublish(ctx, publishEntity, heartbeatKV, otKV, WithAutoRefreshHeartbeatDisabled(), WithPodHeartbeatRate(1)).(*publish)
+	p := NewPublish(ctx, publishEntity, store.BuildWatermarkStore(heartbeatKV, otKV), WithAutoRefreshHeartbeatDisabled(), WithPodHeartbeatRate(1)).(*publish)
 
 	var epoch int64 = 1651161600
 	var location, _ = time.LoadLocation("UTC")