Repository: numaproj/numaflow
Pull request #242 (2022-10-20 07:55:29)

fix: watermark watcher leak
- fix: https://github.com/numaproj/numaflow/issues/220
- also previous fix for the issue https://github.com/numaproj/numaflow/issues/224 has a bug. fixed again in this PR.
- add a stopped channel for the watermark watcher to guarantee the watcher stops before connection closes.
- update keyed_test.go

# Flow (based on my observation)

## If The Pod Terminates Properly

When a pod in Vn-1 terminates, the pod is deleted from the Vn-1_Vn heartbeat store.
https://github.com/numaproj/numaflow/blob/3ec5785e0db1bac5be2f7243eb5b45bcf10279b5/pkg/watermark/publish/publisher.go#L178

The Vn pod heartbeat watcher receives the delete op, sets the pod status to _delete, stops the pod offsetTimeline watcher, and removes the pod from heartbeat map (so refreshing heartbeat go routine no longer updates on this deleted pod).
https://github.com/numaproj/numaflow/blob/3ec5785e0db1bac5be2f7243eb5b45bcf10279b5/pkg/watermark/fetch/processor_manager.go#L182-L193

stop the pod offsetTimeline watcher calls the cancel function and stops the watcher.
https://github.com/numaproj/numaflow/blob/3ec5785e0db1bac5be2f7243eb5b45bcf10279b5/pkg/watermark/fetch/processor_to_fetch.go#L53-L57
https://github.com/numaproj/numaflow/blob/3ec5785e0db1bac5be2f7243eb5b45bcf10279b5/pkg/watermark/fetch/processor_to_fetch.go#L101-L103
https://github.com/numaproj/numaflow/blob/46aae1d62b7a795ad171c45433ef7321794f5c7b/pkg/watermark/store/jetstream/kv_watch.go#L112
https://github.com/numaproj/numaflow/blob/46aae1d62b7a795ad171c45433ef7321794f5c7b/pkg/watermark/store/jetstream/kv_watch.go#L122
https://github.com/numaproj/numaflow/blob/46aae1d62b7a795ad171c45433ef7321794f5c7b/pkg/watermark/fetch/processor_to_fetch.go#L105-L113



Although the Vn-1 pod has been deleted, the events up to some watermark `w1` has already been processed. Therefore, we still preserve the in-mem offset timeline for the deleted pod until the head watermark of this pod has been consumed.
https://github.com/numaproj/numaflow/blob/425b6bbfaaf9885cf210197b84079b72ef7eafcb/pkg/watermark/fetch/edge_fetcher.go#L81-L84

## If The Pod Crashes
If a pod's heartbeat stops after 10 heartbeat refreshes, we consider the pod as crashed. 
https://github.com/numaproj/numaflow/blob/3ec5785e0db1bac5be2f7243eb5b45bcf10279b5/pkg/watermark/fetch/processor_manager.go#L120-L128

# test run:

simple pipeline
cat pod from 1 -> 4 -> 1
<img width="1704" alt="image" src="https://user-images.githubusercontent.com/19543684/196923277-38ab727e-eb06-4ad1-9e43-9af5a14070c6.png">

cat pod from 1 -> 10 -> 1
<img width="1557" alt="image" src="https://user-images.githubusercontent.com/19543684/197071860-f38832a2-3348-438f-879d-bdf54da5f106.png">


out pod from 1 -> 4 -> 1
<img width="1704" alt="image" src="https://user-images.githubusercontent.com/19543684/196923561-60eea860-0fb3-465d-bbcf-6fe31afb8d14.png">

out pod from 10 -> 1
<img width="1704" alt="image" src="https://user-images.githubusercontent.com/19543684/196925049-1d195253-ced3-4f5b-b4a9-27ba78a914af.png">

chromevoid: How to fix e2e tests?… @vigith @whynowy 

