fix: remove from tracker after windows are created and fix concurrency issues with rater (#3051)

Signed-off-by: Yashash H L <[email protected]>

Author: yhl25

pkg/daemon/server/service/rater/helper.go

@@ -31,48 +31,6 @@ const (
 	rateNotAvailable = float64(math.MinInt)
 )
 
-// UpdatePendingCount updates the pending count for a pod at a given time
-func UpdatePendingCount(q *sharedqueue.OverflowQueue[*TimestampedCounts], time int64, podPendingCounts *PodPendingCount) {
-	if podPendingCounts == nil {
-		return
-	}
-
-	items := q.Items()
-	// find the element matching the input timestamp and update it
-	for _, i := range items {
-		if i.timestamp == time {
-			i.UpdatePending(podPendingCounts)
-			return
-		}
-	}
-
-	// if we cannot find a matching element, it means we need to add a new timestamped count to the queue
-	tc := NewTimestampedCounts(time)
-	tc.UpdatePending(podPendingCounts)
-	q.Append(tc)
-}
-
-// UpdateCount updates the count of processed messages for a pod at a given time
-func UpdateCount(q *sharedqueue.OverflowQueue[*TimestampedCounts], time int64, podReadCounts *PodReadCount) {
-	if podReadCounts == nil {
-		return
-	}
-
-	items := q.Items()
-	// find the element matching the input timestamp and update it
-	for _, i := range items {
-		if i.timestamp == time {
-			i.Update(podReadCounts)
-			return
-		}
-	}
-
-	// if we cannot find a matching element, it means we need to add a new timestamped count to the queue
-	tc := NewTimestampedCounts(time)
-	tc.Update(podReadCounts)
-	q.Append(tc)
-}
-
 // CalculatePending calculates the pending messages for a given partition in the last lookback seconds
 func CalculatePending(q *sharedqueue.OverflowQueue[*TimestampedCounts], lookbackSeconds int64, partitionName string) int64 {
 	counts := q.Items()

pkg/daemon/server/service/rater/helper_test.go

@@ -27,84 +27,6 @@ import (
 
 const TestTime = 1620000000
 
-func TestUpdateCount(t *testing.T) {
-	t.Run("givenTimeExistsPodExistsPartitionExistsCountAvailable_whenUpdate_thenUpdatePodPartitionCount", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodReadCount{"pod1", map[string]float64{"partition1": 10.0}})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, &PodReadCount{"pod1", map[string]float64{"partition1": 20.0}})
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 20.0, q.Items()[0].podPartitionCount["pod1"]["partition1"])
-	})
-
-	t.Run("givenTimeExistsPodExistsPartitionNotExistsCountAvailable_whenUpdate_thenAddPodPartitionCount", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodReadCount{"pod1", map[string]float64{"partition1": 10.0}})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, &PodReadCount{"pod1", map[string]float64{"partition1": 20.0, "partition2": 30.0}})
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 20.0, q.Items()[0].podPartitionCount["pod1"]["partition1"])
-		assert.Equal(t, 30.0, q.Items()[0].podPartitionCount["pod1"]["partition2"])
-	})
-
-	t.Run("givenTimeExistsPodNotExistsCountAvailable_whenUpdate_thenAddPodCount", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodReadCount{"pod1", map[string]float64{"partition1": 20.0}})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, &PodReadCount{"pod2", map[string]float64{"partition1": 10.0}})
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 20.0, q.Items()[0].podPartitionCount["pod1"]["partition1"])
-		assert.Equal(t, 10.0, q.Items()[0].podPartitionCount["pod2"]["partition1"])
-	})
-
-	t.Run("givenTimeExistsPodExistsCountNotAvailable_whenUpdate_thenNotUpdatePod", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodReadCount{"pod1", map[string]float64{"partition1": 10.0}})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, nil)
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 1, len(q.Items()[0].podPartitionCount))
-		assert.Equal(t, 10.0, q.Items()[0].podPartitionCount["pod1"]["partition1"])
-	})
-
-	t.Run("givenTimeExistsPodNotExistsCountNotAvailable_whenUpdate_thenNoUpdate", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodReadCount{"pod1", map[string]float64{"partition1": 10.0}})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, nil)
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 10.0, q.Items()[0].podPartitionCount["pod1"]["partition1"])
-	})
-
-	t.Run("givenTimeNotExistsCountAvailable_whenUpdate_thenAddNewItem", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodReadCount{"pod1", map[string]float64{"partition1": 10.0}})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime+1, &PodReadCount{"pod1", map[string]float64{"partition1": 20.0}})
-
-		assert.Equal(t, 2, q.Length())
-		assert.Equal(t, 10.0, q.Items()[0].podPartitionCount["pod1"]["partition1"])
-		assert.Equal(t, 20.0, q.Items()[1].podPartitionCount["pod1"]["partition1"])
-	})
-}
-
 func TestCalculateRate(t *testing.T) {
 	t.Run("givenCollectedTimeLessThanTwo_whenCalculateRate_thenReturnRateNotAvailable", func(t *testing.T) {
 		q := sharedqueue.New[*TimestampedCounts](1800)

pkg/daemon/server/service/rater/options.go

@@ -16,11 +16,13 @@ limitations under the License.
 
 package rater
 
+import "time"
+
 type options struct {
 	// Number of workers working on collecting counts of processed messages.
 	workers int
-	// Time in milliseconds, each element in the work queue will be picked up in an interval of this period of time.
-	taskInterval int
+	// Time in seconds, each element in the work queue will be picked up in an interval of this period of time.
+	taskInterval time.Duration
 }
 
 type Option func(*options)
@@ -34,7 +36,7 @@ func defaultOptions() *options {
 	return &options{
 		workers: 50, // default max replicas is 50
 		// we execute the rater metrics fetching every 5 seconds
-		taskInterval: 5,
+		taskInterval: 5 * time.Second,
 	}
 }
 
@@ -44,8 +46,8 @@ func WithWorkers(n int) Option {
 	}
 }
 
-func WithTaskInterval(n int) Option {
+func WithTaskInterval(duration time.Duration) Option {
 	return func(o *options) {
-		o.taskInterval = n
+		o.taskInterval = duration
 	}
 }

pkg/daemon/server/service/rater/rater.go

@@ -24,6 +24,7 @@ import (
 	"maps"
 	"math"
 	"net/http"
+	"sync"
 	"time"
 
 	dto "github.com/prometheus/client_model/go"
@@ -69,6 +70,7 @@ type Rater struct {
 	timestampedPendingCount map[string]*sharedqueue.OverflowQueue[*TimestampedCounts]
 	// this can be updated dynamically, defaults to user-specified value in the spec
 	lookBackSeconds map[string]*atomic.Float64
+	mu              sync.RWMutex
 	options         *options
 }
 
@@ -184,15 +186,65 @@ func (r *Rater) monitorOnePod(ctx context.Context, key string, worker int, times
 			if podPendingCount == nil {
 				log.Debugf("Failed retrieving pending counts for pod %s vertex %s", podInfo.podName, podInfo.vertexName)
 			}
-			UpdatePendingCount(r.timestampedPendingCount[podInfo.vertexName], timestamp, podPendingCount)
+			r.updatePendingCount(podInfo.vertexName, timestamp, podPendingCount)
 		}
 	} else {
 		log.Debugf("Pod %s does not exist, updating it with nil...", podInfo.podName)
 	}
-	UpdateCount(r.timestampedPodCounts[podInfo.vertexName], timestamp, podReadCount)
+	r.updateCount(podInfo.vertexName, timestamp, podReadCount)
 	return nil
 }
 
+func (r *Rater) updateCount(vertexName string, timestamp int64, podReadCount *PodReadCount) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if podReadCount == nil {
+		return
+	}
+
+	q := r.timestampedPodCounts[vertexName]
+	// Peek the last element
+	items := q.Items()
+	if len(items) > 0 {
+		last := items[len(items)-1]
+		// Merge if same timestamp
+		if last.timestamp == timestamp {
+			last.Update(podReadCount)
+			return
+		}
+	}
+
+	tc := NewTimestampedCounts(timestamp)
+	tc.Update(podReadCount)
+	q.Append(tc)
+}
+
+func (r *Rater) updatePendingCount(vertexName string, timestamp int64, podPendingCounts *PodPendingCount) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if podPendingCounts == nil {
+		return
+	}
+	q := r.timestampedPendingCount[vertexName]
+
+	items := q.Items()
+	if len(items) > 0 {
+		last := items[len(items)-1]
+		// Merge if same timestamp
+		if last.timestamp == timestamp {
+			last.UpdatePending(podPendingCounts)
+			return
+		}
+	}
+
+	// if we cannot find a matching element, it means we need to add a new timestamped count to the queue
+	tc := NewTimestampedCounts(timestamp)
+	tc.UpdatePending(podPendingCounts)
+	q.Append(tc)
+}
+
 func (r *Rater) Start(ctx context.Context) error {
 	r.log.Infof("Starting rater...")
 	taskCh := make(chan *podTask)
@@ -225,7 +277,7 @@ func (r *Rater) Start(ctx context.Context) error {
 		}
 	}
 
