tests done

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

Author: jakobht

service/sharddistributor/client/spectatorclient/client.go

@@ -6,16 +6,20 @@ import (
 
 	"github.com/uber-go/tally"
 	"go.uber.org/fx"
-	csync "github.com/uber/cadence/service/sharddistributor/client/spectatorclient/sync"
 
 	"github.com/uber/cadence/client/sharddistributor"
 	"github.com/uber/cadence/common/clock"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/service/sharddistributor/client/clientcommon"
+	csync "github.com/uber/cadence/service/sharddistributor/client/spectatorclient/sync"
 )
 
 //go:generate mockgen -package $GOPACKAGE -source $GOFILE -destination interface_mock.go . Spectator
 
+// EnabledFunc is a function that returns true if the spectator is enabled
+// This is used to disable the spectator in the case of a migration
+type EnabledFunc func() bool
+
 type Spectators struct {
 	spectators map[string]Spectator
 }
@@ -73,9 +77,7 @@ type Params struct {
 	Config       clientcommon.Config
 	TimeSource   clock.TimeSource
 
-	// Enabled is a function that returns true if the spectator is enabled
-	// This is used to disable the spectator in the case of a migration
-	Enabled func() bool `optional:"true"`
+	Enabled EnabledFunc `optional:"true"`
 }
 
 // NewSpectatorWithNamespace creates a spectator for a specific namespace
@@ -109,14 +111,14 @@ func newSpectatorWithConfig(params Params, namespaceConfig *clientcommon.Namespa
 	}
 
 	impl := &spectatorImpl{
-		namespace:    namespaceConfig.Namespace,
-		config:       *namespaceConfig,
-		client:       params.Client,
-		logger:       params.Logger,
-		scope:        params.MetricsScope,
-		timeSource:   params.TimeSource,
+		namespace:        namespaceConfig.Namespace,
+		config:           *namespaceConfig,
+		client:           params.Client,
+		logger:           params.Logger,
+		scope:            params.MetricsScope,
+		timeSource:       params.TimeSource,
 		firstStateSignal: csync.NewResettableSignal(),
-		enabled:      enabled,
+		enabled:          enabled,
 	}
 
 	return impl, nil

service/sharddistributor/client/spectatorclient/clientimpl.go

