Repository: numaproj/numaflow
Pull request #336 (2022-11-10 07:57:27)

refactor: Close watermark fetcher and publisher correctly
Signed-off-by: Derek Wang <whynowy@gmail.com>

1. Close() function on watermark fetchers and publishers, handle them when application exits;
2. Shutdown reduce application with resource cleaning up.

<!--

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/daemon_server.go b/pkg/daemon/server/daemon_server.go
index 19c96cd4..1d1b9987 100644
--- a/pkg/daemon/server/daemon_server.go
+++ b/pkg/daemon/server/daemon_server.go
@@ -40,6 +40,7 @@ import (
 	redisclient "github.com/numaproj/numaflow/pkg/shared/clients/redis"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	sharedtls "github.com/numaproj/numaflow/pkg/shared/tls"
+	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 )
 
 type daemonServer struct {
@@ -64,12 +65,27 @@ func (ds *daemonServer) Run(ctx context.Context) error {
 	case v1alpha1.ISBSvcTypeJetStream:
 		isbSvcClient, err = isbsvc.NewISBJetStreamSvc(ds.pipeline.Name, isbsvc.WithJetStreamClient(jsclient.NewInClusterJetStreamClient()))
 		if err != nil {
-			log.Errorw("Failed to get a ISB Service client.", zap.Error(err))
+			log.Errorw("Failed to get an ISB Service client.", zap.Error(err))
 			return err
 		}
 	default:
 		return fmt.Errorf("unsupported isbsvc buffer type %q", ds.isbSvcType)
 	}
+	wmFetchers, err := service.GetVertexWatermarkFetchers(ctx, ds.pipeline, isbSvcClient)
+	if err != nil {
+		return fmt.Errorf("failed to get watermark fetchers, %w", err)
+	}
+
+	defer func() {
+		for _, fetcherList := range wmFetchers {
+			for _, f := range fetcherList {
+				if err := f.Close(); err != nil {
+					log.Errorw("Failed to close watermark fetcher", zap.Error(err))
+				}
+			}
+		}
+	}()
+
 	// Start listener
 	var conn net.Listener
 	var listerErr error
@@ -85,7 +101,7 @@ func (ds *daemonServer) Run(ctx context.Context) error {
 	}
 
 	tlsConfig := &tls.Config{Certificates: []tls.Certificate{*cer}, MinVersion: tls.VersionTLS12}
-	grpcServer, err := ds.newGRPCServer(isbSvcClient)
+	grpcServer, err := ds.newGRPCServer(isbSvcClient, wmFetchers)
 	if err != nil {
 		return fmt.Errorf("failed to create grpc server: %w", err)
 	}
@@ -106,7 +122,7 @@ func (ds *daemonServer) Run(ctx context.Context) error {
 	return nil
 }
 
-func (ds *daemonServer) newGRPCServer(isbSvcClient isbsvc.ISBService) (*grpc.Server, error) {
+func (ds *daemonServer) newGRPCServer(isbSvcClient isbsvc.ISBService, wmFetchers map[string][]fetch.Fetcher) (*grpc.Server, error) {
 	// "Prometheus histograms are a great way to measure latency distributions of your RPCs.
 	// However, since it is a bad practice to have metrics of high cardinality the latency monitoring metrics are disabled by default.
 	// To enable them please call the following in your server initialization code:"
@@ -120,7 +136,7 @@ func (ds *daemonServer) newGRPCServer(isbSvcClient isbsvc.ISBService) (*grpc.Ser
 	}
 	grpcServer := grpc.NewServer(sOpts...)
 	grpc_prometheus.Register(grpcServer)
-	pipelineMetadataQuery, err := service.NewPipelineMetadataQuery(isbSvcClient, ds.pipeline)
+	pipelineMetadataQuery, err := service.NewPipelineMetadataQuery(isbSvcClient, ds.pipeline, wmFetchers)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/daemon/server/service/pipeline_metrics_query.go b/pkg/daemon/server/service/pipeline_metrics_query.go
index ced6633c..54474984 100644
--- a/pkg/daemon/server/service/pipeline_metrics_query.go
+++ b/pkg/daemon/server/service/pipeline_metrics_query.go
@@ -34,6 +34,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/isbsvc"
 	metricspkg "github.com/numaproj/numaflow/pkg/metrics"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
+	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 )
 
 // metricsHttpClient interface for the GET call to metrics endpoint.
@@ -44,18 +45,19 @@ type metricsHttpClient interface {
 
 // pipelineMetadataQuery has the metadata required for the pipeline queries
 type pipelineMetadataQuery struct {
-	isbSvcClient    isbsvc.ISBService
-	pipeline        *v1alpha1.Pipeline
-	httpClient      metricsHttpClient
-	vertexWatermark *watermarkFetchers
+	isbSvcClient      isbsvc.ISBService
+	pipeline          *v1alpha1.Pipeline
+	httpClient        metricsHttpClient
+	watermarkFetchers map[string][]fetch.Fetcher
 }
 
 // NewPipelineMetadataQuery returns a new instance of pipelineMetadataQuery
-func NewPipelineMetadataQuery(isbSvcClient isbsvc.ISBService, pipeline *v1alpha1.Pipeline) (*pipelineMetadataQuery, error) {
+func NewPipelineMetadataQuery(isbSvcClient isbsvc.ISBService, pipeline *v1alpha1.Pipeline, wmFetchers map[string][]fetch.Fetcher) (*pipelineMetadataQuery, error) {
 	var err error
 	ps := pipelineMetadataQuery{
-		isbSvcClient: isbSvcClient,
-		pipeline:     pipeline,
+		isbSvcClient:      isbSvcClient,
+		pipeline:          pipeline,
+		watermarkFetchers: wmFetchers,
 		httpClient: &http.Client{
 			Transport: &http.Transport{
 				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
@@ -63,7 +65,6 @@ func NewPipelineMetadataQuery(isbSvcClient isbsvc.ISBService, pipeline *v1alpha1
 			Timeout: time.Second * 3,
 		},
 	}
-	ps.vertexWatermark, err = newVertexWatermarkFetcher(pipeline, isbSvcClient)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/daemon/server/service/pipeline_watermark_query.go b/pkg/daemon/server/service/pipeline_watermark_query.go
index 5f8cd9df..df2ff204 100644
--- a/pkg/daemon/server/service/pipeline_watermark_query.go
+++ b/pkg/daemon/server/service/pipeline_watermark_query.go
@@ -31,41 +31,27 @@ import (
 
 // TODO - Write Unit Tests for this file
 
-// watermarkFetchers used to store watermark metadata for propagation
-type watermarkFetchers struct {
-	// A map used to query watermark for vertices
-	// Key: vertex name
-	// Value: a list of watermark fetchers, each fetcher represents one of the out edges for the vertex
-	fetcherMap map[string][]fetch.Fetcher
-}
-
-// newVertexWatermarkFetcher creates a new instance of watermarkFetchers. This is used to populate a map of vertices to
-// corresponding fetchers. The fetchers are to retrieve vertex level watermarks.
-func newVertexWatermarkFetcher(pipeline *v1alpha1.Pipeline, isbSvcClient isbsvc.ISBService) (*watermarkFetchers, error) {
-	ctx := context.Background()
-	var wmFetcher = new(watermarkFetchers)
-
+// GetVertexWatermarkFetchers returns a map of the watermark fetchers, where key is the vertex name,
+// value is a list of fetchers to all the outgoing buffers of that vertex.
+func GetVertexWatermarkFetchers(ctx context.Context, pipeline *v1alpha1.Pipeline, isbSvcClient isbsvc.ISBService) (map[string][]fetch.Fetcher, error) {
+	var wmFetchers = make(map[string][]fetch.Fetcher)
 	if pipeline.Spec.Watermark.Disabled {
-		return wmFetcher, nil
+		return wmFetchers, nil
 	}
 
-	vertexToFetchersMap := make(map[string][]fetch.Fetcher)
-	pipelineName := pipeline.Name
-
-	// TODO: https://github.com/numaproj/numaflow/pull/120#discussion_r927316015
 	for _, vertex := range pipeline.Spec.Vertices {
 		if vertex.Sink != nil {
-			toBufferName := v1alpha1.GenerateSinkBufferName(pipeline.Namespace, pipelineName, vertex.Name)
-			fetchWatermark, err := isbSvcClient.CreateWatermarkFetcher(ctx, toBufferName)
+			toBufferName := v1alpha1.GenerateSinkBufferName(pipeline.Namespace, pipeline.Name, vertex.Name)
+			wmFetcher, err := isbSvcClient.CreateWatermarkFetcher(ctx, toBufferName)
 			if err != nil {
-				return nil, fmt.Errorf("failed to create watermark fetcher  %w", err)
+				return nil, fmt.Errorf("failed to create watermark fetcher, %w", err)
 			}
-			vertexToFetchersMap[vertex.Name] = []fetch.Fetcher{fetchWatermark}
+			wmFetchers[vertex.Name] = []fetch.Fetcher{wmFetcher}
 		} else {
 			// If the vertex is not a sink, to fetch the watermark, we consult all out edges and grab the latest watermark among them.
 			var wmFetcherList []fetch.Fetcher
 			for _, edge := range pipeline.GetToEdges(vertex.Name) {
-				toBufferNames := v1alpha1.GenerateEdgeBufferNames(pipeline.Namespace, pipelineName, edge)
+				toBufferNames := v1alpha1.GenerateEdgeBufferNames(pipeline.Namespace, pipeline.Name, edge)
 				for _, toBufferName := range toBufferNames {
 					fetchWatermark, err := isbSvcClient.CreateWatermarkFetcher(ctx, toBufferName)
 					if err != nil {
@@ -74,11 +60,10 @@ func newVertexWatermarkFetcher(pipeline *v1alpha1.Pipeline, isbSvcClient isbsvc.
 					wmFetcherList = append(wmFetcherList, fetchWatermark)
 				}
 			}
-			vertexToFetchersMap[vertex.Name] = wmFetcherList
+			wmFetchers[vertex.Name] = wmFetcherList
 		}
 	}
-	wmFetcher.fetcherMap = vertexToFetchersMap
-	return wmFetcher, nil
+	return wmFetchers, nil
 }
 
 // GetVertexWatermark is used to return the head watermark for a given vertex.
@@ -103,7 +88,7 @@ func (ps *pipelineMetadataQuery) GetVertexWatermark(ctx context.Context, request
 	}
 
 	// Watermark is enabled
-	vertexFetchers, ok := ps.vertexWatermark.fetcherMap[vertexName]
+	vertexFetchers, ok := ps.watermarkFetchers[vertexName]
 
 	// Vertex not found
 	if !ok {
diff --git a/pkg/isb/forward/forward.go b/pkg/isb/forward/forward.go
index 6483f68f..4d924a9f 100644
--- a/pkg/isb/forward/forward.go
+++ b/pkg/isb/forward/forward.go
@@ -50,7 +50,7 @@ type InterStepDataForward struct {
 	fromBuffer       isb.BufferReader
 	toBuffers        map[string]isb.BufferWriter
 	FSD              ToWhichStepDecider
-	UDF              udfapplier.Applier
+	UDF              udfapplier.MapApplier
 	fetchWatermark   fetch.Fetcher
 	publishWatermark map[string]publish.Publisher
 	opts             options
@@ -64,7 +64,7 @@ func NewInterStepDataForward(vertex *dfv1.Vertex,
 	fromStep isb.BufferReader,
 	toSteps map[string]isb.BufferWriter,
 	fsd ToWhichStepDecider,
-	applyUDF udfapplier.Applier,
+	applyUDF udfapplier.MapApplier,
 	fetchWatermark fetch.Fetcher,
 	publishWatermark map[string]publish.Publisher,
 	opts ...Option) (*InterStepDataForward, error) {
@@ -154,9 +154,16 @@ func (isdf *InterStepDataForward) Start() <-chan struct{} {
 			}
 		}
 
-		// stop watermark publisher if watermarking is enabled
+		// stop watermark fetcher
+		if err := isdf.fetchWatermark.Close(); err != nil {
+			log.Errorw("Failed to close watermark fetcher", zap.Error(err))
+		}
+
+		// stop watermark publisher
 		for _, publisher := range isdf.publishWatermark {
-			publisher.StopPublisher()
+			if err := publisher.Close(); err != nil {
+				log.Errorw("Failed to close watermark publisher", zap.Error(err))
+			}
 		}
 		close(stopped)
 	}()
@@ -433,7 +440,7 @@ func (isdf *InterStepDataForward) concurrentApplyUDF(ctx context.Context, readMe
 // The UserError retry will be done on the ApplyUDF.
 func (isdf *InterStepDataForward) applyUDF(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.Message, error) {
 	for {
-		writeMessages, err := isdf.UDF.Apply(ctx, readMessage)
+		writeMessages, err := isdf.UDF.ApplyMap(ctx, readMessage)
 		if err != nil {
 			isdf.opts.logger.Errorw("UDF.Apply error", zap.Error(err))
 			// TODO: implement retry with backoff etc.
diff --git a/pkg/isbsvc/jetstream_service.go b/pkg/isbsvc/jetstream_service.go
index 44c4cd2e..95fffd1a 100644
--- a/pkg/isbsvc/jetstream_service.go
+++ b/pkg/isbsvc/jetstream_service.go
@@ -27,7 +27,6 @@ import (
 	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
 	"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"
 	"github.com/spf13/viper"
@@ -128,7 +127,6 @@ func (jss *jetStreamSvc) CreateBuffers(ctx context.Context, buffers []dfv1.Buffe
 			}
 			if _, err := js.CreateKeyValue(&nats.KeyValueConfig{
 				Bucket:       otBucket,
-				Description:  fmt.Sprintf("Offset timeline bucket of buffer [%s]", buffer.Name),
 				MaxValueSize: v.GetInt32("otBucket.maxValueSize"),
 				History:      uint8(v.GetUint("otBucket.history")),
 				TTL:          v.GetDuration("otBucket.ttl"),
@@ -148,7 +146,6 @@ func (jss *jetStreamSvc) CreateBuffers(ctx context.Context, buffers []dfv1.Buffe
 			}
 			if _, err := js.CreateKeyValue(&nats.KeyValueConfig{
 				Bucket:       procBucket,
-				Description:  fmt.Sprintf("Processor bucket of buffer [%s]", buffer.Name),
 				MaxValueSize: v.GetInt32("procBucket.maxValueSize"),
 				History:      uint8(v.GetUint("procBucket.history")),
 				TTL:          v.GetDuration("procBucket.ttl"),
@@ -289,7 +286,7 @@ func (jss *jetStreamSvc) CreateWatermarkFetcher(ctx context.Context, bufferName
 	if err != nil {
 		return nil, err
 	}
-	watermarkFetcher := generic.NewGenericEdgeFetch(ctx, bufferName, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
+	watermarkFetcher := fetch.NewEdgeFetcher(ctx, bufferName, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
 	return watermarkFetcher, nil
 }
 
diff --git a/pkg/isbsvc/redis_service.go b/pkg/isbsvc/redis_service.go
index 734cdf86..36d21f80 100644
--- a/pkg/isbsvc/redis_service.go
+++ b/pkg/isbsvc/redis_service.go
@@ -19,8 +19,8 @@ package isbsvc
 import (
 	"context"
 	"fmt"
+
 	redis2 "github.com/numaproj/numaflow/pkg/isb/stores/redis"
-	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
 	"go.uber.org/multierr"
@@ -146,6 +146,6 @@ func (r *isbsRedisSvc) CreateWatermarkFetcher(ctx context.Context, bufferName st
 	// Watermark fetching is not supported for Redis ATM. Creating noop watermark fetcher.
 	hbWatcher := noop.NewKVOpWatch()
 	otWatcher := noop.NewKVOpWatch()
-	watermarkFetcher := generic.NewGenericEdgeFetch(ctx, bufferName, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher))
+	watermarkFetcher := fetch.NewEdgeFetcher(ctx, bufferName, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher))
 	return watermarkFetcher, nil
 }
diff --git a/pkg/reduce/pnf/processandforward.go b/pkg/reduce/pnf/processandforward.go
index a96472f2..526cee6f 100644
--- a/pkg/reduce/pnf/processandforward.go
+++ b/pkg/reduce/pnf/processandforward.go
@@ -39,14 +39,14 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/pbq"
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	udfreducer "github.com/numaproj/numaflow/pkg/udf/reducer"
+	"github.com/numaproj/numaflow/pkg/udf/applier"
 )
 
 // ProcessAndForward reads messages from pbq, invokes udf using grpc, forwards the results to ISB, and then publishes
 // the watermark for that partition.
 type ProcessAndForward struct {
 	PartitionID      partition.ID
-	UDF              udfreducer.Reducer
+	UDF              applier.ReduceApplier
 	result           []*isb.Message
 	pbqReader        pbq.Reader
 	log              *zap.SugaredLogger
@@ -58,7 +58,7 @@ type ProcessAndForward struct {
 // NewProcessAndForward will return a new ProcessAndForward instance
 func NewProcessAndForward(ctx context.Context,
 	partitionID partition.ID,
-	udf udfreducer.Reducer,
+	udf applier.ReduceApplier,
 	pbqReader pbq.Reader,
 	toBuffers map[string]isb.BufferWriter,
 	whereToDecider forward.ToWhichStepDecider, pw map[string]publish.Publisher) *ProcessAndForward {
@@ -79,7 +79,7 @@ func (p *ProcessAndForward) Process(ctx context.Context) error {
 
 	// FIXME: we need to fix https://github.com/numaproj/numaflow-go/blob/main/pkg/function/service.go#L101
 	// blocking call, only returns the result after it has read all the messages from pbq
-	p.result, err = p.UDF.Reduce(ctx, &p.PartitionID, p.pbqReader.ReadCh())
+	p.result, err = p.UDF.ApplyReduce(ctx, &p.PartitionID, p.pbqReader.ReadCh())
 	return err
 }
 
diff --git a/pkg/reduce/readloop/ordered.go b/pkg/reduce/readloop/ordered.go
index e3d99906..5afddeab 100644
--- a/pkg/reduce/readloop/ordered.go
+++ b/pkg/reduce/readloop/ordered.go
@@ -32,7 +32,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/reduce/pnf"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
-	udfreducer "github.com/numaproj/numaflow/pkg/udf/reducer"
+	"github.com/numaproj/numaflow/pkg/udf/applier"
 )
 
 var retryDelay = 1 * time.Second
@@ -69,7 +69,7 @@ func (of *orderedForwarder) startUp(ctx context.Context) {
 
 // schedulePnF creates and schedules the PnF routine.
 func (of *orderedForwarder) schedulePnF(ctx context.Context,
-	udf udfreducer.Reducer,
+	udf applier.ReduceApplier,
 	pbq pbq.Reader,
 	partitionID partition.ID,
 	toBuffers map[string]isb.BufferWriter,
diff --git a/pkg/reduce/readloop/readloop.go b/pkg/reduce/readloop/readloop.go
index 61822d25..78a3f5b2 100644
--- a/pkg/reduce/readloop/readloop.go
+++ b/pkg/reduce/readloop/readloop.go
@@ -35,15 +35,14 @@ import (
 	"go.uber.org/zap"
 	"k8s.io/apimachinery/pkg/util/wait"
 
-	"github.com/numaproj/numaflow/pkg/isb/forward"
-	"github.com/numaproj/numaflow/pkg/shared/logging"
-	udfReducer "github.com/numaproj/numaflow/pkg/udf/reducer"
-	"github.com/numaproj/numaflow/pkg/watermark/publish"
-
 	"github.com/numaproj/numaflow/pkg/isb"
+	"github.com/numaproj/numaflow/pkg/isb/forward"
 	"github.com/numaproj/numaflow/pkg/pbq"
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/shared/logging"
+	"github.com/numaproj/numaflow/pkg/udf/applier"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
+	"github.com/numaproj/numaflow/pkg/watermark/publish"
 	"github.com/numaproj/numaflow/pkg/window"
 	"github.com/numaproj/numaflow/pkg/window/keyed"
 	"github.com/numaproj/numaflow/pkg/window/strategy/fixed"
@@ -51,7 +50,7 @@ import (
 
 // ReadLoop is responsible for reading and forwarding the message from ISB to PBQ.
 type ReadLoop struct {
-	UDF               udfReducer.Reducer
+	UDF               applier.ReduceApplier
 	pbqManager        *pbq.Manager
 	windowingStrategy window.Windower
 	aw                *fixed.ActiveWindows
@@ -64,7 +63,7 @@ type ReadLoop struct {
 
 // NewReadLoop initializes  and returns ReadLoop.
 func NewReadLoop(ctx context.Context,
-	udf udfReducer.Reducer,
+	udf applier.ReduceApplier,
 	pbqManager *pbq.Manager,
 	windowingStrategy window.Windower,
 	toBuffers map[string]isb.BufferWriter,
diff --git a/pkg/reduce/reduce.go b/pkg/reduce/reduce.go
index ca44f163..d1471ead 100644
--- a/pkg/reduce/reduce.go
+++ b/pkg/reduce/reduce.go
@@ -30,7 +30,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/pbq"
 	"github.com/numaproj/numaflow/pkg/reduce/readloop"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
-	udfReducer "github.com/numaproj/numaflow/pkg/udf/reducer"
+	"github.com/numaproj/numaflow/pkg/udf/applier"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 	"github.com/numaproj/numaflow/pkg/window"
@@ -38,22 +38,24 @@ import (
 
 // DataForward reads data from isb and forwards them to readloop
 type DataForward struct {
-	fromBuffer        isb.BufferReader
-	readloop          *readloop.ReadLoop
-	fetchWatermark    fetch.Fetcher
-	windowingStrategy window.Windower
-	opts              *Options
-	log               *zap.SugaredLogger
+	fromBuffer          isb.BufferReader
+	toBuffers           map[string]isb.BufferWriter
+	readloop            *readloop.ReadLoop
+	watermarkFetcher    fetch.Fetcher
+	watermarkPublishers map[string]publish.Publisher
+	windowingStrategy   window.Windower
+	opts                *Options
+	log                 *zap.SugaredLogger
 }
 
 func NewDataForward(ctx context.Context,
-	udf udfReducer.Reducer,
+	udf applier.ReduceApplier,
 	fromBuffer isb.BufferReader,
 	toBuffers map[string]isb.BufferWriter,
 	pbqManager *pbq.Manager,
 	whereToDecider forward.ToWhichStepDecider,
 	fw fetch.Fetcher,
-	publishWatermark map[string]publish.Publisher,
+	watermarkPublishers map[string]publish.Publisher,
 	windowingStrategy window.Windower, opts ...Option) (*DataForward, error) {
 
 	options := DefaultOptions()
@@ -64,14 +66,16 @@ func NewDataForward(ctx context.Context,
 		}
 	}
 
-	rl := readloop.NewReadLoop(ctx, udf, pbqManager, windowingStrategy, toBuffers, whereToDecider, publishWatermark, options.windowOpts)
+	rl := readloop.NewReadLoop(ctx, udf, pbqManager, windowingStrategy, toBuffers, whereToDecider, watermarkPublishers, options.windowOpts)
 	return &DataForward{
-		fromBuffer:        fromBuffer,
-		readloop:          rl,
-		fetchWatermark:    fw,
-		windowingStrategy: windowingStrategy,
-		log:               logging.FromContext(ctx),
-		opts:              options}, nil
+		fromBuffer:          fromBuffer,
+		toBuffers:           toBuffers,
+		readloop:            rl,
+		watermarkFetcher:    fw,
+		watermarkPublishers: watermarkPublishers,
+		windowingStrategy:   windowingStrategy,
+		log:                 logging.FromContext(ctx),
+		opts:                options}, nil
 }
 
 // Start starts forwarding messages to readloop
@@ -80,7 +84,35 @@ func (d *DataForward) Start(ctx context.Context) {
 	for {
 		select {
 		case <-ctx.Done():
-			d.log.Infow("Stopping reduce data forwarder... ", zap.Error(ctx.Err()))
+			d.log.Infow("Stopping reduce data forwarder...")
+			cctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+			defer cancel()
+			d.readloop.ShutDown(cctx)
+
+			if err := d.fromBuffer.Close(); err != nil {
+				d.log.Errorw("Failed to close buffer reader, shutdown anyways...", zap.Error(err))
+			} else {
+				d.log.Infow("Closed buffer reader", zap.String("bufferFrom", d.fromBuffer.GetName()))
+			}
+			for _, v := range d.toBuffers {
+				if err := v.Close(); err != nil {
+					d.log.Errorw("Failed to close buffer writer, shutdown anyways...", zap.Error(err), zap.String("bufferTo", v.GetName()))
+				} else {
+					d.log.Infow("Closed buffer writer", zap.String("bufferTo", v.GetName()))
+				}
+			}
+
+			// stop watermark fetcher
+			if err := d.watermarkFetcher.Close(); err != nil {
+				d.log.Errorw("Failed to close watermark fetcher", zap.Error(err))
+			}
+
+			// stop watermark publisher
+			for _, publisher := range d.watermarkPublishers {
+				if err := publisher.Close(); err != nil {
+					d.log.Errorw("Failed to close watermark publisher", zap.Error(err))
+				}
+			}
 			return
 		default:
 			d.forwardAChunk(ctx)
@@ -102,7 +134,7 @@ func (d *DataForward) forwardAChunk(ctx context.Context) {
 
 	// fetch watermark using the first element's watermark, because we assign the watermark to all other
 	// elements in the batch based on the watermark we fetch from 0th offset.
-	processorWM := d.fetchWatermark.GetWatermark(readMessages[0].ReadOffset)
+	processorWM := d.watermarkFetcher.GetWatermark(readMessages[0].ReadOffset)
 	for _, m := range readMessages {
 		m.Watermark = time.Time(processorWM)
 	}
diff --git a/pkg/sources/generator/tickgen.go b/pkg/sources/generator/tickgen.go
index 5810a50f..6f3f2e5d 100644
--- a/pkg/sources/generator/tickgen.go
+++ b/pkg/sources/generator/tickgen.go
@@ -257,7 +257,10 @@ func (mg *memgen) Ack(_ context.Context, offsets []isb.Offset) []error {
 	return make([]error, len(offsets))
 }
 
-func (br *memgen) Close() error {
+func (mg *memgen) Close() error {
+	if err := mg.sourcePublishWM.Close(); err != nil {
+		mg.logger.Errorw("Failed to close source vertex watermark publisher", zap.Error(err))
+	}
 	return nil
 }
 
diff --git a/pkg/sources/http/http.go b/pkg/sources/http/http.go
index 4ecea9dd..39c8ee50 100644
--- a/pkg/sources/http/http.go
+++ b/pkg/sources/http/http.go
@@ -251,6 +251,9 @@ func (h *httpSource) Close() error {
 	if err := h.shutdown(context.Background()); err != nil {
 		return err
 	}
+	if err := h.sourcePublishWM.Close(); err != nil {
+		h.logger.Errorw("Failed to close source vertex watermark publisher", zap.Error(err))
+	}
 	h.logger.Info("HTTP source server shutdown")
 	return nil
 }
diff --git a/pkg/sources/kafka/reader.go b/pkg/sources/kafka/reader.go
index 5fa31dd3..0898cd5a 100644
--- a/pkg/sources/kafka/reader.go
+++ b/pkg/sources/kafka/reader.go
@@ -228,6 +228,11 @@ func (r *KafkaSource) Close() error {
 		}
 	}
 	<-r.stopch
+	for _, p := range r.sourcePublishWMs {
+		if err := p.Close(); err != nil {
+			r.logger.Errorw("Failed to close source vertex watermark publisher", zap.Error(err))
+		}
+	}
 	r.logger.Info("Kafka reader closed")
 	return nil
 }
diff --git a/pkg/sources/source.go b/pkg/sources/source.go
index 77fae141..658b42c3 100644
--- a/pkg/sources/source.go
+++ b/pkg/sources/source.go
@@ -86,6 +86,7 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
+
 		sourcePublisherStores, err = jetstream.BuildSourcePublisherStores(ctx, sp.VertexInstance)
 		if err != nil {
 			return err
diff --git a/pkg/udf/applier/applier.go b/pkg/udf/applier/map_applier.go
similarity index 63%
rename from pkg/udf/applier/applier.go
rename to pkg/udf/applier/map_applier.go
index 4998c683..805d6e5a 100644
--- a/pkg/udf/applier/applier.go
+++ b/pkg/udf/applier/map_applier.go
@@ -22,22 +22,22 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb"
 )
 
-// Applier applies the HTTPBasedUDF on the read message and gives back a new message. Any UserError will be retried here, while
+// MapApplier applies the HTTPBasedUDF on the read message and gives back a new message. Any UserError will be retried here, while
 // InternalErr can be returned and could be retried by the callee.
-type Applier interface {
-	Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error)
+type MapApplier interface {
+	ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error)
 }
 
 // ApplyFunc utility function used to create an Applier implementation
-type ApplyFunc func(context.Context, *isb.ReadMessage) ([]*isb.Message, error)
+type ApplyMapFunc func(context.Context, *isb.ReadMessage) ([]*isb.Message, error)
 
-func (a ApplyFunc) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (a ApplyMapFunc) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return a(ctx, message)
 }
 
 var (
 	// Terminal Applier do not make any change to the message
-	Terminal = ApplyFunc(func(ctx context.Context, msg *isb.ReadMessage) ([]*isb.Message, error) {
+	Terminal = ApplyMapFunc(func(ctx context.Context, msg *isb.ReadMessage) ([]*isb.Message, error) {
 		return []*isb.Message{&msg.Message}, nil
 	})
 )
diff --git a/pkg/udf/reducer/reducer.go b/pkg/udf/applier/reducer_applier.go
similarity index 55%
rename from pkg/udf/reducer/reducer.go
rename to pkg/udf/applier/reducer_applier.go
index 68acb367..5e212486 100644
--- a/pkg/udf/reducer/reducer.go
+++ b/pkg/udf/applier/reducer_applier.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package reducer
+package applier
 
 import (
 	"context"
@@ -23,15 +23,15 @@ import (
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
 )
 
-// Reducer applies the HTTPBasedUDF on the read message and gives back a new message. Any UserError will be retried here, while
+// ReduceApplier applies the HTTPBasedUDF on the read message and gives back a new message. Any UserError will be retried here, while
 // InternalErr can be returned and could be retried by the callee.
-type Reducer interface {
-	Reduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error)
+type ReduceApplier interface {
+	ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error)
 }
 
-// ReduceFunc utility function used to create a Reducer implementation
-type ReduceFunc func(context.Context, *partition.ID, <-chan *isb.ReadMessage) ([]*isb.Message, error)
+// ApplyReduceFunc utility function used to create a Reducer implementation
+type ApplyReduceFunc func(context.Context, *partition.ID, <-chan *isb.ReadMessage) ([]*isb.Message, error)
 
-func (a ReduceFunc) Reduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (a ApplyReduceFunc) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	return a(ctx, partitionID, messageStream)
 }
diff --git a/pkg/udf/function/uds_grpc.go b/pkg/udf/function/uds_grpc.go
index e5866f7d..7558a395 100644
--- a/pkg/udf/function/uds_grpc.go
+++ b/pkg/udf/function/uds_grpc.go
@@ -25,7 +25,6 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
-	"github.com/numaproj/numaflow/pkg/udf/reducer"
 
 	functionpb "github.com/numaproj/numaflow-go/pkg/apis/proto/function/v1"
 	functionsdk "github.com/numaproj/numaflow-go/pkg/function"
@@ -41,8 +40,8 @@ type udsGRPCBasedUDF struct {
 	client functionsdk.Client
 }
 
-var _ applier.Applier = (*udsGRPCBasedUDF)(nil)
-var _ reducer.Reducer = (*udsGRPCBasedUDF)(nil)
+var _ applier.MapApplier = (*udsGRPCBasedUDF)(nil)
+var _ applier.ReduceApplier = (*udsGRPCBasedUDF)(nil)
 
 // NewUDSGRPCBasedUDF returns a new udsGRPCBasedUDF object.
 func NewUDSGRPCBasedUDF() (*udsGRPCBasedUDF, error) {
@@ -78,7 +77,7 @@ func (u *udsGRPCBasedUDF) WaitUntilReady(ctx context.Context) error {
 	}
 }
 
-func (u *udsGRPCBasedUDF) Apply(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.Message, error) {
+func (u *udsGRPCBasedUDF) ApplyMap(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.Message, error) {
 	key := readMessage.Key
 	payload := readMessage.Body.Payload
 	offset := readMessage.ReadOffset
@@ -124,7 +123,7 @@ func (u *udsGRPCBasedUDF) Apply(ctx context.Context, readMessage *isb.ReadMessag
 // should we pass metadata information ?
 
 // Reduce accepts a channel of isbMessages and returns the aggregated result
-func (u *udsGRPCBasedUDF) Reduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (u *udsGRPCBasedUDF) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	datumCh := make(chan *functionpb.Datum)
 	var wg sync.WaitGroup
 	var result []*functionpb.Datum
diff --git a/pkg/watermark/fetch/edge_fetcher.go b/pkg/watermark/fetch/edge_fetcher.go
index 98ea39ae..1b545eef 100644
--- a/pkg/watermark/fetch/edge_fetcher.go
+++ b/pkg/watermark/fetch/edge_fetcher.go
@@ -28,24 +28,28 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
 // edgeFetcher is a fetcher between two vertices.
 type edgeFetcher struct {
 	ctx              context.Context
-	edgeName         string
+	bufferName       string
+	storeWatcher     store.WatermarkStoreWatcher
 	processorManager *ProcessorManager
 	log              *zap.SugaredLogger
 }
 
-// NewEdgeFetcher returns a new edge fetcher, processorManager has the details about the processors responsible for writing to this
-// edge.
-func NewEdgeFetcher(ctx context.Context, edgeName string, processorManager *ProcessorManager) Fetcher {
+// NewEdgeFetcher returns a new edge fetcher.
+func NewEdgeFetcher(ctx context.Context, bufferName string, storeWatcher store.WatermarkStoreWatcher) Fetcher {
+	log := logging.FromContext(ctx).With("bufferName", bufferName)
+	log.Info("Creating a new edge watermark fetcher")
 	return &edgeFetcher{
 		ctx:              ctx,
-		edgeName:         edgeName,
-		processorManager: processorManager,
-		log:              logging.FromContext(ctx).With("edgeName", edgeName),
+		bufferName:       bufferName,
+		storeWatcher:     storeWatcher,
+		processorManager: NewProcessorManager(ctx, storeWatcher),
+		log:              log,
 	}
 }
 
@@ -107,7 +111,17 @@ func (e *edgeFetcher) GetWatermark(inputOffset isb.Offset) processor.Watermark {
 	if epoch == math.MaxInt64 {
 		epoch = -1
 	}
-	e.log.Debugf("%s[%s] get watermark for offset %d: %+v", debugString.String(), e.edgeName, offset, epoch)
+	e.log.Debugf("%s[%s] get watermark for offset %d: %+v", debugString.String(), e.bufferName, offset, epoch)
 
 	return processor.Watermark(time.UnixMilli(epoch))
 }
+
+// Close function closes the watchers.
+func (e *edgeFetcher) Close() error {
+	e.log.Infof("Closing edge watermark fetcher")
+	if e.storeWatcher != nil {
+		e.storeWatcher.HeartbeatWatcher().Close()
+		e.storeWatcher.OffsetTimelineWatcher().Close()
+	}
+	return nil
+}
diff --git a/pkg/watermark/fetch/interface.go b/pkg/watermark/fetch/interface.go
index 7e1f6d1f..8aba4f3b 100644
--- a/pkg/watermark/fetch/interface.go
+++ b/pkg/watermark/fetch/interface.go
@@ -17,12 +17,15 @@ limitations under the License.
 package fetch
 
 import (
+	"io"
+
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 )
 
 // Fetcher fetches watermark data from Vn-1 vertex.
 type Fetcher interface {
+	io.Closer
 	// GetWatermark returns the inorder monotonically increasing watermark of the edge connected to Vn-1.
 	GetWatermark(offset isb.Offset) processor.Watermark
 	// GetHeadWatermark returns the latest watermark based on the head offset
diff --git a/pkg/watermark/fetch/processor_manager.go b/pkg/watermark/fetch/processor_manager.go
index 3e3f0e4d..9028d777 100644
--- a/pkg/watermark/fetch/processor_manager.go
+++ b/pkg/watermark/fetch/processor_manager.go
@@ -58,7 +58,6 @@ func NewProcessorManager(ctx context.Context, watermarkStoreWatcher store.Waterm
 	for _, opt := range inputOpts {
 		opt(opts)
 	}
-
 	v := &ProcessorManager{
 		ctx:        ctx,
 		hbWatcher:  watermarkStoreWatcher.HeartbeatWatcher(),
@@ -163,10 +162,6 @@ func (v *ProcessorManager) startHeatBeatWatcher() {
 	for {
 		select {
 		case <-stopped:
-			// main process exit
-			// close both watcher
-			v.otWatcher.Close()
-			v.hbWatcher.Close()
 			return
 		case value := <-watchCh:
 			if value == nil {
diff --git a/pkg/watermark/fetch/source_fetcher.go b/pkg/watermark/fetch/source_fetcher.go
index 774c73bb..92344dd1 100644
--- a/pkg/watermark/fetch/source_fetcher.go
+++ b/pkg/watermark/fetch/source_fetcher.go
@@ -21,28 +21,34 @@ import (
 	"math"
 	"time"
 
+	"go.uber.org/zap"
+
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
-	"go.uber.org/zap"
+	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
 // sourceFetcher is a fetcher on source buffers.
 type sourceFetcher struct {
 	ctx              context.Context
 	sourceBufferName string
+	storeWatcher     store.WatermarkStoreWatcher
 	processorManager *ProcessorManager
 	log              *zap.SugaredLogger
 }
 
 // NewSourceFetcher returns a new source fetcher, processorManager has the details about the processors responsible for writing to the
 // buckets of the source buffer.
-func NewSourceFetcher(ctx context.Context, sourceBufferName string, processorManager *ProcessorManager) Fetcher {
+func NewSourceFetcher(ctx context.Context, sourceBufferName string, storeWatcher store.WatermarkStoreWatcher) Fetcher {
+	log := logging.FromContext(ctx).With("sourceBufferName", sourceBufferName)
+	log.Info("Creating a new source watermark fetcher")
 	return &sourceFetcher{
 		ctx:              ctx,
 		sourceBufferName: sourceBufferName,
-		processorManager: processorManager,
-		log:              logging.FromContext(ctx).With("sourceBufferName", sourceBufferName),
+		storeWatcher:     storeWatcher,
+		processorManager: NewProcessorManager(ctx, storeWatcher),
+		log:              log,
 	}
 }
 
@@ -81,3 +87,13 @@ func (e *sourceFetcher) GetWatermark(_ isb.Offset) processor.Watermark {
 	}
 	return processor.Watermark(time.UnixMilli(epoch))
 }
+
+// Close function closes the watchers.
+func (e *sourceFetcher) Close() error {
+	e.log.Infof("Closing source watermark fetcher")
+	if e.storeWatcher != nil {
+		e.storeWatcher.HeartbeatWatcher().Close()
+		e.storeWatcher.OffsetTimelineWatcher().Close()
+	}
+	return nil
+}
diff --git a/pkg/watermark/generic/generic.go b/pkg/watermark/generic/generic.go
deleted file mode 100644
index b5cc970c..00000000
--- a/pkg/watermark/generic/generic.go
+++ /dev/null
@@ -1,36 +0,0 @@
-/*
-Copyright 2022 The Numaproj Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
-package generic
-
-import (
-	"context"
-
-	"github.com/numaproj/numaflow/pkg/watermark/fetch"
-	"github.com/numaproj/numaflow/pkg/watermark/store"
-)
-
-// NewGenericEdgeFetch returns a Fetcher, which fetches watermarks from bufferName by consulting the corresponding storeWatcher
-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, 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 8be57346..f0b3bce0 100644
--- a/pkg/watermark/generic/jetstream/generic.go
+++ b/pkg/watermark/generic/jetstream/generic.go
@@ -61,9 +61,9 @@ func BuildWatermarkProgressors(ctx context.Context, vertexInstance *v1alpha1.Ver
 	}
 
 	if fromBuffer.Type == v1alpha1.SourceBuffer {
-		fetchWatermark = generic.NewGenericSourceFetch(ctx, fromBuffer.Name, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
+		fetchWatermark = fetch.NewSourceFetcher(ctx, fromBuffer.Name, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
 	} else {
-		fetchWatermark = generic.NewGenericEdgeFetch(ctx, fromBuffer.Name, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
+		fetchWatermark = fetch.NewEdgeFetcher(ctx, fromBuffer.Name, store.BuildWatermarkStoreWatcher(hbWatch, otWatch))
 	}
 
 	// Publisher map creation, we need a publisher per out edge.
diff --git a/pkg/watermark/generic/noop.go b/pkg/watermark/generic/noop.go
index ae7770a3..efc6c9c4 100644
--- a/pkg/watermark/generic/noop.go
+++ b/pkg/watermark/generic/noop.go
@@ -56,8 +56,9 @@ func (n NoOpWMProgressor) GetHeadWatermark() processor.Watermark {
 	return processor.Watermark{}
 }
 
-// StopPublisher stops the no-op publisher.
-func (n NoOpWMProgressor) StopPublisher() {
+// Close stops the no-op progressor.
+func (n NoOpWMProgressor) Close() error {
+	return nil
 }
 
 func BuildNoOpWatermarkProgressorsFromEdgeList(bufferList []string) (fetch.Fetcher, map[string]publish.Publisher) {
diff --git a/pkg/watermark/publish/publisher.go b/pkg/watermark/publish/publisher.go
index ea7f9cf2..319f63a4 100644
--- a/pkg/watermark/publish/publisher.go
+++ b/pkg/watermark/publish/publisher.go
@@ -19,6 +19,7 @@ package publish
 import (
 	"context"
 	"fmt"
+	"io"
 	"time"
 
 	"github.com/numaproj/numaflow/pkg/isb"
@@ -31,12 +32,11 @@ import (
 
 // Publisher interface defines how to publish Watermark for a ProcessorEntitier.
 type Publisher interface {
+	io.Closer
 	// PublishWatermark publishes the watermark.
 	PublishWatermark(processor.Watermark, isb.Offset)
 	// GetLatestWatermark returns the latest published watermark.
 	GetLatestWatermark() processor.Watermark
-	// StopPublisher stops the publisher
-	StopPublisher()
 }
 
 // publish publishes the watermark for a processor entity.
@@ -54,9 +54,10 @@ type publish struct {
 
 // NewPublish returns `Publish`.
 func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier, watermarkStores store.WatermarkStorer, inputOpts ...PublishOption) Publisher {
-
-	log := logging.FromContext(ctx)
-
+	log := logging.FromContext(ctx).With("entityID", processorEntity.GetID()).
+		With("otStore", watermarkStores.OffsetTimelineStore().GetStoreName()).
+		With("hbStore", watermarkStores.HeartbeatStore().GetStoreName())
+	log.Info("Creating a new watermark publisher")
 	opts := &publishOptions{
 		autoRefreshHeartbeat: true,
 		podHeartbeatRate:     5,
@@ -81,7 +82,6 @@ func NewPublish(ctx context.Context, processorEntity processor.ProcessorEntitier
 	if opts.autoRefreshHeartbeat {
 		go p.publishHeartbeat()
 	}
-
 	return p
 }
 
@@ -183,20 +183,25 @@ func (p *publish) publishHeartbeat() {
 	}
 }
 
-// StopPublisher stops the publisher and cleans up the data associated with key.
-func (p *publish) StopPublisher() {
+// Close stops the publisher and cleans up the data associated with key.
+func (p *publish) Close() error {
+	p.log.Info("Closing watermark publisher")
+	defer func() {
+		if p.otStore != nil {
+			p.otStore.Close()
+		}
+		if p.heartbeatStore != nil {
+			p.heartbeatStore.Close()
+		}
+	}()
 	// TODO: cleanup after processor dies
 	//   - delete the Offset-Timeline bucket
 	//   - remove itself from heartbeat bucket
 
-	p.log.Infow("Stopping publisher", zap.String("bucket", p.heartbeatStore.GetStoreName()))
-
 	// clean up heartbeat bucket
-	err := p.heartbeatStore.DeleteKey(p.ctx, p.entity.GetID())
-	if err != nil {
+	if err := p.heartbeatStore.DeleteKey(p.ctx, p.entity.GetID()); err != nil {
 		p.log.Errorw("Failed to delete the key in the heartbeat bucket", zap.String("bucket", p.heartbeatStore.GetStoreName()), zap.String("key", p.entity.GetID()), zap.Error(err))
+		return err
 	}
-
-	p.otStore.Close()
-	p.heartbeatStore.Close()
+	return nil
 }
diff --git a/pkg/watermark/store/inmem/kv_store.go b/pkg/watermark/store/inmem/kv_store.go
index 060c127a..cbbd538a 100644
--- a/pkg/watermark/store/inmem/kv_store.go
+++ b/pkg/watermark/store/inmem/kv_store.go
@@ -31,11 +31,6 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
-var (
-	buckets     = make(map[string]*inMemStore)
-	bucketsLock sync.RWMutex
-)
-
 // kvEntry is each key-value entry in the store and the operation associated with the kv pair.
 type kvEntry struct {
 	key   string
@@ -79,9 +74,6 @@ func NewKVInMemKVStore(ctx context.Context, pipelineName string, bucketName stri
 		kvEntryCh:    make(chan store.WatermarkKVEntry, 10),
 		log:          logging.FromContext(ctx).With("pipeline", pipelineName).With("bucketName", bucketName),
 	}
-	bucketsLock.Lock()
-	buckets[bucketName] = s
-	bucketsLock.Unlock()
 	return s, s.kvEntryCh, nil
 }
 
@@ -150,7 +142,4 @@ func (kv *inMemStore) PutKV(_ context.Context, k string, v []byte) error {
 // Close closes the channel connection and clean up the bucket.
 func (kv *inMemStore) Close() {
 	close(kv.kvEntryCh)
-	bucketsLock.Lock()
-	delete(buckets, kv.bucketName)
-	bucketsLock.Unlock()
 }
diff --git a/pkg/watermark/store/jetstream/kv_watch.go b/pkg/watermark/store/jetstream/kv_watch.go
index a0fe00ae..cd4446c1 100644
--- a/pkg/watermark/store/jetstream/kv_watch.go
+++ b/pkg/watermark/store/jetstream/kv_watch.go
@@ -48,8 +48,7 @@ func NewKVJetStreamKVWatch(ctx context.Context, pipelineName string, kvBucketNam
 		return nil, fmt.Errorf("failed to get nats connection, %w", err)
 	}
 
-	// do we need to specify any opts? if yes, send it via options.
-	js, err := conn.JetStream(nats.PublishAsyncMaxPending(256))
+	js, err := conn.JetStream()
 	if err != nil {
 		if !conn.IsClosed() {
 			conn.Close()


Test patch:
diff --git a/pkg/daemon/server/service/pipeline_metrics_query_test.go b/pkg/daemon/server/service/pipeline_metrics_query_test.go
index 08a61fe2..e0e34fee 100644
--- a/pkg/daemon/server/service/pipeline_metrics_query_test.go
+++ b/pkg/daemon/server/service/pipeline_metrics_query_test.go
@@ -77,7 +77,7 @@ func TestGetVertexMetrics(t *testing.T) {
 		ObjectMeta: metav1.ObjectMeta{Name: pipelineName},
 	}
 	client, _ := isbsvc.NewISBJetStreamSvc(pipelineName)
-	pipelineMetricsQueryService, err := NewPipelineMetadataQuery(client, pipeline)
+	pipelineMetricsQueryService, err := NewPipelineMetadataQuery(client, pipeline, nil)
 	assert.NoError(t, err)
 
 	metricsResponse := `# HELP vertex_processing_rate Message processing rate in the last period of seconds, tps. It represents the rate of a vertex instead of a pod.
@@ -152,7 +152,7 @@ func TestGetBuffer(t *testing.T) {
 	}
 
 	ms := &mockIsbSvcClient{}
-	pipelineMetricsQueryService, err := NewPipelineMetadataQuery(ms, pipeline)
+	pipelineMetricsQueryService, err := NewPipelineMetadataQuery(ms, pipeline, nil)
 	assert.NoError(t, err)
 
 	bufferName := "numaflow-system-simple-pipeline-in-cat"
@@ -194,7 +194,7 @@ func TestListBuffers(t *testing.T) {
 	}
 
 	ms := &mockIsbSvcClient{}
-	pipelineMetricsQueryService, err := NewPipelineMetadataQuery(ms, pipeline)
+	pipelineMetricsQueryService, err := NewPipelineMetadataQuery(ms, pipeline, nil)
 	assert.NoError(t, err)
 
 	req := &daemon.ListBuffersRequest{Pipeline: &pipelineName}
diff --git a/pkg/isb/forward/forward_test.go b/pkg/isb/forward/forward_test.go
index e4b0afca..8082544a 100644
--- a/pkg/isb/forward/forward_test.go
+++ b/pkg/isb/forward/forward_test.go
@@ -48,7 +48,7 @@ func (f myForwardTest) WhereTo(_ string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
-func (f myForwardTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (f myForwardTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return testutils.CopyUDFTestApply(ctx, message)
 }
 
@@ -110,7 +110,7 @@ func (f myForwardDropTest) WhereTo(_ string) ([]string, error) {
 	return []string{"__DROP__"}, nil
 }
 
-func (f myForwardDropTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (f myForwardDropTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return testutils.CopyUDFTestApply(ctx, message)
 }
 
@@ -163,7 +163,7 @@ func (f myForwardApplyErrTest) WhereTo(_ string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
-func (f myForwardApplyErrTest) Apply(_ context.Context, _ *isb.ReadMessage) ([]*isb.Message, error) {
+func (f myForwardApplyErrTest) ApplyMap(_ context.Context, _ *isb.ReadMessage) ([]*isb.Message, error) {
 	return nil, udfapplier.ApplyUDFErr{
 		UserUDFErr: false,
 		InternalErr: struct {
@@ -215,7 +215,7 @@ func (f myForwardApplyWhereToErrTest) WhereTo(_ string) ([]string, error) {
 	return []string{"to1"}, fmt.Errorf("whereToStep failed")
 }
 
-func (f myForwardApplyWhereToErrTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (f myForwardApplyWhereToErrTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return testutils.CopyUDFTestApply(ctx, message)
 }
 
@@ -262,7 +262,7 @@ func (f myForwardApplyUDFErrTest) WhereTo(_ string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
-func (f myForwardApplyUDFErrTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (f myForwardApplyUDFErrTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 
 	return nil, fmt.Errorf("UDF error")
 }
@@ -311,7 +311,7 @@ func (f myForwardToAllTest) WhereTo(_ string) ([]string, error) {
 	return []string{dfv1.MessageKeyAll}, nil
 }
 
-func (f myForwardToAllTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (f myForwardToAllTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return testutils.CopyUDFTestApply(ctx, message)
 }
 
diff --git a/pkg/isb/forward/shutdown_test.go b/pkg/isb/forward/shutdown_test.go
index f15d0e02..a3b68884 100644
--- a/pkg/isb/forward/shutdown_test.go
+++ b/pkg/isb/forward/shutdown_test.go
@@ -38,7 +38,7 @@ func (s myShutdownTest) WhereTo(_ string) ([]string, error) {
 	return []string{dfv1.MessageKeyAll}, nil
 }
 
-func (s myShutdownTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (s myShutdownTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return testutils.CopyUDFTestApply(ctx, message)
 }
 
diff --git a/pkg/isb/stores/jetstream/writer_test.go b/pkg/isb/stores/jetstream/writer_test.go
index 45718c72..cad8576a 100644
--- a/pkg/isb/stores/jetstream/writer_test.go
+++ b/pkg/isb/stores/jetstream/writer_test.go
@@ -42,7 +42,7 @@ func (f myForwardJetStreamTest) WhereTo(_ string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
-func (f myForwardJetStreamTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (f myForwardJetStreamTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return testutils.CopyUDFTestApply(ctx, message)
 }
 
diff --git a/pkg/isb/stores/redis/read_test.go b/pkg/isb/stores/redis/read_test.go
index dcc44d04..c3c69c3b 100644
--- a/pkg/isb/stores/redis/read_test.go
+++ b/pkg/isb/stores/redis/read_test.go
@@ -296,7 +296,7 @@ func (f forwardReadWritePerformance) WhereTo(_ string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
-func (f forwardReadWritePerformance) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (f forwardReadWritePerformance) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return testutils.CopyUDFTestApply(ctx, message)
 }
 
diff --git a/pkg/isb/stores/redis/write_test.go b/pkg/isb/stores/redis/write_test.go
index dadf5913..b4160873 100644
--- a/pkg/isb/stores/redis/write_test.go
+++ b/pkg/isb/stores/redis/write_test.go
@@ -306,7 +306,7 @@ func (f myForwardRedisTest) WhereTo(_ string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
-func (f myForwardRedisTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
+func (f myForwardRedisTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.Message, error) {
 	return testutils.CopyUDFTestApply(ctx, message)
 }
 
diff --git a/pkg/reduce/readloop/ordered_test.go b/pkg/reduce/readloop/ordered_test.go
index 7bb47ce1..267d3576 100644
--- a/pkg/reduce/readloop/ordered_test.go
+++ b/pkg/reduce/readloop/ordered_test.go
@@ -31,7 +31,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/pbq"
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/pbq/store"
-	"github.com/numaproj/numaflow/pkg/udf/reducer"
+	"github.com/numaproj/numaflow/pkg/udf/applier"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -49,7 +49,7 @@ func (f myForwardTest) Apply(ctx context.Context, message *isb.ReadMessage) ([]*
 func TestOrderedProcessing(t *testing.T) {
 
 	// Test Reducer returns the messages as is
-	identityReducer := reducer.ReduceFunc(func(ctx context.Context, partitionID *partition.ID, input <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+	identityReducer := applier.ApplyReduceFunc(func(ctx context.Context, partitionID *partition.ID, input <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 		messages := make([]*isb.Message, 0)
 		for msg := range input {
 			messages = append(messages, &msg.Message)
diff --git a/pkg/reduce/reduce_test.go b/pkg/reduce/reduce_test.go
index f8f822e1..a10ec907 100644
--- a/pkg/reduce/reduce_test.go
+++ b/pkg/reduce/reduce_test.go
@@ -50,7 +50,8 @@ func (e EventTypeWMProgressor) GetLatestWatermark() processor.Watermark {
 	return processor.Watermark{}
 }
 
-func (e EventTypeWMProgressor) StopPublisher() {
+func (e EventTypeWMProgressor) Close() error {
+	return nil
 }
 
 func (e EventTypeWMProgressor) GetWatermark(_ isb.Offset) processor.Watermark {
@@ -71,7 +72,7 @@ type CounterReduceTest struct {
 }
 
 // Reduce returns a result with the count of messages
-func (f CounterReduceTest) Reduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (f CounterReduceTest) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	count := 0
 	for range messageStream {
 		count += 1
@@ -103,7 +104,7 @@ func (f CounterReduceTest) WhereTo(s string) ([]string, error) {
 type SumReduceTest struct {
 }
 
-func (s SumReduceTest) Reduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (s SumReduceTest) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	sum := 0
 	for msg := range messageStream {
 		var payload PayloadForTest
@@ -133,7 +134,7 @@ func (s SumReduceTest) Reduce(ctx context.Context, partitionID *partition.ID, me
 type MaxReduceTest struct {
 }
 
-func (m MaxReduceTest) Reduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (m MaxReduceTest) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	mx := math.MinInt64
 	for msg := range messageStream {
 		var payload PayloadForTest
@@ -553,6 +554,8 @@ func fetcherAndPublisher(ctx context.Context, toBuffers map[string]isb.BufferWri
 	// create publisher for to Buffers
 	for key := range toBuffers {
 		publishEntity := processor.NewProcessorEntity(key)
+		hb, _, _ := inmem.NewKVInMemKVStore(ctx, pipelineName, key+"_PROCESSORS")
+		ot, _, _ := inmem.NewKVInMemKVStore(ctx, pipelineName, key+"_OT")
 		p := publish.NewPublish(ctx, publishEntity, wmstore.BuildWatermarkStore(hb, ot), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
 		publishers[key] = p
 	}
@@ -576,9 +579,7 @@ func fetcherAndPublisher(ctx context.Context, toBuffers map[string]isb.BufferWri
 	hbWatcher, _ := inmem.NewInMemWatch(ctx, pipelineName, keyspace+"_PROCESSORS", hbWatcherCh)
 	otWatcher, _ := inmem.NewInMemWatch(ctx, pipelineName, keyspace+"_OT", otWatcherCh)
 
-	// fetcher for reduce
-	var pm = fetch.NewProcessorManager(ctx, wmstore.BuildWatermarkStoreWatcher(hbWatcher, otWatcher), fetch.WithPodHeartbeatRate(1))
-	var f = fetch.NewEdgeFetcher(ctx, fromBuffer.GetName(), pm)
+	var f = fetch.NewEdgeFetcher(ctx, fromBuffer.GetName(), wmstore.BuildWatermarkStoreWatcher(hbWatcher, otWatcher))
 	return f, publishers
 }
 
diff --git a/pkg/udf/function/uds_grpc_test.go b/pkg/udf/function/uds_grpc_test.go
index 9676efa9..418f7404 100644
--- a/pkg/udf/function/uds_grpc_test.go
+++ b/pkg/udf/function/uds_grpc_test.go
@@ -111,7 +111,7 @@ func TestGRPCBasedUDF_BasicApplyWithMockClient(t *testing.T) {
 		}()
 
 		u := NewMockUDSGRPCBasedUDF(mockClient)
-		got, err := u.Apply(ctx, &isb.ReadMessage{
+		got, err := u.ApplyMap(ctx, &isb.ReadMessage{
 			Message: isb.Message{
 				Header: isb.Header{
 					PaneInfo: isb.PaneInfo{
@@ -155,7 +155,7 @@ func TestGRPCBasedUDF_BasicApplyWithMockClient(t *testing.T) {
 		}()
 
 		u := NewMockUDSGRPCBasedUDF(mockClient)
-		_, err := u.Apply(ctx, &isb.ReadMessage{
+		_, err := u.ApplyMap(ctx, &isb.ReadMessage{
 			Message: isb.Message{
 				Header: isb.Header{
 					PaneInfo: isb.PaneInfo{
@@ -235,7 +235,7 @@ func TestHGRPCBasedUDF_ApplyWithMockClient(t *testing.T) {
 	var results = make([][]byte, len(readMessages))
 	var resultKeys = make([]string, len(readMessages))
 	for idx, readMessage := range readMessages {
-		apply, err := u.Apply(ctx, &readMessage)
+		apply, err := u.ApplyMap(ctx, &readMessage)
 		assert.NoError(t, err)
 		results[idx] = apply[0].Payload
 		resultKeys[idx] = string(apply[0].Header.Key)
@@ -307,7 +307,7 @@ func TestGRPCBasedUDF_BasicReduceWithMockClient(t *testing.T) {
 			Key:   "test",
 		}
 
-		got, err := u.Reduce(ctx, partitionID, messageCh)
+		got, err := u.ApplyReduce(ctx, partitionID, messageCh)
 
 		assert.Len(t, got, 1)
 		assert.NoError(t, err)
@@ -358,7 +358,7 @@ func TestGRPCBasedUDF_BasicReduceWithMockClient(t *testing.T) {
 			Key:   "test",
 		}
 
-		_, err := u.Reduce(ctx, partitionID, messageCh)
+		_, err := u.ApplyReduce(ctx, partitionID, messageCh)
 
 		assert.ErrorIs(t, err, ApplyUDFErr{
 			UserUDFErr: false,
@@ -400,7 +400,7 @@ func TestGRPCBasedUDF_BasicReduceWithMockClient(t *testing.T) {
 			End:   time.Unix(120, 0),
 			Key:   "test",
 		}
-		_, err := u.Reduce(ctx, partitionID, messageCh)
+		_, err := u.ApplyReduce(ctx, partitionID, messageCh)
 
 		assert.Error(t, err, ctx.Err())
 	})
@@ -470,7 +470,7 @@ func TestHGRPCBasedUDF_Reduce(t *testing.T) {
 		Key:   "test",
 	}
 
-	result, err := u.Reduce(ctx, partitionID, messageCh)
+	result, err := u.ApplyReduce(ctx, partitionID, messageCh)
 
 	var resultPayload testutils.PayloadForTest
 	_ = json.Unmarshal(result[0].Payload, &resultPayload)
diff --git a/pkg/watermark/fetch/edge_fetcher_test.go b/pkg/watermark/fetch/edge_fetcher_test.go
index 484d4861..9a4f6922 100644
--- a/pkg/watermark/fetch/edge_fetcher_test.go
+++ b/pkg/watermark/fetch/edge_fetcher_test.go
@@ -134,7 +134,7 @@ func TestBuffer_GetWatermark(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			b := &edgeFetcher{
 				ctx:              ctx,
-				edgeName:         "testBuffer",
+				bufferName:       "testBuffer",
 				processorManager: tt.processorManager,
 				log:              zaptest.NewLogger(t).Sugar(),
 			}
diff --git a/pkg/watermark/fetch/processor_manager_inmem_test.go b/pkg/watermark/fetch/processor_manager_inmem_test.go
index b6db5793..240c1a29 100644
--- a/pkg/watermark/fetch/processor_manager_inmem_test.go
+++ b/pkg/watermark/fetch/processor_manager_inmem_test.go
@@ -77,8 +77,7 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 	assert.NoError(t, err)
 	otWatcher, err := inmem.NewInMemWatch(ctx, "testFetch", keyspace+"_OT", otWatcherCh)
 	assert.NoError(t, err)
-	var testVertex = NewProcessorManager(ctx, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher), WithPodHeartbeatRate(1), WithRefreshingProcessorsRate(1))
-	var testBuffer = NewEdgeFetcher(ctx, "testBuffer", testVertex).(*edgeFetcher)
+	var testBuffer = NewEdgeFetcher(ctx, "testBuffer", store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher)).(*edgeFetcher)
 
 	// start p1 heartbeat for 3 loops
 	wg.Add(1)
@@ -88,7 +87,7 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 		for i := 0; i < 3; i++ {
 			err = hbStore.PutKV(ctx, "p1", []byte(fmt.Sprintf("%d", time.Now().Unix())))
 			assert.NoError(t, err)
-			time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+			time.Sleep(1 * time.Second)
 		}
 		err = hbStore.DeleteKey(ctx, "p1")
 		assert.NoError(t, err)
@@ -102,7 +101,7 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 		for i := 0; i < 20; i++ {
 			err = hbStore.PutKV(ctx, "p2", []byte(fmt.Sprintf("%d", time.Now().Unix())))
 			assert.NoError(t, err)
-			time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+			time.Sleep(1 * time.Second)
 		}
 	}()
 
@@ -162,12 +161,12 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 		for i := 0; i < 5; i++ {
 			err = hbStore.PutKV(ctx, "p1", []byte(fmt.Sprintf("%d", time.Now().Unix())))
 			assert.NoError(t, err)
-			time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+			time.Sleep(1 * time.Second)
 		}
 	}()
 
 	// wait until p1 becomes active
-	time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+	time.Sleep(1 * time.Second)
 	allProcessors = testBuffer.processorManager.GetAllProcessors()
 	for !allProcessors["p1"].IsActive() {
 		select {
@@ -234,7 +233,7 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 		for i := 0; i < 3; i++ {
 			err = hbStore.PutKV(ctx, "p1", []byte(fmt.Sprintf("%d", time.Now().Unix())))
 			assert.NoError(t, err)
-			time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+			time.Sleep(1 * time.Second)
 		}
 	}()
 
diff --git a/pkg/watermark/fetch/processor_manager_test.go b/pkg/watermark/fetch/processor_manager_test.go
index 9e534d12..c75b3385 100644
--- a/pkg/watermark/fetch/processor_manager_test.go
+++ b/pkg/watermark/fetch/processor_manager_test.go
@@ -127,8 +127,7 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 	assert.NoError(t, err)
 	otWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", keyspace+"_OT", defaultJetStreamClient)
 	assert.NoError(t, err)
-	var testVertex = NewProcessorManager(ctx, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher), WithPodHeartbeatRate(1), WithRefreshingProcessorsRate(1))
-	var testBuffer = NewEdgeFetcher(ctx, "testBuffer", testVertex).(*edgeFetcher)
+	var testBuffer = NewEdgeFetcher(ctx, "testBuffer", store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher)).(*edgeFetcher)
 
 	wg.Add(1)
 	go func() {
@@ -137,7 +136,7 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 		for i := 0; i < 3; i++ {
 			err = hbStore.PutKV(ctx, "p1", []byte(fmt.Sprintf("%d", time.Now().Unix())))
 			assert.NoError(t, err)
-			time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+			time.Sleep(1 * time.Second)
 		}
 		err = hbStore.DeleteKey(ctx, "p1")
 		assert.NoError(t, err)
@@ -151,7 +150,7 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 		for i := 0; i < 20; i++ {
 			err = hbStore.PutKV(ctx, "p2", []byte(fmt.Sprintf("%d", time.Now().Unix())))
 			assert.NoError(t, err)
-			time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+			time.Sleep(1 * time.Second)
 		}
 	}()
 
@@ -222,12 +221,12 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 		for i := 0; i < 5; i++ {
 			err = hbStore.PutKV(ctx, "p1", []byte(fmt.Sprintf("%d", time.Now().Unix())))
 			assert.NoError(t, err)
-			time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+			time.Sleep(1 * time.Second)
 		}
 	}()
 
 	// wait until p1 becomes active
-	time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+	time.Sleep(1 * time.Second)
 	allProcessors = testBuffer.processorManager.GetAllProcessors()
 	for !allProcessors["p1"].IsActive() {
 		select {
@@ -293,7 +292,7 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 		for i := 0; i < 3; i++ {
 			err = hbStore.PutKV(ctx, "p1", []byte(fmt.Sprintf("%d", time.Now().Unix())))
 			assert.NoError(t, err)
-			time.Sleep(time.Duration(testVertex.opts.podHeartbeatRate) * time.Second)
+			time.Sleep(1 * time.Second)
 		}
 	}()
 
diff --git a/pkg/watermark/publish/publisher_inmem_test.go b/pkg/watermark/publish/publisher_inmem_test.go
index 2dcd2a19..f2699c66 100644
--- a/pkg/watermark/publish/publisher_inmem_test.go
+++ b/pkg/watermark/publish/publisher_inmem_test.go
@@ -68,7 +68,7 @@ func TestPublisherWithSharedOTBuckets_InMem(t *testing.T) {
 	head := p.GetLatestWatermark()
 	assert.Equal(t, processor.Watermark(time.UnixMilli(epoch-60000).In(location)).String(), head.String())
 
-	p.StopPublisher()
+	_ = p.Close()
 
 	_, err = p.heartbeatStore.GetValue(ctx, publishEntity.GetID())
 	assert.Equal(t, fmt.Errorf("key publisherTestPod1 not found"), err)
diff --git a/pkg/watermark/publish/publisher_test.go b/pkg/watermark/publish/publisher_test.go
index b4d70b3d..65824a6c 100644
--- a/pkg/watermark/publish/publisher_test.go
+++ b/pkg/watermark/publish/publisher_test.go
@@ -92,7 +92,7 @@ func TestPublisherWithSharedOTBucket(t *testing.T) {
 	head := p.GetLatestWatermark()
 	assert.Equal(t, processor.Watermark(time.UnixMilli(epoch-60000).In(location)).String(), head.String())
 
-	p.StopPublisher()
+	_ = p.Close()
 
 	_, err = p.heartbeatStore.GetValue(ctx, publishEntity.GetID())
 	assert.Equal(t, nats.ErrConnectionClosed, err)