Repository: numaproj/numaflow
Pull request #939 (2023-08-09 07:05:32)

fix: retry logic for fetching last updated kv time
if there are no keys in the store we should not retry while trying to calculate the last updated kv time.

Signed-off-by: Yashash H L <yashashhl25@gmail.com>


Patch:
diff --git a/pkg/shared/kvs/inmem/kv_watch.go b/pkg/shared/kvs/inmem/kv_watch.go
index 341db6d5..329cf7b8 100644
--- a/pkg/shared/kvs/inmem/kv_watch.go
+++ b/pkg/shared/kvs/inmem/kv_watch.go
@@ -34,6 +34,7 @@ type inMemWatch struct {
 	kvEntryCh    <-chan kvs.KVEntry
 	kvHistory    []kvs.KVEntry
 	updatesChMap map[string]chan kvs.KVEntry
+	doneCh       chan struct{}
 	lock         sync.Mutex // the lock for kvHistory and updatesChMap
 	log          *zap.SugaredLogger
 }
@@ -48,6 +49,7 @@ func NewInMemWatch(ctx context.Context, pipelineName string, bucketName string,
 		kvEntryCh:    kvEntryCh,
 		kvHistory:    []kvs.KVEntry{},
 		updatesChMap: make(map[string]chan kvs.KVEntry),
+		doneCh:       make(chan struct{}),
 		log:          logging.FromContext(ctx).With("pipeline", pipelineName).With("bucketName", bucketName),
 	}
 	return k, nil
@@ -95,6 +97,10 @@ func (k *inMemWatch) Watch(ctx context.Context) (<-chan kvs.KVEntry, <-chan stru
 					ch <- value
 				}
 				k.lock.Unlock()
+			case <-k.doneCh:
+				close(updates)
+				close(stopped)
+				return
 			}
 		}
 	}()