-	ticker := time.NewTicker(time.Duration(r.options.taskInterval) * time.Millisecond)
+	ticker := time.NewTicker(r.options.taskInterval)
 	defer ticker.Stop()
 
 	// Following for loop keeps calling assign() function to assign monitoring tasks to the workers.

pkg/daemon/server/service/rater/rater_test.go

@@ -108,7 +108,7 @@ func TestRater_Start(t *testing.T) {
 			},
 		},
 	}
-	r := NewRater(ctx, pipeline, WithTaskInterval(1000))
+	r := NewRater(ctx, pipeline, WithTaskInterval(time.Second))
 	podTracker := NewPodTracker(ctx, pipeline, WithRefreshInterval(time.Second*1))
 	podTracker.httpClient = &raterMockHttpClient{podOneCount: 0, podTwoCount: 0, lock: &sync.RWMutex{}}
 	r.httpClient = &raterMockHttpClient{podOneCount: 0, podTwoCount: 0, lock: &sync.RWMutex{}}

pkg/mvtxdaemon/server/service/rater/helper.go

@@ -32,27 +32,6 @@ const (
 	rateNotAvailable = float64(math.MinInt)
 )
 
-// UpdateCount updates the count for a given timestamp in the queue.
-func UpdateCount(q *sharedqueue.OverflowQueue[*TimestampedCounts], time int64, podReadCounts *PodMetricsCount) {
-	if podReadCounts == nil {
-		return
-	}
-
-	items := q.Items()
-	// find the element matching the input timestamp and update it
-	for _, i := range items {
-		if i.timestamp == time {
-			i.Update(podReadCounts)
-			return
-		}
-	}
-
-	// if we cannot find a matching element, it means we need to add a new timestamped count to the queue
-	tc := NewTimestampedCounts(time)
-	tc.Update(podReadCounts)
-	q.Append(tc)
-}
-
 // CalculateRate calculates the rate of a MonoVertex for a given lookback period.
 func CalculateRate(q *sharedqueue.OverflowQueue[*TimestampedCounts], lookbackSeconds int64) float64 {
 	counts := q.Items()

pkg/mvtxdaemon/server/service/rater/helper_test.go

@@ -30,71 +30,6 @@ import (
 
 const TestTime = 1620000000
 
-func TestUpdateCount(t *testing.T) {
-	t.Run("givenTimeExistsPodExistsCountAvailable_whenUpdate_thenUpdatePodPartitionCount", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodMetricsCount{"pod1", 10.0})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, &PodMetricsCount{"pod1", 20.0})
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 20.0, q.Items()[0].podReadCounts["pod1"])
-	})
-
-	t.Run("givenTimeExistsPodNotExistsCountAvailable_whenUpdate_thenAddPodCount", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodMetricsCount{"pod1", 20.0})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, &PodMetricsCount{"pod2", 10.0})
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 20.0, q.Items()[0].podReadCounts["pod1"])
-		assert.Equal(t, 10.0, q.Items()[0].podReadCounts["pod2"])
-	})
-
-	t.Run("givenTimeExistsPodExistsCountNotAvailable_whenUpdate_thenNotUpdatePod", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodMetricsCount{"pod1", 10.0})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, nil)
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 1, len(q.Items()[0].podReadCounts))
-		assert.Equal(t, 10.0, q.Items()[0].podReadCounts["pod1"])
-	})
-
-	t.Run("givenTimeExistsPodNotExistsCountNotAvailable_whenUpdate_thenNoUpdate", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodMetricsCount{"pod1", 10.0})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime, nil)
-
-		assert.Equal(t, 1, q.Length())
-		assert.Equal(t, 10.0, q.Items()[0].podReadCounts["pod1"])
-	})
-
-	t.Run("givenTimeNotExistsCountAvailable_whenUpdate_thenAddNewItem", func(t *testing.T) {
-		q := sharedqueue.New[*TimestampedCounts](1800)
-		tc := NewTimestampedCounts(TestTime)
-		tc.Update(&PodMetricsCount{"pod1", 10.0})
-		q.Append(tc)
-
-		UpdateCount(q, TestTime+1, &PodMetricsCount{"pod1", 20.0})
-
-		assert.Equal(t, 2, q.Length())
-		assert.Equal(t, 10.0, q.Items()[0].podReadCounts["pod1"])
-		assert.Equal(t, 20.0, q.Items()[1].podReadCounts["pod1"])
-	})
-}
-
 func TestCalculatePending(t *testing.T) {
 	t.Run("givenCollectedTimeLessThanTwo_whenCalculateRate_thenReturnPendingNotAvailable", func(t *testing.T) {
 		q := sharedqueue.New[*TimestampedCounts](1800)