Add sanity check for shard id validation

common/persistence/cassandra/test.go

@@ -93,6 +93,7 @@ func (s *TestCluster) Config() config.Persistence {
 			"test": {Cassandra: &cfg, FaultInjection: s.faultInjection},
 		},
 		TransactionSizeLimit: dynamicconfig.GetIntPropertyFn(common.DefaultTransactionSizeLimit),
+		NumHistoryShards:     1024,
 	}
 }
 

common/persistence/client/factory.go

@@ -123,7 +123,7 @@ func (f *factoryImpl) NewShardManager() (p.ShardManager, error) {
 		return nil, err
 	}
 
-	result := p.NewShardManager(shardStore, f.serializer)
+	result := p.NewShardManager(shardStore, f.serializer, f.config.NumHistoryShards)
 	if f.ratelimiter != nil {
 		result = p.NewShardPersistenceRateLimitedClient(result, f.ratelimiter, f.logger)
 	}

common/persistence/shard_manager.go

@@ -29,25 +29,29 @@ import (
 
 	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
+	"go.temporal.io/api/serviceerror"
 
 	persistencespb "go.temporal.io/server/api/persistence/v1"
 	"go.temporal.io/server/common/persistence/serialization"
 	"go.temporal.io/server/common/primitives/timestamp"
 )
 
 type shardManagerImpl struct {
-	shardStore ShardStore
-	serializer serialization.Serializer
+	shardStore    ShardStore
+	serializer    serialization.Serializer
+	maxShardCount int32
 }
 
 // NewShardManager create a new instance of ShardManager
 func NewShardManager(
 	shardStore ShardStore,
 	serializer serialization.Serializer,
+	maxShardCount int32,
 ) ShardManager {
 	return &shardManagerImpl{
-		shardStore: shardStore,
-		serializer: serializer,
+		shardStore:    shardStore,
+		serializer:    serializer,
+		maxShardCount: maxShardCount,
 	}
 }
 
@@ -63,6 +67,10 @@ func (m *shardManagerImpl) GetOrCreateShard(
 	ctx context.Context,
 	request *GetOrCreateShardRequest,
 ) (*GetOrCreateShardResponse, error) {
+	if err := m.validateShardId(request.ShardID); err != nil {
+		return nil, err
+	}
+
 	createShardInfo := func() (int64, *commonpb.DataBlob, error) {
 		shardInfo := request.InitialShardInfo
 		if shardInfo == nil {
@@ -99,6 +107,9 @@ func (m *shardManagerImpl) UpdateShard(
 ) error {
 	shardInfo := request.ShardInfo
 	shardInfo.UpdateTime = timestamp.TimeNowPtrUtc()
+	if err := m.validateShardId(shardInfo.ShardId); err != nil {
+		return err
+	}
 
 	shardInfoBlob, err := m.serializer.ShardInfoToBlob(shardInfo, enumspb.ENCODING_TYPE_PROTO3)
 	if err != nil {
@@ -117,5 +128,19 @@ func (m *shardManagerImpl) AssertShardOwnership(
 	ctx context.Context,
 	request *AssertShardOwnershipRequest,
 ) error {
+	if err := m.validateShardId(request.ShardID); err != nil {
+		return err
+	}
 	return m.shardStore.AssertShardOwnership(ctx, request)
 }
+
+func (m *shardManagerImpl) validateShardId(shardId int32) error {
+	if shardId <= 0 {
+		return serviceerror.NewInvalidArgument("Cannot handle shard Id below the lower bound")
+	}
+
+	if shardId > m.maxShardCount {
+		return serviceerror.NewInvalidArgument("Cannot handle shard Id exceed the upper bound")
+	}
+	return nil
+}

common/persistence/sql/sqlPersistenceTest.go

@@ -116,6 +116,7 @@ func (s *TestCluster) Config() config.Persistence {
 			"test": {SQL: &cfg, FaultInjection: s.faultInjection},
 		},
 		TransactionSizeLimit: dynamicconfig.GetIntPropertyFn(common.DefaultTransactionSizeLimit),
+		NumHistoryShards:     1024,
 	}
 }
 

common/persistence/tests/execution_mutable_state.go

@@ -79,6 +79,7 @@ func NewExecutionMutableStateSuite(
 		ShardManager: p.NewShardManager(
 			shardStore,
 			serializer,
+			1024,
 		),
 		ExecutionManager: p.NewExecutionManager(
 			executionStore,

common/persistence/tests/execution_mutable_state_task.go

@@ -89,6 +89,7 @@ func NewExecutionMutableStateTaskSuite(
 		ShardManager: p.NewShardManager(
 			shardStore,
 			serializer,
+			1024,
 		),
 		ExecutionManager: p.NewExecutionManager(
 			executionStore,

common/persistence/tests/shard.go

@@ -62,6 +62,7 @@ func NewShardSuite(
 		ShardManager: p.NewShardManager(
 			shardStore,
 			serializer,
+			4,
 		),
 		Logger: logger,
 	}

service/history/shard/context_impl.go

@@ -1468,9 +1468,11 @@ func (s *ContextImpl) handleWriteErrorAndUpdateMaxReadLevelLocked(err error, new
 		// No special handling required for these errors.
 		return err
 
-	case *persistence.ShardOwnershipLostError:
+	case *persistence.ShardOwnershipLostError,
+		*serviceerror.InvalidArgument:
 		// Shard is stolen, trigger shutdown of history engine.
 		// Handling of max read level doesn't matter here.
+		// Or the shard id is invalid. Stop the shard.
 		s.transition(contextRequestStop{})
 		return err