Add SequenceTracker for efficiently tracking insertions.

To be more exact, it can:
- Insert a number. Complexity: Amortized O(1).
- Check if all numbers upto a given number are inserted. Complexity:
  O(1).
- Efficiently wait until all numbers inserted up to a given number.

Author: walidghallab

cluster-autoscaler/simulator/clustersnapshot/predicate/sequence_tracker.go

@@ -0,0 +1,76 @@
+package predicate
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+)
+
+// sequenceTracker is a helper structure that tracks a sequence of integers and allows waiting until all
+// non-negative integers up to a certain value have been inserted.
+type sequenceTracker struct {
+	received          map[int]bool
+	currentContiguous int                   // All non-negative numbers until this is guaranteed to be inserted.
+	broadcaster       map[int]chan struct{} // Used to efficiently signal when a all the non-negative numbers until the number in the key are inserted.
+	mutex             sync.Mutex
+}
+
+func newSequenceTracker() *sequenceTracker {
+	return &sequenceTracker{
+		received:          make(map[int]bool),
+		currentContiguous: -1,
+		broadcaster:       make(map[int]chan struct{}),
+	}
+}
+
+// Complexity: Amortized O(1).
+func (s *sequenceTracker) Insert(i int) {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	s.received[i] = true
+
+	// Slide the window forward
+	for s.received[s.currentContiguous+1] {
+		s.currentContiguous++
+		if ch, ok := s.broadcaster[s.currentContiguous]; ok {
+			close(ch)
+			delete(s.broadcaster, s.currentContiguous)
+		}
+		delete(s.received, s.currentContiguous)
+	}
+}
+
+// Complexity: O(1).
+func (s *sequenceTracker) allUpToInserted(n int) bool {
+	return s.currentContiguous >= n
+}
+
+// WaitUntilAllUpToInserted waits until all non-negative integers up to n have been inserted into the tracker.
+// If all numbers are already inserted, it returns immediately.
+// If not, it waits until the condition is met, the context is cancelled, or the timeout is reached.
+func (s *sequenceTracker) WaitUntilAllUpToInserted(ctx context.Context, n int, timeout time.Duration) error {
+	ctx, _ = context.WithTimeout(ctx, timeout)
+	s.mutex.Lock()
+	// After entering the mutex, there are only 2 cases:
+	// - Either all up to the given number is inserted. In such case we exit quickly.
+	// - Not all up to the given number is inserted. In such case, we create a channel if it doesn't exist, and rely on Insert to close it when all below are inserted.
+	if s.allUpToInserted(n) {
+		s.mutex.Unlock()
+		return nil
+	}
+	ch, ok := s.broadcaster[n]
+	if !ok {
+		ch = make(chan struct{})
+		s.broadcaster[n] = ch
+	}
+	s.mutex.Unlock()
+
+	select {
+	case <-ctx.Done():
+		return fmt.Errorf("context is done: %v", ctx.Err())
+	case <-ch:
+		return nil
+	}
+}

cluster-autoscaler/simulator/clustersnapshot/predicate/sequence_tracker_test.go

@@ -0,0 +1,204 @@
+package predicate
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	testWaitTimeout = 3 * time.Second
+)
+
+func TestNewSequenceTracker(t *testing.T) {
+	tracker := newSequenceTracker()
+	assert.NotNil(t, tracker)
+	assert.Equal(t, -1, tracker.currentContiguous)
+	assert.NotNil(t, tracker.received)
+	assert.NotNil(t, tracker.broadcaster)
+}
+
+func TestAllUpToInserted(t *testing.T) {
+	type testStep struct {
+		numToInsert int
+		assertions  map[int]bool
+	}
+
+	testCases := []struct {
+		name       string
+		insertions []testStep
+	}{
+		{
+			name: "multiple insertions",
+			insertions: []testStep{
+				{
+					numToInsert: 0,
+					assertions: map[int]bool{
+						0: true,
+						1: false,
+						2: false,
+						3: false,
+						4: false,
+						5: false,
+					},
+				},
+				{
+					numToInsert: 1,
+					assertions: map[int]bool{
+						0: true,
+						1: true,
+						2: false,
+						3: false,
+						4: false,
+						5: false,
+					},
+				},
+				{
+					numToInsert: 2,
+					assertions: map[int]bool{
+						0: true,
+						1: true,
+						2: true,
+						3: false,
+						4: false,
+						5: false,
+					},
+				},
+				{
+					numToInsert: 5,
+					assertions: map[int]bool{
+						0: true,
+						1: true,
+						2: true,
+						3: false,
+						4: false,
+						5: false,
+					},
+				},
+			},
+		},
+		{
+			name: "multiple insertions - reverse order",
+			insertions: []testStep{
+				{
+					numToInsert: 5,
+					assertions: map[int]bool{
+						0: false,
+						1: false,
+						2: false,
+						3: false,
+						4: false,
+						5: false,
+					},
+				},
+				{
+					numToInsert: 2,
+					assertions: map[int]bool{
+						0: false,
+						1: false,
+						2: false,
+						3: false,
+						4: false,
+						5: false,
+					},
+				},
+				{
+					numToInsert: 1,
+					assertions: map[int]bool{
+						0: false,
+						1: false,
+						2: false,
+						3: false,
+						4: false,
+						5: false,
+					},
+				},
+				{
+					numToInsert: 0,
+					assertions: map[int]bool{
+						0: true,
+						1: true,
+						2: true,
+						3: false,
+						4: false,
+						5: false,
+					},
+				},
+			},
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			tracker := newSequenceTracker()
+			for _, insertion := range tc.insertions {
+				tracker.Insert(insertion.numToInsert)
+				for i, expected := range insertion.assertions {
+					assert.Equal(t, expected, tracker.allUpToInserted(i), fmt.Sprintf("allUpToInserted(%d) should be %v after inserting %d", i, expected, insertion.numToInsert))
+				}
+			}
+		})
+	}
+}
+
+func TestWaitUntilAllUpToInserted(t *testing.T) {
+	tracker := newSequenceTracker()
+	var wg sync.WaitGroup
+	wg.Add(1)
+	waitFinished := false
+	go func() {
+		defer wg.Done()
+		ctx := context.Background()
+		err := tracker.WaitUntilAllUpToInserted(ctx, 2, testWaitTimeout)
+		assert.NoError(t, err)
+		waitFinished = true
+	}()
+
+	assert.False(t, waitFinished)
+
+	tracker.Insert(0)
+	assert.False(t, waitFinished)
+
+	tracker.Insert(1)
+	assert.False(t, waitFinished)
+
+	tracker.Insert(2)
+	wg.Wait()
+	assert.True(t, waitFinished)
+}
+
+func TestWaitUntilAllUpToInserted_AlreadyInserted(t *testing.T) {
+	tracker := newSequenceTracker()
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	tracker.Insert(0)
+	tracker.Insert(1)
+	tracker.Insert(2)
+
+	go func() {
+		defer wg.Done()
+		ctx := context.Background()
+		err := tracker.WaitUntilAllUpToInserted(ctx, 2, testWaitTimeout)
+		assert.NoError(t, err)
+	}()
+
+	wg.Wait()
+}
+
+func TestWaitUntilAllUpToInserted_WithContext(t *testing.T) {
+	tracker := newSequenceTracker()
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		err := tracker.WaitUntilAllUpToInserted(ctx, 0, testWaitTimeout)
+		assert.Error(t, err)
+	}()
+	cancel()
+	wg.Wait()
+}