Added a resettable signal

Signed-off-by: Jakob Haahr Taankvist <jht@uber.com>

Author: jakobht

service/sharddistributor/client/spectatorclient/sync/resettable_signal.go

@@ -0,0 +1,87 @@
+package sync
+
+import (
+	"context"
+	"errors"
+	"sync"
+)
+
+// ErrReset is returned by Wait() when Reset() is called while goroutines are waiting
+var ErrReset = errors.New("signal was reset")
+
+// resettableSignal is a synchronization primitive that allows waiting for a one-time
+// signal with context support, and can be reset to wait for a new signal.
+// Similar to sync.WaitGroup but for single completion events with reset capability.
+type resettableSignal struct {
+	mu      sync.Mutex
+	doneCh  chan struct{}
+	resetCh chan struct{}
+	done    bool
+}
+
+// ResettableSignal must be created via NewResettableSignal(). Zero value is invalid.
+type ResettableSignal = *resettableSignal
+
+// NewResettableSignal creates a new resettable signal in waiting state
+func NewResettableSignal() *resettableSignal {
+	return &resettableSignal{
+		doneCh:  make(chan struct{}),
+		resetCh: make(chan struct{}),
+	}
+}
+
+// Done signals that the event has completed. Safe to call multiple times (idempotent).
+func (s *resettableSignal) Done() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if !s.done {
+		s.done = true
+		close(s.doneCh)
+	}
+}
+
+// Wait blocks until either Done() is called, the context is cancelled, or Reset() is called.
+// Returns:
+//   - nil if Done() was called
+//   - ctx.Err() if context was cancelled
+//   - ErrReset if Reset() was called while waiting
+func (s *resettableSignal) Wait(ctx context.Context) error {
+	s.mu.Lock()
+	doneCh := s.doneCh
+	resetCh := s.resetCh
+	done := s.done
+	s.mu.Unlock()
+
+	// Fast path: already done
+	if done {
+		return nil
+	}
+
+	select {
+	case <-doneCh:
+		return nil
+	case <-resetCh:
+		return ErrReset
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+// Reset resets the signal to waiting state. Any goroutines currently blocked in Wait()
+// will immediately be unblocked with ErrReset.
+func (s *resettableSignal) Reset() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// Close reset channel to unblock any waiters (they'll get ErrReset)
+	// Only close if not already done (to avoid closing a closed channel)
+	if !s.done {
+		close(s.resetCh)
+	}
+
+	// Create new channels and reset done flag
+	s.doneCh = make(chan struct{})
+	s.resetCh = make(chan struct{})
+	s.done = false
+}

service/sharddistributor/client/spectatorclient/sync/resettable_signal_test.go

@@ -0,0 +1,153 @@
+package sync
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestResettableSignal_DoneWait(t *testing.T) {
+	signal := NewResettableSignal()
+
+	signal.Done()
+
+	ctx := context.Background()
+	err := signal.Wait(ctx)
+	assert.NoError(t, err)
+}
+
+func TestResettableSignal_WaitDone(t *testing.T) {
+	signal := NewResettableSignal()
+
+	done := make(chan error)
+	go func() {
+		done <- signal.Wait(context.Background())
+	}()
+
+	// Give goroutine time to start waiting
+	time.Sleep(10 * time.Millisecond)
+
+	signal.Done()
+
+	select {
+	case err := <-done:
+		assert.NoError(t, err)
+	case <-time.After(1 * time.Second):
+		t.Fatal("Wait did not complete after Done")
+	}
+}
+
+func TestResettableSignal_ContextCancellation(t *testing.T) {
+	signal := NewResettableSignal()
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	done := make(chan error)
+	go func() {
+		done <- signal.Wait(ctx)
+	}()
+
+	time.Sleep(10 * time.Millisecond)
+
+	cancel()
+
+	select {
+	case err := <-done:
+		assert.Error(t, err)
+		assert.Equal(t, context.Canceled, err)
+	case <-time.After(1 * time.Second):
+		t.Fatal("Wait did not complete after context cancellation")
+	}
+}
+
+func TestResettableSignal_ResetWhileWaiting(t *testing.T) {
+	signal := NewResettableSignal()
+
+	done := make(chan error)
+	go func() {
+		done <- signal.Wait(context.Background())
+	}()
+
+	time.Sleep(10 * time.Millisecond)
+
+	signal.Reset()
+
+	select {
+	case err := <-done:
+		assert.Error(t, err)
+		assert.True(t, errors.Is(err, ErrReset), "expected ErrReset, got %v", err)
+	case <-time.After(1 * time.Second):
+		t.Fatal("Wait did not complete after Reset")
+	}
+}
+
+func TestResettableSignal_MultipleWaitersReset(t *testing.T) {
+	signal := NewResettableSignal()
+
+	const numWaiters = 5
+	results := make([]chan error, numWaiters)
+	for i := 0; i < numWaiters; i++ {
+		results[i] = make(chan error)
+		go func() {
+			results[i] <- signal.Wait(context.Background())
+		}()
+	}
+
+	// Give goroutines time to start waiting
+	time.Sleep(10 * time.Millisecond)
+
+	signal.Reset()
+
+	for i, ch := range results {
+		select {
+		case err := <-ch:
+			assert.Error(t, err, "waiter %d should get error", i)
+			assert.True(t, errors.Is(err, ErrReset), "waiter %d: expected ErrReset, got %v", i, err)
+		case <-time.After(1 * time.Second):
+			t.Fatalf("Waiter %d did not complete after Reset", i)
+		}
+	}
+}
+
+func TestResettableSignal_ResetAfterDone(t *testing.T) {
+	signal := NewResettableSignal()
+
+	// Terminate and reset the signal
+	signal.Done()
+	signal.Reset()
+
+	// Now signal should be waiting again
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	err := signal.Wait(ctx)
+	assert.Error(t, err)
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+}
+
+func TestResettableSignal_ResetThenDoneThenWait(t *testing.T) {
+	signal := NewResettableSignal()
+
+	// Do a full cycle
+	signal.Done()
+	signal.Reset()
+	signal.Done()
+
+	err := signal.Wait(context.Background())
+	assert.NoError(t, err)
+}
+
+func TestResettableSignal_IdempotentDone(t *testing.T) {
+	signal := NewResettableSignal()
+
+	// Call Done multiple times
+	signal.Done()
+	signal.Done()
+	signal.Done()
+
+	err := signal.Wait(context.Background())
+	assert.NoError(t, err)
+}