@@ -15,6 +15,7 @@ import (
 	"github.com/uber/cadence/common/log/tag"
 	"github.com/uber/cadence/common/types"
 	"github.com/uber/cadence/service/sharddistributor/client/clientcommon"
+	csync "github.com/uber/cadence/service/sharddistributor/client/spectatorclient/sync"
 )
 
 const (
@@ -30,6 +31,7 @@ type ShardOwner struct {
 
 type spectatorImpl struct {
 	namespace  string
+	enabled    EnabledFunc
 	config     clientcommon.NamespaceConfig
 	client     sharddistributor.Client
 	scope      tally.Scope
@@ -73,27 +75,37 @@ func (s *spectatorImpl) Stop() {
 }
 
 func (s *spectatorImpl) watchLoop() {
+	defer s.logger.Info("Shutting down, stopping watch loop", tag.ShardNamespace(s.namespace))
+
 	s.logger.Info("Starting watch loop for namespace", tag.ShardNamespace(s.namespace))
 
 	for {
 		if s.ctx.Err() != nil {
-			s.logger.Info("Shutting down, stopping watch loop", tag.ShardNamespace(s.namespace))
 			return
 		}
 
+		if !s.enabled() {
+			// If spectator is disabled, sleep for a second and continue
+			s.firstStateSignal.Reset()
+
+			err := s.timeSource.SleepWithContext(s.ctx, backoff.JitDuration(streamRetryInterval, streamRetryJitterCoeff))
+			if err != nil {
+				return
+			}
+			continue
+		}
+
 		// Create new stream
 		stream, err := s.client.WatchNamespaceState(s.ctx, &types.WatchNamespaceStateRequest{
 			Namespace: s.namespace,
 		})
 		if err != nil {
 			if s.ctx.Err() != nil {
-				s.logger.Info("Shutting down during stream creation, exiting watch loop", tag.ShardNamespace(s.namespace))
 				return
 			}
 
 			s.logger.Error("Failed to create stream, retrying", tag.Error(err), tag.ShardNamespace(s.namespace))
 			if err := s.timeSource.SleepWithContext(s.ctx, backoff.JitDuration(streamRetryInterval, streamRetryJitterCoeff)); err != nil {
-				s.logger.Info("Shutting down waiting to retry stream creation, exiting watch loop", tag.ShardNamespace(s.namespace))
 				return // Context cancelled during sleep
 			}
 			continue
@@ -102,7 +114,6 @@ func (s *spectatorImpl) watchLoop() {
 		s.receiveLoop(stream)
 
 		if s.ctx.Err() != nil {
-			s.logger.Info("Shutting down, exiting watch loop", tag.ShardNamespace(s.namespace))
 			return
 		}
 
@@ -122,6 +133,12 @@ func (s *spectatorImpl) receiveLoop(stream sharddistributor.WatchNamespaceStateC
 	}()
 
 	for {
+		if !s.enabled() {
+			// the loop was disabled, exit
+			s.firstStateSignal.Reset()
+			return
+		}
+
 		response, err := stream.Recv()
 		if err != nil {
 			if s.ctx.Err() != nil {
@@ -176,7 +193,9 @@ func (s *spectatorImpl) handleResponse(response *types.WatchNamespaceStateRespon
 // If not found in cache, it falls back to querying the shard distributor directly.
 func (s *spectatorImpl) GetShardOwner(ctx context.Context, shardKey string) (*ShardOwner, error) {
 	// Wait for first state to be received to avoid flooding shard distributor on startup
-	s.firstStateSignal.Wait(ctx)
+	if err := s.firstStateSignal.Wait(ctx); err != nil {
+		return nil, fmt.Errorf("wait for first state: %w", err)
+	}
 
 	// Check cache first
 	s.stateMu.RLock()

service/sharddistributor/client/spectatorclient/clientimpl_test.go

@@ -3,6 +3,7 @@ package spectatorclient
 import (
 	"context"
 	"errors"
+	"sync"
 	"testing"
 	"time"
 
@@ -16,6 +17,7 @@ import (
 	"github.com/uber/cadence/common/clock"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/types"
+	csync "github.com/uber/cadence/service/sharddistributor/client/spectatorclient/sync"
 )
 
 func TestWatchLoopBasicFlow(t *testing.T) {
@@ -26,12 +28,13 @@ func TestWatchLoopBasicFlow(t *testing.T) {
 	mockStream := sharddistributor.NewMockWatchNamespaceStateClient(ctrl)
 
 	spectator := &spectatorImpl{
-		namespace:    "test-ns",
-		client:       mockClient,
-		logger:       log.NewNoop(),
-		scope:        tally.NoopScope,
-		timeSource:   clock.NewRealTimeSource(),
-		firstStateCh: make(chan struct{}),
+		namespace:        "test-ns",
+		client:           mockClient,
+		logger:           log.NewNoop(),
+		scope:            tally.NoopScope,
+		timeSource:       clock.NewRealTimeSource(),
+		firstStateSignal: csync.NewResettableSignal(),
+		enabled:          func() bool { return true },
 	}
 
 	// Expect stream creation
@@ -70,7 +73,7 @@ func TestWatchLoopBasicFlow(t *testing.T) {
 	defer spectator.Stop()
 
 	// Wait for first state
-	<-spectator.firstStateCh
+	require.NoError(t, spectator.firstStateSignal.Wait(context.Background()))
 
 	// Query shard owner
 	owner, err := spectator.GetShardOwner(context.Background(), "shard-1")
@@ -91,12 +94,13 @@ func TestGetShardOwner_CacheMiss_FallbackToRPC(t *testing.T) {
 	mockStream := sharddistributor.NewMockWatchNamespaceStateClient(ctrl)
 
 	spectator := &spectatorImpl{
-		namespace:    "test-ns",
-		client:       mockClient,
-		logger:       log.NewNoop(),
-		scope:        tally.NoopScope,
-		timeSource:   clock.NewRealTimeSource(),
-		firstStateCh: make(chan struct{}),
+		namespace:        "test-ns",
+		client:           mockClient,
+		logger:           log.NewNoop(),
+		scope:            tally.NoopScope,
+		timeSource:       clock.NewRealTimeSource(),
+		firstStateSignal: csync.NewResettableSignal(),
+		enabled:          func() bool { return true },
 	}
 
 	// Setup stream
@@ -142,7 +146,7 @@ func TestGetShardOwner_CacheMiss_FallbackToRPC(t *testing.T) {
 	spectator.Start(context.Background())
 	defer spectator.Stop()
 
-	<-spectator.firstStateCh
+	require.NoError(t, spectator.firstStateSignal.Wait(context.Background()))
 
 	// Cache hit
 	owner, err := spectator.GetShardOwner(context.Background(), "shard-1")
@@ -166,12 +170,13 @@ func TestStreamReconnection(t *testing.T) {
 	mockTimeSource := clock.NewMockedTimeSource()
 
 	spectator := &spectatorImpl{
-		namespace:    "test-ns",
-		client:       mockClient,
-		logger:       log.NewNoop(),
-		scope:        tally.NoopScope,
-		timeSource:   mockTimeSource,
-		firstStateCh: make(chan struct{}),
+		namespace:        "test-ns",
+		client:           mockClient,
+		logger:           log.NewNoop(),
+		scope:            tally.NoopScope,
+		timeSource:       mockTimeSource,
+		firstStateSignal: csync.NewResettableSignal(),
+		enabled:          func() bool { return true },
 	}
 
 	// First stream fails immediately
@@ -208,7 +213,7 @@ func TestStreamReconnection(t *testing.T) {
 	mockTimeSource.BlockUntil(1) // Wait for 1 goroutine to be blocked in Sleep
 	mockTimeSource.Advance(2 * time.Second)
 
-	<-spectator.firstStateCh
+	require.NoError(t, spectator.firstStateSignal.Wait(context.Background()))
 }
 
 func TestGetShardOwner_TimeoutBeforeFirstState(t *testing.T) {
@@ -218,12 +223,13 @@ func TestGetShardOwner_TimeoutBeforeFirstState(t *testing.T) {
 	mockClient := sharddistributor.NewMockClient(ctrl)
 
 	spectator := &spectatorImpl{
-		namespace:    "test-ns",
-		client:       mockClient,
-		logger:       log.NewNoop(),
-		scope:        tally.NoopScope,
-		timeSource:   clock.NewRealTimeSource(),
-		firstStateCh: make(chan struct{}),
+		namespace:        "test-ns",
+		client:           mockClient,
+		logger:           log.NewNoop(),
+		scope:            tally.NoopScope,
+		timeSource:       clock.NewRealTimeSource(),
+		firstStateSignal: csync.NewResettableSignal(),
+		enabled:          func() bool { return true },
 	}
 
 	// Create a context with a short timeout
@@ -234,5 +240,49 @@ func TestGetShardOwner_TimeoutBeforeFirstState(t *testing.T) {
 	// Should timeout and return an error
 	_, err := spectator.GetShardOwner(ctx, "shard-1")
 	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "context cancelled while waiting for first state")
+	assert.Contains(t, err.Error(), "wait for first state")
+}
+
+func TestWatchLoopDisabled(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	stateSignal := csync.NewResettableSignal()
+	timeSource := clock.NewMockedTimeSource()
+
+	spectator := &spectatorImpl{
+		firstStateSignal: stateSignal,
+		timeSource:       timeSource,
+		logger:           log.NewNoop(),
+		enabled:          func() bool { return false },
+	}
+
+	spectator.Start(context.Background())
+	defer spectator.Stop()
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		err := stateSignal.Wait(context.Background())
+		assert.Error(t, err)
+
+		// Second wait might return ErrReset (if it observes the second reset)
+		// or nil (if Stop() is called first). Both are acceptable.
+		_ = stateSignal.Wait(context.Background())
+	}()
+
+	// First sleep should reset the signal
+	timeSource.BlockUntil(1)
+	timeSource.Advance(1200 * time.Millisecond)
+
+	// Second sleep should reset the signal
+	timeSource.BlockUntil(1)
+
+	// Ensure the loop is exited
+	timeSource.Advance(1200 * time.Millisecond)
+
+	// Stop the spectator to unblock any waiting goroutines
+	spectator.Stop()
+
+	wg.Wait()
 }

service/sharddistributor/client/spectatorclient/peer_chooser_test.go

@@ -105,6 +105,8 @@ func TestSpectatorPeerChooser_Choose_Success(t *testing.T) {
 
 	mockSpectator := NewMockSpectator(ctrl)
 	peerTransport := grpc.NewTransport()
+	require.NoError(t, peerTransport.Start())
+	defer peerTransport.Stop()
 
 	chooser := &SpectatorPeerChooser{
 		transport: peerTransport,
@@ -150,6 +152,8 @@ func TestSpectatorPeerChooser_Choose_ReusesPeer(t *testing.T) {
 
 	mockSpectator := NewMockSpectator(ctrl)
 	peerTransport := grpc.NewTransport()
+	require.NoError(t, peerTransport.Start())
+	defer peerTransport.Stop()
 
 	chooser := &SpectatorPeerChooser{
 		transport: peerTransport,