@@ -108,4 +114,5 @@ func (k *inMemWatch) GetKVName() string {
 
 // Close does nothing here because the updates channel has already been closed when exiting the Watch.
 func (k *inMemWatch) Close() {
+	close(k.doneCh)
 }
diff --git a/pkg/shared/kvs/jetstream/kv_watch.go b/pkg/shared/kvs/jetstream/kv_watch.go
index 8de6ffc7..8be96860 100644
--- a/pkg/shared/kvs/jetstream/kv_watch.go
+++ b/pkg/shared/kvs/jetstream/kv_watch.go
@@ -161,9 +161,11 @@ func (jsw *jetStreamWatch) Watch(ctx context.Context) (<-chan kvs.KVEntry, <-cha
 			case <-jsw.kvwTimer.C:
 				// if the timer expired, it means that the watcher is not receiving any updates
 				kvLastUpdatedTime := jsw.lastUpdateKVTime()
+
+				// if the last update time is zero, it means that there are no key-value pairs in the store yet or ctx was canceled both the cases we should not recreate the watcher
 				// if the last update time is before the previous fetch time, it means that the store is not getting any updates
 				// therefore, we don't have to recreate the watcher
-				if kvLastUpdatedTime.Before(jsw.previousFetchTime) {
+				if kvLastUpdatedTime.IsZero() || kvLastUpdatedTime.Before(jsw.previousFetchTime) {
 					jsw.log.Debug("The watcher is not receiving any updates, but the store is not getting any updates either", zap.String("watcher", jsw.GetKVName()), zap.Time("lastUpdateKVTime", kvLastUpdatedTime), zap.Time("previousFetchTime", jsw.previousFetchTime))
 				} else {
 					// if the last update time is after the previous fetch time, it means that the store is getting updates but the watcher is not receiving any
@@ -203,17 +205,39 @@ func (jsw *jetStreamWatch) newWatcher(ctx context.Context) nats.KeyWatcher {
 	return kvWatcher
 }
 
-// lastUpdateKVTime returns the last update time of the kv store
+// lastUpdateKVTime returns the last update time of the kv store. if the key is missing, it will return time.Zero.
 func (jsw *jetStreamWatch) lastUpdateKVTime() time.Time {
-	keys, err := jsw.kvStore.Keys()
-	for err != nil {
-		jsw.log.Errorw("Failed to get keys", zap.String("watcher", jsw.GetKVName()), zap.Error(err))
-		keys, err = jsw.kvStore.Keys()
-		time.Sleep(100 * time.Millisecond)
+	var (
+		keys       []string
+		err        error
+		lastUpdate time.Time
+		value      nats.KeyValueEntry
+	)
+
+retryLoop:
+	for {
+		select {
+		case <-jsw.ctx.Done():
+			return time.Time{}
+		default:
+			keys, err = jsw.kvStore.Keys()
+			if err == nil {
+				break retryLoop
+			} else {
+				// if there are no keys in the store, return zero time because there are no updates
+				// upstream will handle it
+				if err == nats.ErrNoKeysFound {
+					return time.Time{}
+				}
+				jsw.log.Errorw("Failed to get keys", zap.String("watcher", jsw.GetKVName()), zap.Error(err))
+			}
+			time.Sleep(100 * time.Millisecond)
+		}
+
 	}
-	var lastUpdate = time.Time{}
+
 	for _, key := range keys {
-		value, err := jsw.kvStore.Get(key)
+		value, err = jsw.kvStore.Get(key)
 		for err != nil {
 			select {
 			case <-jsw.ctx.Done():


Test patch:
diff --git a/pkg/sources/forward/data_forward_test.go b/pkg/sources/forward/data_forward_test.go
index 93ee148e..d235ae13 100644
--- a/pkg/sources/forward/data_forward_test.go
+++ b/pkg/sources/forward/data_forward_test.go
@@ -225,6 +225,12 @@ func TestNewDataForward(t *testing.T) {
 			fetchWatermark := &testForwardFetcher{}
 			publishWatermark, otStores := buildPublisherMapAndOTStore(toSteps)
 
+			defer func() {
+				for _, p := range publishWatermark {
+					_ = p.Close()
+				}
+			}()
+
 			f, err := NewDataForward(vertex, fromStep, toSteps, &myForwardToAllTest{}, &myForwardToAllTest{}, fetchWatermark, publishWatermark, TestSourceWatermarkPublisher{}, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
@@ -371,6 +377,13 @@ func TestNewDataForward(t *testing.T) {
 
 			fetchWatermark := &testForwardFetcher{}
 			publishWatermark, otStores := buildPublisherMapAndOTStore(toSteps)
+
+			defer func() {
+				for _, p := range publishWatermark {
+					_ = p.Close()
+				}
+			}()
+
 			f, err := NewDataForward(vertex, fromStep, toSteps, myForwardDropTest{}, myForwardDropTest{}, fetchWatermark, publishWatermark, TestSourceWatermarkPublisher{}, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
@@ -529,6 +542,13 @@ func TestNewDataForward(t *testing.T) {
 
 			fetchWatermark := &testForwardFetcher{}
 			publishWatermark, otStores := buildPublisherMapAndOTStore(toSteps)
+
+			defer func() {
+				for _, p := range publishWatermark {
+					_ = p.Close()
+				}
+			}()
+
 			f, err := NewDataForward(vertex, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, fetchWatermark, publishWatermark, TestSourceWatermarkPublisher{}, WithReadBatchSize(batchSize))
 
 			assert.NoError(t, err)
diff --git a/pkg/sources/nats/nats_test.go b/pkg/sources/nats/nats_test.go
index a6a1d697..4a600ee4 100644
--- a/pkg/sources/nats/nats_test.go
+++ b/pkg/sources/nats/nats_test.go
@@ -90,8 +90,8 @@ func Test_Single(t *testing.T) {
 	defer server.Shutdown()
 
 	url := "127.0.0.1"
-	testSubject := "test"
-	testQueue := "test-queue"
+	testSubject := "test-single"
+	testQueue := "test-queue-single"
 	vi := testVertex(t, url, testSubject, testQueue, "test-host", 0)
 	ns, err := newInstance(t, vi)
 	assert.NoError(t, err)
@@ -102,13 +102,15 @@ func Test_Single(t *testing.T) {
 	nc, err := natslib.Connect(url)
 	assert.NoError(t, err)
 	defer nc.Close()
-	_ = nc.Publish(testSubject, []byte("1"))
-	_ = nc.Publish(testSubject, []byte("2"))
-	_ = nc.Publish(testSubject, []byte("3"))
 
-	msgs, err := ns.Read(context.Background(), 5)
-	readMessagesCount := len(msgs)
-	assert.NoError(t, err)
+	for i := 0; i < 3; i++ {
+		err = nc.Publish(testSubject, []byte(fmt.Sprintf("%d", i)))
+		assert.NoError(t, err)
+	}
+
+	var msgs []*isb.ReadMessage
+	var readMessagesCount int
+
 loop:
 	for {
 		select {
@@ -116,7 +118,7 @@ loop:
 			t.Fatal("timeout waiting for messages")
 			return
 		default:
-			msgs, err = ns.Read(context.Background(), 5)
+			msgs, err = ns.Read(ctx, 5)
 			assert.NoError(t, err)
 			readMessagesCount += len(msgs)
 			if readMessagesCount == 3 {
@@ -132,8 +134,8 @@ func Test_Multiple(t *testing.T) {
 	defer server.Shutdown()
 
 	url := "127.0.0.1"
-	testSubject := "test"
-	testQueue := "test-queue"
+	testSubject := "test-multiple"
+	testQueue := "test-queue-multiple"
 	v1 := testVertex(t, url, testSubject, testQueue, "test-host1", 0)
 	ns1, err := newInstance(t, v1)
 	assert.NoError(t, err)
@@ -150,7 +152,7 @@ func Test_Multiple(t *testing.T) {
 	assert.NoError(t, err)
 	defer nc.Close()
 	for i := 0; i < 5; i++ {
-		err := nc.Publish(testSubject, []byte(fmt.Sprint(i)))
+		err = nc.Publish(testSubject, []byte(fmt.Sprint(i)))
 		assert.NoError(t, err)
 	}