Patch:
diff --git a/pkg/watermark/fetch/edge_fetcher.go b/pkg/watermark/fetch/edge_fetcher.go
index 98ddb80c..54a59b87 100644
--- a/pkg/watermark/fetch/edge_fetcher.go
+++ b/pkg/watermark/fetch/edge_fetcher.go
@@ -73,14 +73,15 @@ func (e *edgeFetcher) GetWatermark(inputOffset isb.Offset) processor.Watermark {
 	var epoch int64 = math.MaxInt64
 	var allProcessors = e.processorManager.GetAllProcessors()
 	for _, p := range allProcessors {
-		if !p.IsActive() {
-			continue
-		}
 		debugString.WriteString(fmt.Sprintf("[Processor: %v] \n", p))
 		var t = p.offsetTimeline.GetEventTime(inputOffset)
 		if t != -1 && t < epoch {
 			epoch = t
 		}
+		if p.IsDeleted() && (offset > p.offsetTimeline.GetHeadOffset()) {
+			// if the pod is not active and the current offset is ahead of all offsets in Timeline
+			e.processorManager.DeleteProcessor(p.entity.GetID())
+		}
 	}
 	// if the offset is smaller than every offset in the timeline, set the value to be -1
 	if epoch == math.MaxInt64 {
diff --git a/pkg/watermark/fetch/processor_manager.go b/pkg/watermark/fetch/processor_manager.go
index 4c22c8e4..b0c8d060 100644
--- a/pkg/watermark/fetch/processor_manager.go
+++ b/pkg/watermark/fetch/processor_manager.go
@@ -18,10 +18,13 @@ import (
 // ProcessorManager manages the point of view of Vn-1 from Vn vertex processors (or source processor). The code is running on Vn vertex.
 // It has the mapping of all the processors which in turn has all the information about each processor timelines.
 type ProcessorManager struct {
-	ctx        context.Context
-	hbWatcher  store.WatermarkKVWatcher
-	otWatcher  store.WatermarkKVWatcher
-	heartbeat  *ProcessorHeartbeat
+	ctx       context.Context
+	hbWatcher store.WatermarkKVWatcher
+	otWatcher store.WatermarkKVWatcher
+	// heartbeat just tracks the heartbeat of each processing unit. we use it to mark a processing unit's status (e.g, inactive)
+	heartbeat *ProcessorHeartbeat
+	// processors has reference to the actual processing unit (ProcessorEntity) which includes offset timeline which is
+	// used for tracking watermark.
 	processors map[string]*ProcessorToFetch
 	lock       sync.RWMutex
 	log        *zap.SugaredLogger
@@ -73,7 +76,6 @@ func (v *ProcessorManager) GetProcessor(processor string) *ProcessorToFetch {
 }
 
 // DeleteProcessor deletes a processor.
-// Note: This operation is not used anywhere ATM.
 func (v *ProcessorManager) DeleteProcessor(processor string) {
 	v.lock.Lock()
 	defer v.lock.Unlock()
@@ -118,7 +120,16 @@ func (v *ProcessorManager) refreshingProcessors() {
 		}
 		// default heartbeat rate is every 5 seconds
 		// TODO: tolerance?
-		if time.Now().Unix()-pTime > v.opts.podHeartbeatRate {
+		if time.Now().Unix()-pTime > 10*v.opts.podHeartbeatRate {
+			// if the pod doesn't come back after 10 heartbeats,
+			// it's possible the pod has exited unexpectedly, so we need to delete the pod
+			// NOTE: the pod entry still remains in the heartbeat store (bucket)
+			// TODO: how to delete the pod from the heartbeat store?
+			v.log.Infow("Processor has been inactive for 10 heartbeats, deleting...", zap.String("key", pName), zap.String(pName, p.String()))
+			p.setStatus(_deleted)
+			p.stopTimeLineWatcher()
+			v.heartbeat.Delete(pName)
+		} else if time.Now().Unix()-pTime > v.opts.podHeartbeatRate {
 			// if the pod's last heartbeat is greater than podHeartbeatRate
 			// then the pod is not considered as live
 			p.setStatus(_inactive)
@@ -133,10 +144,13 @@ func (v *ProcessorManager) refreshingProcessors() {
 // startHeatBeatWatcher starts the processor Heartbeat Watcher to listen to the processor bucket and update the processor
 // Heartbeat map. In the processor heartbeat bucket we have the structure key: processor-name, value: processor-heartbeat.
 func (v *ProcessorManager) startHeatBeatWatcher() {
-	watchCh := v.hbWatcher.Watch(v.ctx)
+	watchCh, stopped := v.hbWatcher.Watch(v.ctx)
 	for {
 		select {
-		case <-v.ctx.Done():
+		case <-stopped:
+			// main process exit
+			// close both watcher
+			v.otWatcher.Close()
 			v.hbWatcher.Close()
 			return
 		case value := <-watchCh:
@@ -176,6 +190,7 @@ func (v *ProcessorManager) startHeatBeatWatcher() {
 				} else {
 					v.log.Infow("Deleting", zap.String("key", value.Key()), zap.String(value.Key(), p.String()))
 					p.setStatus(_deleted)
+					p.stopTimeLineWatcher()
 					v.heartbeat.Delete(value.Key())
 				}
 			case store.KVPurge:
diff --git a/pkg/watermark/fetch/processor_to_fetch.go b/pkg/watermark/fetch/processor_to_fetch.go
index a9bc1a0f..b7843950 100644
--- a/pkg/watermark/fetch/processor_to_fetch.go
+++ b/pkg/watermark/fetch/processor_to_fetch.go
@@ -36,6 +36,7 @@ func (s status) String() string {
 // ProcessorToFetch is the smallest unit of entity (from which we fetch data) that does inorder processing or contains inorder data.
 type ProcessorToFetch struct {
 	ctx            context.Context
+	cancel         context.CancelFunc
 	entity         processor.ProcessorEntitier
 	status         status
 	offsetTimeline *OffsetTimeline
@@ -50,8 +51,10 @@ func (p *ProcessorToFetch) String() string {
 
 // NewProcessorToFetch creates ProcessorToFetch.
 func NewProcessorToFetch(ctx context.Context, processor processor.ProcessorEntitier, capacity int, watcher store.WatermarkKVWatcher) *ProcessorToFetch {
+	ctx, cancel := context.WithCancel(ctx)
 	p := &ProcessorToFetch{
 		ctx:            ctx,
+		cancel:         cancel,
 		entity:         processor,
 		status:         _active,
 		offsetTimeline: NewOffsetTimeline(ctx, capacity),
@@ -95,12 +98,18 @@ func (p *ProcessorToFetch) IsDeleted() bool {
 	return p.status == _deleted
 }
 
+func (p *ProcessorToFetch) stopTimeLineWatcher() {
+	p.cancel()
+}
+
 func (p *ProcessorToFetch) startTimeLineWatcher() {
-	watchCh := p.otWatcher.Watch(p.ctx)
+	watchCh, stopped := p.otWatcher.Watch(p.ctx)
+
 	for {
 		select {
-		case <-p.ctx.Done():
-			p.otWatcher.Close()
+		case <-stopped:
+			// no need to close ot watcher here because the ot watcher is shared for the given vertex
+			// the parent ctx will close the ot watcher
 			return
 		case value := <-watchCh:
 			// TODO: why will value will be nil?
@@ -131,6 +140,5 @@ func (p *ProcessorToFetch) startTimeLineWatcher() {
 				// skip
 			}
 		}
-
 	}
 }
diff --git a/pkg/watermark/publish/publisher.go b/pkg/watermark/publish/publisher.go
index 9d079128..446abb8e 100644
--- a/pkg/watermark/publish/publisher.go
+++ b/pkg/watermark/publish/publisher.go
@@ -149,7 +149,7 @@ func (p *publish) GetLatestWatermark() processor.Watermark {
 func (p *publish) publishHeartbeat() {
 	ticker := time.NewTicker(time.Second * time.Duration(p.opts.podHeartbeatRate))
 	defer ticker.Stop()
-	p.log.Infow("Refreshing ActiveProcessors ticker started")
+	p.log.Infow("Publishing Heartbeat ticker started")
 	for {
 		select {
 		case <-p.ctx.Done():
diff --git a/pkg/watermark/store/inmem/kv_watch.go b/pkg/watermark/store/inmem/kv_watch.go
index 97f21fc3..da5a5fc6 100644
--- a/pkg/watermark/store/inmem/kv_watch.go
+++ b/pkg/watermark/store/inmem/kv_watch.go
@@ -37,10 +37,11 @@ func NewInMemWatch(ctx context.Context, pipelineName string, bucketName string,
 }
 
 // Watch watches the key-value store.
-func (k *inMemWatch) Watch(ctx context.Context) <-chan store.WatermarkKVEntry {
+func (k *inMemWatch) Watch(ctx context.Context) (<-chan store.WatermarkKVEntry, <-chan struct{}) {
 	// create a new updates channel and fill in the history
 	var id = util.RandomString(10)
 	var updates = make(chan store.WatermarkKVEntry)
+	var stopped = make(chan struct{})
 
 	// for new updates channel initialization
 	go func() {
@@ -63,6 +64,7 @@ func (k *inMemWatch) Watch(ctx context.Context) <-chan store.WatermarkKVEntry {
 				delete(k.updatesChMap, id)
 				close(updates)
 				k.lock.Unlock()
+				close(stopped)
 				return
 			case value := <-k.kvEntryCh:
 				k.log.Debug(value.Key(), value.Value(), value.Operation())
@@ -75,7 +77,7 @@ func (k *inMemWatch) Watch(ctx context.Context) <-chan store.WatermarkKVEntry {
 			}
 		}
 	}()
-	return updates
+	return updates, stopped
 }
 
 // GetKVName returns the KV store (bucket) name.
diff --git a/pkg/watermark/store/interface.go b/pkg/watermark/store/interface.go
index 6d0968d0..4b08f48a 100644
--- a/pkg/watermark/store/interface.go
+++ b/pkg/watermark/store/interface.go
@@ -64,7 +64,8 @@ type WatermarkKVEntry interface {
 
 // WatermarkKVWatcher watches the KV bucket for watermark progression.
 type WatermarkKVWatcher interface {
-	Watch(context.Context) <-chan WatermarkKVEntry
+	// Watch starts the watermark kv watcher and returns a kv updates channel and a watcher stopped channel.
+	Watch(context.Context) (<-chan WatermarkKVEntry, <-chan struct{})
 	GetKVName() string
 	Close()
 }
diff --git a/pkg/watermark/store/jetstream/kv_watch.go b/pkg/watermark/store/jetstream/kv_watch.go
index 2e2d5b81..fc0e175e 100644
--- a/pkg/watermark/store/jetstream/kv_watch.go
+++ b/pkg/watermark/store/jetstream/kv_watch.go
@@ -96,7 +96,7 @@ func (k kvEntry) Operation() store.KVWatchOp {
 }
 
 // Watch watches the key-value store (aka bucket).
-func (k *jetStreamWatch) Watch(ctx context.Context) <-chan store.WatermarkKVEntry {
+func (k *jetStreamWatch) Watch(ctx context.Context) (<-chan store.WatermarkKVEntry, <-chan struct{}) {
 	kvWatcher, err := k.kv.WatchAll()
 	for err != nil {
 		k.log.Errorw("WatchAll failed", zap.String("watcher", k.GetKVName()), zap.Error(err))
@@ -105,6 +105,7 @@ func (k *jetStreamWatch) Watch(ctx context.Context) <-chan store.WatermarkKVEntr
 	}
 
 	var updates = make(chan store.WatermarkKVEntry)
+	var stopped = make(chan struct{})
 	go func() {
 		for {
 			select {
@@ -118,6 +119,7 @@ func (k *jetStreamWatch) Watch(ctx context.Context) <-chan store.WatermarkKVEntr
 					k.log.Infow("WatchAll successfully stopped", zap.String("watcher", k.GetKVName()))
 				}
 				close(updates)
+				close(stopped)
 				return
 			case value := <-kvWatcher.Updates():
 				// if channel is closed, nil could come in
@@ -144,7 +146,7 @@ func (k *jetStreamWatch) Watch(ctx context.Context) <-chan store.WatermarkKVEntr
 			}
 		}
 	}()
-	return updates
+	return updates, stopped
 }
 
 // GetKVName returns the KV store (bucket) name.
diff --git a/pkg/watermark/store/noop/kv_watch.go b/pkg/watermark/store/noop/kv_watch.go
index e71b7bb2..228d074c 100644
--- a/pkg/watermark/store/noop/kv_watch.go
+++ b/pkg/watermark/store/noop/kv_watch.go
@@ -16,9 +16,15 @@ func NewKVOpWatch() store.WatermarkKVWatcher {
 }
 
 // Watch returns a blocking channel.
-func (no noOpWatch) Watch(ctx context.Context) <-chan store.WatermarkKVEntry {
+func (no noOpWatch) Watch(ctx context.Context) (<-chan store.WatermarkKVEntry, <-chan struct{}) {
 	retChan := make(chan store.WatermarkKVEntry)
-	return retChan
+	stopped := make(chan struct{})
+	go func() {
+		<-ctx.Done()
+		close(retChan)
+		close(stopped)
+	}()
+	return retChan, stopped
 }
 
 func (no noOpWatch) GetKVName() string {


Test patch:
diff --git a/pkg/watermark/fetch/processor_manager_inmem_test.go b/pkg/watermark/fetch/processor_manager_inmem_test.go
index a5679748..b00ca13b 100644
--- a/pkg/watermark/fetch/processor_manager_inmem_test.go
+++ b/pkg/watermark/fetch/processor_manager_inmem_test.go
@@ -50,10 +50,8 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 
 	hbWatcher, err := inmem.NewInMemWatch(ctx, "testFetch", keyspace+"_PROCESSORS", hbWatcherCh)
 	assert.NoError(t, err)
-	defer hbWatcher.Close()
 	otWatcher, err := inmem.NewInMemWatch(ctx, "testFetch", keyspace+"_OT", otWatcherCh)
 	assert.NoError(t, err)
-	defer otWatcher.Close()
 	var testVertex = NewProcessorManager(ctx, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher), WithPodHeartbeatRate(1), WithRefreshingProcessorsRate(1), WithSeparateOTBuckets(false))
 	var testBuffer = NewEdgeFetcher(ctx, "testBuffer", testVertex).(*edgeFetcher)
 
@@ -122,11 +120,11 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 	assert.Equal(t, 2, len(allProcessors))
 	assert.True(t, allProcessors["p1"].IsDeleted())
 	assert.True(t, allProcessors["p2"].IsActive())
+	// "p1" should be deleted after this GetWatermark offset=101
+	// because "p1" offsetTimeline's head offset=100, which is < inputOffset 103
 	_ = testBuffer.GetWatermark(isb.SimpleOffset(func() string { return strconv.FormatInt(testOffset+3, 10) }))
 	allProcessors = testBuffer.processorManager.GetAllProcessors()
-	// we don't delete inactive processors from processor manager, so the length should still be 2
-	assert.Equal(t, 2, len(allProcessors))
-	assert.True(t, allProcessors["p1"].IsDeleted())
+	assert.Equal(t, 1, len(allProcessors))
 	assert.True(t, allProcessors["p2"].IsActive())
 
 	time.Sleep(time.Second)
@@ -159,9 +157,27 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 	}
 	assert.True(t, allProcessors["p1"].IsActive())
 	assert.True(t, allProcessors["p2"].IsActive())
-	// "p1" should still have the same offsetTimeline
+	// "p1" has been deleted from vertex.Processors
+	// so "p1" will be considered as a new processors and a new offsetTimeline watcher for "p1" will be created
 	_ = testBuffer.GetWatermark(isb.SimpleOffset(func() string { return strconv.FormatInt(testOffset+1, 10) }))
 	p1 := testBuffer.processorManager.GetProcessor("p1")
+
+	assert.NotNil(t, p1)
+	assert.True(t, p1.IsActive())
+	assert.NotNil(t, p1.offsetTimeline)
+	// wait till the offsetTimeline has been populated
+	newP1head := p1.offsetTimeline.GetHeadOffset()
+	for newP1head == -1 {
+		select {
+		case <-ctx.Done():
+			t.Fatalf("expected head offset to not be equal to -1, %s", ctx.Err())
+		default:
+			time.Sleep(1 * time.Millisecond)
+			newP1head = p1.offsetTimeline.GetHeadOffset()
+		}
+	}
+	// because we keep the kv updates history in the in mem watermark
+	// the new watcher will read all the history data to create this new offsetTimeline
 	assert.Equal(t, int64(100), p1.offsetTimeline.GetHeadOffset())
 
 	// publish a new watermark 101
diff --git a/pkg/watermark/fetch/processor_manager_test.go b/pkg/watermark/fetch/processor_manager_test.go
index 0e842951..451e3a26 100644
--- a/pkg/watermark/fetch/processor_manager_test.go
+++ b/pkg/watermark/fetch/processor_manager_test.go
@@ -95,10 +95,8 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 	// create watchers for heartbeat and offset timeline
 	hbWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", keyspace+"_PROCESSORS", defaultJetStreamClient)
 	assert.NoError(t, err)
-	defer hbWatcher.Close()
 	otWatcher, err := jetstream.NewKVJetStreamKVWatch(ctx, "testFetch", keyspace+"_OT", defaultJetStreamClient)
 	assert.NoError(t, err)
-	defer otWatcher.Close()
 	var testVertex = NewProcessorManager(ctx, store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher), WithPodHeartbeatRate(1), WithRefreshingProcessorsRate(1), WithSeparateOTBuckets(false))
 	var testBuffer = NewEdgeFetcher(ctx, "testBuffer", testVertex).(*edgeFetcher)
 
@@ -166,12 +164,12 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 	assert.Equal(t, 2, len(allProcessors))
 	assert.True(t, allProcessors["p1"].IsDeleted())
 	assert.True(t, allProcessors["p2"].IsActive())
+	// "p1" should be deleted after this GetWatermark offset=101
+	// because "p1" offsetTimeline's head offset=100, which is < inputOffset 103
 	_ = testBuffer.GetWatermark(isb.SimpleOffset(func() string { return strconv.FormatInt(testOffset+3, 10) }))
 	allProcessors = testBuffer.processorManager.GetAllProcessors()
-	// we don't delete inactive processors, so the length should still be 2
-	assert.Equal(t, 2, len(allProcessors))
+	assert.Equal(t, 1, len(allProcessors))
 	assert.True(t, allProcessors["p2"].IsActive())
-	assert.False(t, allProcessors["p1"].IsActive())
 
 	time.Sleep(time.Second)
 	// resume after one second
@@ -203,12 +201,25 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 
 	assert.True(t, allProcessors["p1"].IsActive())
 	assert.True(t, allProcessors["p2"].IsActive())
-	// "p1" status is back to active, the offset timeline should be preserved
+	// "p1" has been deleted from vertex.Processors
+	// so "p1" will be considered as a new processors and a new offsetTimeline watcher for "p1" will be created
 	_ = testBuffer.GetWatermark(isb.SimpleOffset(func() string { return strconv.FormatInt(testOffset+1, 10) }))
 	p1 := testBuffer.processorManager.GetProcessor("p1")
 	assert.NotNil(t, p1)
 	assert.True(t, p1.IsActive())
 	assert.NotNil(t, p1.offsetTimeline)
+	// wait till the offsetTimeline has been populated
+	newP1head := p1.offsetTimeline.GetHeadOffset()
+	for newP1head == -1 {
+		select {
+		case <-ctx.Done():
+			t.Fatalf("expected head offset to not be equal to -1, %s", ctx.Err())
+		default:
+			time.Sleep(1 * time.Millisecond)
+			newP1head = p1.offsetTimeline.GetHeadOffset()
+		}
+	}
+	// because the bucket hasn't been cleaned up, the new watcher will read all the history data to create this new offsetTimeline
 	assert.Equal(t, int64(100), p1.offsetTimeline.GetHeadOffset())
 
 	// publish a new watermark 101
@@ -259,12 +270,6 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 
 	wg.Wait()
 	cancel()
-
-	// wait for all the ctx.Done() is completed
-	// we don't have wg on the Watch functions go routine..
-	// TODO: opened an issue https://github.com/numaproj/numaflow/issues/224
-	//   fix tests after we revisit the watermark store and watcher closing logic
-	time.Sleep(3 * time.Second)
 }
 
 func TestFetcherWithSeparateOTBucket(t *testing.T) {
diff --git a/pkg/window/keyed/keyed_test.go b/pkg/window/keyed/keyed_test.go
index 8785ccca..7f496cd0 100644
--- a/pkg/window/keyed/keyed_test.go
+++ b/pkg/window/keyed/keyed_test.go
@@ -1,6 +1,7 @@
 package keyed
 
 import (
+	"sort"
 	"testing"
 	"time"
 
@@ -45,6 +46,7 @@ func TestKeyedWindow_AddKey(t *testing.T) {
 				kw.AddKey(k)
 			}
 			kw.AddKey(tt.input)
+			assert.Equal(t, len(tt.expectedKeys), len(kw.Keys))
 			for k := range tt.expectedKeys {
 				_, ok := kw.Keys[k]
 				assert.True(t, ok)
@@ -73,7 +75,7 @@ func TestKeyedWindow_Partitions(t *testing.T) {
 		{
 			name: "with_some_existing_keys",
 			given: &KeyedWindow{
-				Keys: map[string]string{"key2": "key2", "key3": "key3"},
+				Keys: map[string]string{"key2": "key2", "key3": "key3", "key4": "key4"},
 			},
 			expected: []partition.ID{
 				{
@@ -86,6 +88,11 @@ func TestKeyedWindow_Partitions(t *testing.T) {
 					Start: time.Unix(60, 0),
 					End:   time.Unix(120, 0),
 				},
+				{
+					Key:   "key4",
+					Start: time.Unix(60, 0),
+					End:   time.Unix(120, 0),
+				},
 			},
 		},
 	}
@@ -94,6 +101,11 @@ func TestKeyedWindow_Partitions(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			kw.Keys = tt.given.Keys
 			ret := kw.Partitions()
+			// the kw.Keys is a map so the order of the output is random
+			// use sort to sort the ret array by key
+			sort.Slice(ret, func(i int, j int) bool {
+				return ret[i].Key < ret[j].Key
+			})
 			for idx, s := range tt.expected {
 				assert.EqualValues(t, ret[idx], s)
 			}