Repository: numaproj/numaflow
Pull request #93 (2022-07-06 21:34:51)

chore: refactor source watermark
refactor source watermark


Patch:
diff --git a/pkg/sources/generator/tickgen.go b/pkg/sources/generator/tickgen.go
index a2eeab84..a0881dbc 100644
--- a/pkg/sources/generator/tickgen.go
+++ b/pkg/sources/generator/tickgen.go
@@ -6,7 +6,6 @@ import (
 	"context"
 	"encoding/binary"
 	"encoding/json"
-	"os"
 	"strconv"
 	"sync/atomic"
 	"time"
@@ -16,6 +15,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb/forward"
 	"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"
@@ -78,7 +78,11 @@ type memgen struct {
 	// read timeout for the reader
 	readTimeout time.Duration
 
-	// watermark progressor
+	// vertex instance
+	vertexInstance *dfv1.VertexInstance
+	// source watermark publisher
+	sourcePublishWM publish.Publisher
+	// TODO: delete this watermark progressor
 	progressor *watermark
 
 	logger *zap.SugaredLogger
@@ -114,14 +118,21 @@ func WithReadTimeOut(timeout time.Duration) Option {
 // msgSize - size of each generated message
 // timeunit - unit of time per tick. could be any golang time.Duration.
 // writers - destinations to write to
-func NewMemGen(vertexInstance *dfv1.VertexInstance, rpu int, msgSize int32, timeunit time.Duration, writers []isb.BufferWriter, opts ...Option) (*memgen, error) {
+func NewMemGen(vertexInstance *dfv1.VertexInstance,
+	rpu int,
+	msgSize int32,
+	timeunit time.Duration,
+	writers []isb.BufferWriter,
+	fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores *generic.PublishWMStores, // watermarks
+	opts ...Option) (*memgen, error) {
 	gensrc := &memgen{
-		rpu:          rpu,
-		msgSize:      msgSize,
-		timeunit:     timeunit,
-		name:         vertexInstance.Vertex.Spec.Name,
-		pipelineName: vertexInstance.Vertex.Spec.PipelineName,
-		genfn:        recordGenerator,
+		rpu:            rpu,
+		msgSize:        msgSize,
+		timeunit:       timeunit,
+		name:           vertexInstance.Vertex.Spec.Name,
+		pipelineName:   vertexInstance.Vertex.Spec.PipelineName,
+		genfn:          recordGenerator,
+		vertexInstance: vertexInstance,
 		progressor: &watermark{
 			sourcePublish: nil,
 			wmProgressor:  nil,
@@ -157,20 +168,11 @@ func NewMemGen(vertexInstance *dfv1.VertexInstance, rpu int, msgSize int32, time
 		}
 	}
 
-	var wmProgressor generic.Progressor = nil
-	var err error
-	// TODO: pass it as option
-	if val, ok := os.LookupEnv(dfv1.EnvWatermarkOn); ok && val == "true" {
-		err = gensrc.buildWMProgressor(vertexInstance)
-		if err != nil {
-			return nil, err
-		}
-		wmProgressor = gensrc.progressor.wmProgressor
-	}
+	// attach a source publisher so the source can assign the watermarks.
+	gensrc.sourcePublishWM = gensrc.buildSourceWatermarkPublisher(publishWMStores)
 
-	_ = wmProgressor
 	// we pass in the context to forwarder as well so that it can shut down when we cancel the context
-	forwarder, err := forward.NewInterStepDataForward(vertexInstance.Vertex, gensrc, destinations, forward.All, applier.Terminal, nil, nil, forwardOpts...)
+	forwarder, err := forward.NewInterStepDataForward(vertexInstance.Vertex, gensrc, destinations, forward.All, applier.Terminal, fetchWM, publishWM, forwardOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -209,12 +211,11 @@ func (mg *memgen) Read(ctx context.Context, count int64) ([]*isb.ReadMessage, er
 			// publish the last message's offset with watermark, this is an optimization to avoid too many insert calls
 			// into the offset timeline store.
 			// Please note that we are inserting the watermark before the data has been persisted into ISB by the forwarder.
-			if mg.progressor.sourcePublish != nil {
-				o := msgs[len(msgs)-1].ReadOffset
-				nanos, _ := o.Sequence()
-				// remove the nanosecond precision
-				mg.progressor.sourcePublish.PublishWatermark(processor.Watermark(time.Unix(0, nanos)), o)
-			}
+			o := msgs[len(msgs)-1].ReadOffset
+			nanos, _ := o.Sequence()
+			// remove the nanosecond precision
+			mg.sourcePublishWM.PublishWatermark(processor.Watermark(time.Unix(0, nanos)), o)
+
 		case <-timeout:
 			mg.logger.Debugw("Timed out waiting for messages to read.", zap.Duration("waited", mg.readTimeout))
 			return msgs, nil
diff --git a/pkg/sources/generator/watermark.go b/pkg/sources/generator/watermark.go
index 0c5f5ed8..cc21aa11 100644
--- a/pkg/sources/generator/watermark.go
+++ b/pkg/sources/generator/watermark.go
@@ -3,68 +3,14 @@ package generator
 import (
 	"fmt"
 
-	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
-	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 )
 
-// buildWMProgressor builds the watermark progressor. It will create a sourcePublisher so it can publish the watermark
-// it reads from the source. The sourcePublisher is passed in as the input for generic progressor.
-func (mg *memgen) buildWMProgressor(vertexInstance *dfv1.VertexInstance) error {
-	ctx := mg.lifecycleCtx
-
-	js, err := generic.GetJetStreamConnection(mg.lifecycleCtx)
-	if err != nil {
-		return err
-	}
-
-	// publish source watermark and this is very much dependent on the source
-	sourcePublishKeySpace := fmt.Sprintf("source-%s", generic.GetPublishKeySpace(vertexInstance.Vertex))
-	// TODO: remove this once bucket creation has been moved to controller
-	err = generic.CreateProcessorBucketIfMissing(fmt.Sprintf("%s_PROCESSORS", sourcePublishKeySpace), js)
-	if err != nil {
-		return err
-	}
-	publishEntity := processor.NewProcessorEntity(fmt.Sprintf("source-%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica), sourcePublishKeySpace, processor.WithSeparateOTBuckets(false))
-
-	if err != nil {
-		return err
-	}
-	// use this while reading the data from the source.
-	mg.progressor.sourcePublish = publish.NewPublish(mg.lifecycleCtx, publishEntity, nil, nil)
-
-	// fall back on the generic progressor and use the source publisher as the input to the generic progressor.
-	// use the source Publisher as the source
-
-	// TODO: remove this once bucket creation has been moved to controller
-	err = generic.CreateProcessorBucketIfMissing(fmt.Sprintf("%s_PROCESSORS", generic.GetPublishKeySpace(vertexInstance.Vertex)), js)
-	if err != nil {
-		return err
-	}
-	var fetchWM = generic.BuildFetchWM(nil, nil)
-	var publishWM = generic.BuildPublishWM(nil, nil)
-	var wmProgressor = generic.NewGenericProgress(ctx, fmt.Sprintf("%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica), sourcePublishKeySpace, generic.GetPublishKeySpace(vertexInstance.Vertex), publishWM, fetchWM)
-	mg.progressor.wmProgressor = wmProgressor
-
-	mg.logger.Info("Initialized watermark progressor")
-
-	return nil
-}
-
-func (w *watermark) PublishWatermark(watermark processor.Watermark, offset isb.Offset) {
-	w.wmProgressor.PublishWatermark(watermark, offset)
-}
-
-func (w *watermark) GetLatestWatermark() processor.Watermark {
-	return w.wmProgressor.GetLatestWatermark()
-}
-
-func (w *watermark) StopPublisher() {
-	w.wmProgressor.StopPublisher()
-}
-
-func (w *watermark) GetWatermark(offset isb.Offset) processor.Watermark {
-	return w.wmProgressor.GetWatermark(offset)
+func (mg *memgen) buildSourceWatermarkPublisher(publishWMStores *generic.PublishWMStores) 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, mg.vertexInstance.Vertex.Name+"_SOURCE")
+	return publish.NewPublish(mg.lifecycleCtx, processorEntity, publishWMStores.HBStore, publishWMStores.OTStore)
 }
diff --git a/pkg/sources/source.go b/pkg/sources/source.go
index 5348ac6f..5f71d44e 100644
--- a/pkg/sources/source.go
+++ b/pkg/sources/source.go
@@ -5,6 +5,10 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/numaproj/numaflow/pkg/watermark/fetch"
+	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/publish"
+	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 	"go.uber.org/zap"
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
@@ -24,14 +28,24 @@ type SourceProcessor struct {
 	VertexInstance *dfv1.VertexInstance
 }
 
-func (u *SourceProcessor) Start(ctx context.Context) error {
+func (sp *SourceProcessor) Start(ctx context.Context) error {
 	log := logging.FromContext(ctx)
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 	var writers []isb.BufferWriter
-	switch u.ISBSvcType {
+
+	// watermark variables no-op initialization
+	var fetchWatermark fetch.Fetcher = generic.NewNoOpWMProgressor()
+	// publishWatermark is a map representing a progressor per edge, we are initializing them to a no-op progressor
+	publishWatermark := make(map[string]publish.Publisher)
+	for _, buffer := range sp.VertexInstance.Vertex.GetToBuffers() {
+		publishWatermark[buffer.Name] = generic.NewNoOpWMProgressor()
+	}
+	var publishWMStore = generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+
+	switch sp.ISBSvcType {
 	case dfv1.ISBSvcTypeRedis:
-		for _, e := range u.VertexInstance.Vertex.Spec.ToEdges {
+		for _, e := range sp.VertexInstance.Vertex.Spec.ToEdges {
 			writeOpts := []redisisb.Option{}
 			if x := e.Limits; x != nil && x.BufferMaxLength != nil {
 				writeOpts = append(writeOpts, redisisb.WithMaxLength(int64(*x.BufferMaxLength)))
@@ -39,14 +53,18 @@ func (u *SourceProcessor) Start(ctx context.Context) error {
 			if x := e.Limits; x != nil && x.BufferUsageLimit != nil {
 				writeOpts = append(writeOpts, redisisb.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 			}
-			buffer := dfv1.GenerateEdgeBufferName(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, e.From, e.To)
+			buffer := dfv1.GenerateEdgeBufferName(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, e.From, e.To)
 			group := buffer + "-group"
 			redisClient := clients.NewInClusterRedisClient()
 			writer := redisisb.NewBufferWrite(ctx, redisClient, buffer, group, writeOpts...)
 			writers = append(writers, writer)
 		}
 	case dfv1.ISBSvcTypeJetStream:
-		for _, e := range u.VertexInstance.Vertex.Spec.ToEdges {
+		// build the right watermark progessor if watermark is enabled
+		// build watermark progressors
+		fetchWatermark, publishWatermark, publishWMStore = generic.BuildJetStreamWatermarkProgressorsForSource(ctx, sp.VertexInstance)
+
+		for _, e := range sp.VertexInstance.Vertex.Spec.ToEdges {
 			writeOpts := []jetstreamisb.WriteOption{
 				jetstreamisb.WithUsingWriteInfoAsRate(true),
 			}
@@ -56,8 +74,8 @@ func (u *SourceProcessor) Start(ctx context.Context) error {
 			if x := e.Limits; x != nil && x.BufferUsageLimit != nil {
 				writeOpts = append(writeOpts, jetstreamisb.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 			}
-			buffer := dfv1.GenerateEdgeBufferName(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, e.From, e.To)
-			streamName := fmt.Sprintf("%s-%s", u.VertexInstance.Vertex.Spec.PipelineName, buffer)
+			buffer := dfv1.GenerateEdgeBufferName(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, e.From, e.To)
+			streamName := fmt.Sprintf("%s-%s", sp.VertexInstance.Vertex.Spec.PipelineName, buffer)
 			jetStreamClient := clients.NewInClusterJetStreamClient()
 			writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jetStreamClient, buffer, streamName, streamName, writeOpts...)
 			if err != nil {
@@ -66,14 +84,14 @@ func (u *SourceProcessor) Start(ctx context.Context) error {
 			writers = append(writers, writer)
 		}
 	default:
-		return fmt.Errorf("unrecognized isb svc type %q", u.ISBSvcType)
+		return fmt.Errorf("unrecognized isb svc type %q", sp.ISBSvcType)
 	}
 
-	sourcer, err := u.getSourcer(writers, log)
+	sourcer, err := sp.getSourcer(writers, fetchWatermark, publishWatermark, &publishWMStore, log)
 	if err != nil {
 		return fmt.Errorf("failed to find a sourcer, error: %w", err)
 	}
-	log.Infow("Start processing source messages", zap.String("isbs", string(u.ISBSvcType)), zap.Any("to", u.VertexInstance.Vertex.GetToBuffers()))
+	log.Infow("Start processing source messages", zap.String("isbs", string(sp.ISBSvcType)), zap.Any("to", sp.VertexInstance.Vertex.GetToBuffers()))
 	stopped := sourcer.Start()
 	wg := &sync.WaitGroup{}
 	wg.Add(1)
@@ -87,7 +105,7 @@ func (u *SourceProcessor) Start(ctx context.Context) error {
 	}()
 
 	metricsOpts := []metrics.Option{}
-	if s := u.VertexInstance.Vertex.Spec.Scale.LookbackSeconds; s != nil {
+	if s := sp.VertexInstance.Vertex.Spec.Scale.LookbackSeconds; s != nil {
 		metricsOpts = append(metricsOpts, metrics.WithLookbackSeconds(int64(*s)))
 	}
 	if x, ok := sourcer.(isb.LagReader); ok {
@@ -96,7 +114,7 @@ func (u *SourceProcessor) Start(ctx context.Context) error {
 	if x, ok := writers[0].(isb.Ratable); ok { // Only need to use the rate of one of the writer
 		metricsOpts = append(metricsOpts, metrics.WithRater(x))
 	}
-	ms := metrics.NewMetricsServer(u.VertexInstance.Vertex, metricsOpts...)
+	ms := metrics.NewMetricsServer(sp.VertexInstance.Vertex, metricsOpts...)
 	if shutdown, err := ms.Start(ctx); err != nil {
 		return fmt.Errorf("failed to start metrics server, error: %w", err)
 	} else {
@@ -112,14 +130,14 @@ func (u *SourceProcessor) Start(ctx context.Context) error {
 }
 
 // getSourcer is used to send the sourcer information
-func (u *SourceProcessor) getSourcer(writers []isb.BufferWriter, logger *zap.SugaredLogger) (Sourcer, error) {
-	src := u.VertexInstance.Vertex.Spec.Source
+func (sp *SourceProcessor) getSourcer(writers []isb.BufferWriter, fetchWM fetch.Fetcher, publishWM map[string]publish.Publisher, publishWMStores *generic.PublishWMStores, logger *zap.SugaredLogger) (Sourcer, error) {
+	src := sp.VertexInstance.Vertex.Spec.Source
 	if x := src.Generator; x != nil {
-		return generator.NewMemGen(u.VertexInstance, int(*x.RPU), *x.MsgSize, x.Duration.Duration, writers, generator.WithLogger(logger))
+		return generator.NewMemGen(sp.VertexInstance, int(*x.RPU), *x.MsgSize, x.Duration.Duration, writers, fetchWM, publishWM, publishWMStores, generator.WithLogger(logger))
 	} else if x := src.Kafka; x != nil {
-		return kafka.NewKafkaSource(u.VertexInstance.Vertex, writers, kafka.WithGroupName(x.ConsumerGroupName), kafka.WithLogger(logger))
+		return kafka.NewKafkaSource(sp.VertexInstance.Vertex, writers, kafka.WithGroupName(x.ConsumerGroupName), kafka.WithLogger(logger))
 	} else if x := src.HTTP; x != nil {
-		return http.New(u.VertexInstance.Vertex, writers, http.WithLogger(logger))
+		return http.New(sp.VertexInstance.Vertex, writers, http.WithLogger(logger))
 	}
 	return nil, fmt.Errorf("invalid source spec")
 }
diff --git a/pkg/watermark/generic/fetcher.go b/pkg/watermark/generic/fetcher.go
index e8f3390f..f49748dc 100644
--- a/pkg/watermark/generic/fetcher.go
+++ b/pkg/watermark/generic/fetcher.go
@@ -16,24 +16,25 @@ type GenericFetch struct {
 
 var _ fetch.Fetcher = (*GenericFetch)(nil)
 
-type FetchWM struct {
-	hbWatch store.WatermarkKVWatcher
-	otWatch store.WatermarkKVWatcher
+// FetchWMWatchers has the watcher information required for fetching watermarks.
+type FetchWMWatchers struct {
+	HBWatch store.WatermarkKVWatcher
+	OTWatch store.WatermarkKVWatcher
 }
 
-// BuildFetchWM builds the FetchWM
-func BuildFetchWM(hbWatch store.WatermarkKVWatcher, otWatch store.WatermarkKVWatcher) FetchWM {
-	return FetchWM{
-		hbWatch: hbWatch,
-		otWatch: otWatch,
+// BuildFetchWMWatchers builds the FetchWMWatchers
+func BuildFetchWMWatchers(hbWatch store.WatermarkKVWatcher, otWatch store.WatermarkKVWatcher) FetchWMWatchers {
+	return FetchWMWatchers{
+		HBWatch: hbWatch,
+		OTWatch: otWatch,
 	}
 }
 
 // NewGenericFetch returns GenericFetch. vertexName is the vertex currently processing.
 // fetchKeyspace is obsolete, and will be removed in subsequent iterations. fetchWM is a struct for retrieving both the heartbeat
 // and the offset watermark timeline (Vn-1 vertex).
-func NewGenericFetch(ctx context.Context, vertexName string, fetchKeyspace string, fetchWM FetchWM) *GenericFetch {
-	fromVertex := fetch.NewFromVertex(ctx, fetchKeyspace, fetchWM.hbWatch, fetchWM.otWatch)
+func NewGenericFetch(ctx context.Context, vertexName string, fetchKeyspace string, fetchWM FetchWMWatchers) *GenericFetch {
+	fromVertex := fetch.NewFromVertex(ctx, fetchKeyspace, fetchWM.HBWatch, fetchWM.OTWatch)
 	fromEdge := fetch.NewEdgeBuffer(ctx, vertexName, fromVertex)
 
 	gf := &GenericFetch{
diff --git a/pkg/watermark/generic/generic.go b/pkg/watermark/generic/generic.go
index e9be4b2e..cc91579c 100644
--- a/pkg/watermark/generic/generic.go
+++ b/pkg/watermark/generic/generic.go
@@ -14,6 +14,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
+	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 	"go.uber.org/zap"
 )
 
@@ -41,7 +42,7 @@ func WithSeparateOTBuckets(separate bool) GenericProgressOption {
 }
 
 // NewGenericProgress will move the watermark for all the vertices once consumed fromEdge the source.
-func NewGenericProgress(ctx context.Context, processorName string, fetchKeyspace string, publishKeyspace string, publishWM PublishWM, fetchWM FetchWM, inputOpts ...GenericProgressOption) *GenericProgress {
+func NewGenericProgress(ctx context.Context, processorName string, fetchKeyspace string, publishKeyspace string, publishWM PublishWMStores, fetchWM FetchWMWatchers, inputOpts ...GenericProgressOption) *GenericProgress {
 	opts := &genericProgressOptions{
 		separateOTBucket: false,
 	}
@@ -55,9 +56,9 @@ func NewGenericProgress(ctx context.Context, processorName string, fetchKeyspace
 	// to generic watermark for a UDF, it has to start the Fetcher and the Publisher
 
 	publishEntity := processor.NewProcessorEntity(processorName, publishKeyspace)
-	udfPublish := publish.NewPublish(ctx, publishEntity, publishWM.hbStore, publishWM.otStore)
+	udfPublish := publish.NewPublish(ctx, publishEntity, publishWM.HBStore, publishWM.OTStore)
 
-	udfFromVertex := fetch.NewFromVertex(ctx, fetchKeyspace, fetchWM.hbWatch, fetchWM.otWatch)
+	udfFromVertex := fetch.NewFromVertex(ctx, fetchKeyspace, fetchWM.HBWatch, fetchWM.OTWatch)
 	udfFetch := fetch.NewEdgeBuffer(ctx, processorName, udfFromVertex)
 
 	u := &GenericProgress{
@@ -99,8 +100,7 @@ func BuildJetStreamWatermarkProgressors(ctx context.Context, vertexInstance *v1a
 		fetchWatermark = NewNoOpWMProgressor()
 		publishWatermark = make(map[string]publish.Publisher)
 		for _, buffer := range vertexInstance.Vertex.GetToBuffers() {
-			streamName := isbsvc.JetStreamName(vertexInstance.Vertex.Spec.PipelineName, buffer.Name)
-			publishWatermark[streamName] = NewNoOpWMProgressor()
+			publishWatermark[buffer.Name] = NewNoOpWMProgressor()
 		}
 
 		return fetchWatermark, publishWatermark
@@ -123,27 +123,63 @@ func BuildJetStreamWatermarkProgressors(ctx context.Context, vertexInstance *v1a
 		log.Fatalw("JetStreamKVWatch failed", zap.String("OTBucket", otBucket), zap.Error(err))
 	}
 
-	var fetchWmWatchers = BuildFetchWM(hbWatch, otWatch)
+	var fetchWmWatchers = BuildFetchWMWatchers(hbWatch, otWatch)
 	fetchWatermark = NewGenericFetch(ctx, vertexInstance.Vertex.Name, GetFetchKeyspace(vertexInstance.Vertex), fetchWmWatchers)
 
 	// Publisher map creation
 	// We do not separate Heartbeat bucket for each edge, can be reused
-	hbStore, err := jetstream.NewKVJetStreamKVStore(ctx, vertexInstance.Vertex.Spec.PipelineName, hbBucket, clients.NewInClusterJetStreamClient())
+	hbStore, err := jetstream.NewKVJetStreamKVStore(ctx, pipelineName, hbBucket, clients.NewInClusterJetStreamClient())
 	if err != nil {
 		log.Fatalw("JetStreamKVStore failed", zap.String("HeartbeatBucket", hbBucket), zap.Error(err))
 	}
 
 	for _, buffer := range vertexInstance.Vertex.GetToBuffers() {
-		var processorName = fmt.Sprintf("%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica)
-		streamName := isbsvc.JetStreamName(pipelineName, buffer.Name)
-		otStoreBucket := isbsvc.JetStreamOTBucket(pipelineName, streamName)
-		otStore, err := jetstream.NewKVJetStreamKVStore(ctx, vertexInstance.Vertex.Spec.PipelineName, otStoreBucket, clients.NewInClusterJetStreamClient())
+		otStoreBucket := isbsvc.JetStreamOTBucket(pipelineName, buffer.Name)
+		otStore, err := jetstream.NewKVJetStreamKVStore(ctx, pipelineName, otStoreBucket, clients.NewInClusterJetStreamClient())
 		if err != nil {
 			log.Fatalw("JetStreamKVStore failed", zap.String("OTBucket", otStoreBucket), zap.Error(err))
 		}
-		var publishStores = BuildPublishWM(hbStore, otStore)
-		publishWatermark[streamName] = NewGenericPublish(ctx, processorName, GetPublishKeySpace(vertexInstance.Vertex), publishStores)
+
+		var publishStores = BuildPublishWMStores(hbStore, otStore)
+		var processorName = fmt.Sprintf("%s-%d", vertexInstance.Vertex.Name, vertexInstance.Replica)
+		publishWatermark[buffer.Name] = NewGenericPublish(ctx, processorName, GetPublishKeySpace(vertexInstance.Vertex), publishStores)
 	}
 
 	return fetchWatermark, publishWatermark
 }
+
+// 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) (fetchWatermark fetch.Fetcher, publishWatermark map[string]publish.Publisher, publishWM PublishWMStores) {
+	fetchWatermark, publishWatermark = BuildJetStreamWatermarkProgressors(ctx, vertexInstance)
+
+	// return no-ops if not enabled!
+	if !sharedutil.IsWatermarkEnabled() {
+		publishWM = BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+		return fetchWatermark, publishWatermark, publishWM
+	}
+
+	log := logging.FromContext(ctx)
+	pipelineName := vertexInstance.Vertex.Spec.PipelineName
+
+	sourceBufferName := vertexInstance.Vertex.GetFromBuffers()[0].Name
+	// hearbeat
+	hbBucket := isbsvc.JetStreamProcessorBucket(pipelineName, sourceBufferName)
+	hbKVStore, err := jetstream.NewKVJetStreamKVStore(ctx, pipelineName, hbBucket, clients.NewInClusterJetStreamClient())
+	if err != nil {
+		log.Fatalw("JetStreamKVStore failed", zap.String("HeartbeatBucket", hbBucket), zap.Error(err))
+	}
+
+	// OT
+	otStoreBucket := isbsvc.JetStreamOTBucket(pipelineName, sourceBufferName)
+	otKVStore, err := jetstream.NewKVJetStreamKVStore(ctx, pipelineName, otStoreBucket, clients.NewInClusterJetStreamClient())
+	if err != nil {
+		log.Fatalw("JetStreamKVStore failed", zap.String("OTBucket", otStoreBucket), zap.Error(err))
+	}
+
+	// interface for publisher store (HB and OT)
+	publishWM = BuildPublishWMStores(hbKVStore, otKVStore)
+
+	return fetchWatermark, publishWatermark, publishWM
+}
diff --git a/pkg/watermark/generic/publisher.go b/pkg/watermark/generic/publisher.go
index 4d3b9484..25068837 100644
--- a/pkg/watermark/generic/publisher.go
+++ b/pkg/watermark/generic/publisher.go
@@ -9,17 +9,17 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
-// PublishWM stores the store information for publishing the watermark.
-type PublishWM struct {
-	hbStore store.WatermarkKVStorer
-	otStore store.WatermarkKVStorer
+// PublishWMStores stores the store information for publishing the watermark.
+type PublishWMStores struct {
+	HBStore store.WatermarkKVStorer
+	OTStore store.WatermarkKVStorer
 }
 
-// BuildPublishWM builds the PublishWM
-func BuildPublishWM(hbStore store.WatermarkKVStorer, otStore store.WatermarkKVStorer) PublishWM {
-	return PublishWM{
-		hbStore: hbStore,
-		otStore: otStore,
+// BuildPublishWMStores builds the PublishWMStores.
+func BuildPublishWMStores(hbStore store.WatermarkKVStorer, otStore store.WatermarkKVStorer) PublishWMStores {
+	return PublishWMStores{
+		HBStore: hbStore,
+		OTStore: otStore,
 	}
 }
 
@@ -33,9 +33,9 @@ var _ publish.Publisher = (*GenericPublish)(nil)
 // NewGenericPublish returns GenericPublish. processorName is the unique processor (pod) that is running on this vertex.
 // publishKeyspace is obsolete, and will be removed in subsequent iterations. 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, publishKeyspace string, publishWM PublishWM) *GenericPublish {
+func NewGenericPublish(ctx context.Context, processorName string, publishKeyspace string, publishWM PublishWMStores) *GenericPublish {
 	publishEntity := processor.NewProcessorEntity(processorName, publishKeyspace)
-	udfPublish := publish.NewPublish(ctx, publishEntity, publishWM.hbStore, publishWM.otStore)
+	udfPublish := publish.NewPublish(ctx, publishEntity, publishWM.HBStore, publishWM.OTStore)
 	gp := &GenericPublish{
 		toEdge: udfPublish,
 	}
diff --git a/pkg/watermark/generic/utils.go b/pkg/watermark/generic/utils.go
index c9fb0097..1b705b5e 100644
--- a/pkg/watermark/generic/utils.go
+++ b/pkg/watermark/generic/utils.go
@@ -26,11 +26,6 @@ func GetJetStreamConnection(ctx context.Context) (nats.JetStreamContext, error)
 	return js, nil
 }
 
-// GetHeartbeatBucket returns the heartbeat bucket.
-func GetHeartbeatBucket(js nats.JetStreamContext, publishKeyspace string) (nats.KeyValue, error) {
-	return js.KeyValue(publishKeyspace + "_PROCESSORS")
-}
-
 // GetFetchKeyspace gets the fetch keyspace name fromEdge the vertex.
 func GetFetchKeyspace(v *dfv1.Vertex) string {
 	if len(v.Spec.FromEdges) > 0 {
@@ -44,7 +39,11 @@ func GetFetchKeyspace(v *dfv1.Vertex) string {
 
 // GetPublishKeySpace gets the publish keyspace name fromEdge the vertex
 func GetPublishKeySpace(v *dfv1.Vertex) string {
-	return fmt.Sprintf("%s-%s-%s", v.Namespace, v.Spec.PipelineName, v.Spec.Name)
+	if v.IsASource() {
+		return dfv1.GenerateSourceBufferName(v.Namespace, v.Spec.PipelineName, v.Spec.Name)
+	} else {
+		return fmt.Sprintf("%s-%s-%s", v.Namespace, v.Spec.PipelineName, v.Spec.Name)
+	}
 }
 
 // CreateProcessorBucketIfMissing creates the KV bucket if missing
diff --git a/pkg/watermark/publish/publisher.go b/pkg/watermark/publish/publisher.go
index 41cf7024..10fe0dff 100644
--- a/pkg/watermark/publish/publisher.go
+++ b/pkg/watermark/publish/publisher.go
@@ -66,6 +66,7 @@ func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier
 
 	p.initialSetup()
 
+	// TODO: i do not think we need this autoRefreshHeartbeat flag anymore. Remove it?
 	if p.autoRefreshHeartbeat {
 		go p.publishHeartbeat()
 	}
@@ -109,6 +110,8 @@ func (p *Publish) PublishWatermark(wm processor.Watermark, offset isb.Offset) {
 	}
 }
 
+// loadLatestFromStore loads the latest watermark stored in the watermark store.
+// TODO: how to repopulate if the processing unit is down for a really long time?
 func (p *Publish) loadLatestFromStore() processor.Watermark {
 	var watermarks = p.getAllOTKeysFromBucket()
 	var latestWatermark int64 = math.MinInt64
diff --git a/pkg/watermark/store/jetstream/kv_store.go b/pkg/watermark/store/jetstream/kv_store.go
index 57a5660e..b0dc0d9d 100644
--- a/pkg/watermark/store/jetstream/kv_store.go
+++ b/pkg/watermark/store/jetstream/kv_store.go
@@ -25,6 +25,7 @@ type KVJetStreamStore struct {
 
 var _ store.WatermarkKVStorer = (*KVJetStreamStore)(nil)
 
+// NewKVJetStreamKVStore returns KVJetStreamStore.
 func NewKVJetStreamKVStore(ctx context.Context, pipelineName string, bucketName string, client clients.JetStreamClient, opts ...JSKVStoreOption) (*KVJetStreamStore, error) {
 	var err error
 	conn, err := client.Connect(ctx)
diff --git a/pkg/watermark/store/noop/kv_store.go b/pkg/watermark/store/noop/kv_store.go
index ccdc5064..dbd4a1b5 100644
--- a/pkg/watermark/store/noop/kv_store.go
+++ b/pkg/watermark/store/noop/kv_store.go
@@ -8,12 +8,13 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
-// KVNoOpStore does a no-op store.
+// KVNoOpStore is a no-op store which does not do any operation but can be safely invoked.
 type KVNoOpStore struct {
 }
 
 var _ store.WatermarkKVStorer = (*KVNoOpStore)(nil)
 
+// NewKVNoOpStore returns KVNoOpStore.
 func NewKVNoOpStore() *KVNoOpStore {
 	return &KVNoOpStore{}
 }


Test patch:
diff --git a/pkg/sources/generator/tickgen_test.go b/pkg/sources/generator/tickgen_test.go
index 9d89a292..ac6f33d4 100644
--- a/pkg/sources/generator/tickgen_test.go
+++ b/pkg/sources/generator/tickgen_test.go
@@ -2,6 +2,8 @@ package generator
 
 import (
 	"context"
+	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 	"testing"
 	"time"
 
@@ -24,7 +26,8 @@ func TestRead(t *testing.T) {
 		Hostname: "TestRead",
 		Replica:  0,
 	}
-	mgen, err := NewMemGen(m, 5, 8, time.Millisecond, []isb.BufferWriter{dest})
+	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	mgen, err := NewMemGen(m, 5, 8, time.Millisecond, []isb.BufferWriter{dest}, nil, nil, &publishWMStore)
 	assert.NoError(t, err)
 	_ = mgen.Start()
 
@@ -52,7 +55,8 @@ func TestStop(t *testing.T) {
 		Hostname: "TestRead",
 		Replica:  0,
 	}
-	mgen, err := NewMemGen(m, 5, 8, time.Millisecond, []isb.BufferWriter{dest})
+	publishWMStore := generic.BuildPublishWMStores(noop.NewKVNoOpStore(), noop.NewKVNoOpStore())
+	mgen, err := NewMemGen(m, 5, 8, time.Millisecond, []isb.BufferWriter{dest}, nil, nil, &publishWMStore)
 	assert.NoError(t, err)
 	stop := mgen.Start()
 
diff --git a/pkg/sources/generator/watermark_test.go b/pkg/sources/generator/watermark_test.go
index dd3a5c62..762b53f8 100644
--- a/pkg/sources/generator/watermark_test.go
+++ b/pkg/sources/generator/watermark_test.go
@@ -2,6 +2,8 @@ package generator
 
 import (
 	"context"
+	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 	"os"
 	"testing"
 	"time"
@@ -31,7 +33,8 @@ func TestWatermark(t *testing.T) {
 		Hostname: "TestRead",
 		Replica:  0,
 	}
-	mgen, err := NewMemGen(m, 1, 8, time.Millisecond, []isb.BufferWriter{dest})
+	publishWMStore := generic.BuildPublishWMStores(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()