Repository: numaproj/numaflow
Pull request #902 (2023-07-28 17:42:09)

refactor: avoid exposing internal data structures of pod tracker to the rater
This is a small refactor to modularize the pod tracker. The data structure of `UniqueStringList` is the implementation detail of pod tracker, hence should not be exposed to the rater. This makes less work for the rater as rater doesn't need to call `getActivePods` followed by `.Contains` to check if a pod is active. This also removes the risk of the rater directly manipulating the `UniqueStringList`. Lastly, if in the future, we decide to replace `UniqueStringList` with another data structure, no code change needs to be made on the rater side.


Patch:
diff --git a/pkg/daemon/server/service/rater/pod_tracker.go b/pkg/daemon/server/service/rater/pod_tracker.go
index 2a288ee9..45ae8d72 100644
--- a/pkg/daemon/server/service/rater/pod_tracker.go
+++ b/pkg/daemon/server/service/rater/pod_tracker.go
@@ -114,15 +114,31 @@ func (pt *PodTracker) Start(ctx context.Context) error {
 	return nil
 }
 
+// LeastRecentlyUsed returns the least recently used pod from the active pod list.
+// if there is no active pods, it returns an empty string.
+func (pt *PodTracker) LeastRecentlyUsed() string {
+	if e := pt.activePods.Front(); e != "" {
+		pt.activePods.MoveToBack(e)
+		return e
+	}
+	return ""
+}
+
+// IsActive returns true if the pod is active, false otherwise.
+func (pt *PodTracker) IsActive(podKey string) bool {
+	return pt.activePods.Contains(podKey)
+}
+
+// GetActivePodsCount returns the number of active pods.
+func (pt *PodTracker) GetActivePodsCount() int {
+	return pt.activePods.Length()
+}
+
 func (pt *PodTracker) getPodKey(index int, vertexName string, vertexType string) string {
 	// podKey is used as a unique identifier for the pod, it is used by worker to determine the count of processed messages of the pod.
 	return strings.Join([]string{pt.pipeline.Name, vertexName, fmt.Sprintf("%d", index), vertexType}, PodInfoSeparator)
 }
 
-func (pt *PodTracker) GetActivePods() *UniqueStringList {
-	return pt.activePods
-}
-
 func (pt *PodTracker) isActive(vertexName, podName string) bool {
 	// using the vertex headless service to check if a pod exists or not.
 	// example for 0th pod : https://simple-pipeline-in-0.simple-pipeline-in-headless.default.svc:2469/metrics
diff --git a/pkg/daemon/server/service/rater/rater.go b/pkg/daemon/server/service/rater/rater.go
index 1b7b1b2b..b05805b8 100644
--- a/pkg/daemon/server/service/rater/rater.go
+++ b/pkg/daemon/server/service/rater/rater.go
@@ -135,8 +135,7 @@ func (r *Rater) monitorOnePod(ctx context.Context, key string, worker int) error
 	vertexType := podInfo[3]
 	podName := strings.Join([]string{podInfo[0], podInfo[1], podInfo[2]}, "-")
 	var podReadCount *PodReadCount
-	activePods := r.podTracker.GetActivePods()
-	if activePods.Contains(key) {
+	if r.podTracker.IsActive(key) {
 		podReadCount = r.getPodReadCounts(vertexName, vertexType, podName)
 		if podReadCount == nil {
 			log.Debugf("Failed retrieving total podReadCount for pod %s", podName)
@@ -168,12 +167,9 @@ func (r *Rater) Start(ctx context.Context) error {
 		go r.monitor(ctx, i, keyCh)
 	}
 
-	// Function assign() moves an element in the list from the front to the back,
-	// and send to the channel so that it can be picked up by a worker.
+	// Function assign() sends the least recently used podKey to the channel so that it can be picked up by a worker.
 	assign := func() {
-		activePods := r.podTracker.GetActivePods()
-		if e := activePods.Front(); e != "" {
-			activePods.MoveToBack(e)
+		if e := r.podTracker.LeastRecentlyUsed(); e != "" {
 			keyCh <- e
 			return
 		}
@@ -190,7 +186,7 @@ func (r *Rater) Start(ctx context.Context) error {
 			assign()
 			// Make sure each of the key will be assigned at least every taskInterval milliseconds.
 			sleep(ctx, time.Millisecond*time.Duration(func() int {
-				l := r.podTracker.GetActivePods().Length()
+				l := r.podTracker.GetActivePodsCount()
 				if l == 0 {
 					return r.options.taskInterval
 				}


Test patch:
diff --git a/pkg/daemon/server/service/rater/pod_tracker_test.go b/pkg/daemon/server/service/rater/pod_tracker_test.go
index c8ed868b..95aaff5e 100644
--- a/pkg/daemon/server/service/rater/pod_tracker_test.go
+++ b/pkg/daemon/server/service/rater/pod_tracker_test.go
@@ -29,6 +29,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
@@ -95,7 +96,7 @@ func TestPodTracker_Start(t *testing.T) {
 		}
 	}()
 
-	for tracker.activePods.Length() != 10 {
+	for tracker.GetActivePodsCount() != 10 {
 		select {
 		case <-ctx.Done():
 			t.Fatalf("incorrect active pods %v", ctx.Err())
@@ -106,7 +107,7 @@ func TestPodTracker_Start(t *testing.T) {
 
 	tracker.httpClient.(*trackerMockHttpClient).setPodsCount(5)
 
-	for tracker.activePods.Length() != 5 {
+	for tracker.GetActivePodsCount() != 5 {
 		select {
 		case <-ctx.Done():
 			t.Fatalf("incorrect active pods %v", ctx.Err())
@@ -116,4 +117,9 @@ func TestPodTracker_Start(t *testing.T) {
 	}
 	cancel()
 	wg.Wait()
+
+	assert.Equal(t, "p*v*0*other", tracker.LeastRecentlyUsed())
+	assert.Equal(t, "p*v*1*other", tracker.LeastRecentlyUsed())
+	assert.Equal(t, true, tracker.IsActive("p*v*4*other"))
+	assert.Equal(t, false, tracker.IsActive("p*v*5*other"))
 }