Support for creating ephemeral shards (#7209)

jakobht · web-flow · commit 1249e4e4eedf · 2025-08-27T10:35:47.000+02:00
What changed?
We now support ephemeral shards

Shards are created if they do not already exist.
An executors shards are no longer deleted when the executor disappears
We now load the set of shards from the database for ephemeral shards
Future PRs will:

Support removing shards when they are no longer needed
Support the executor library to heartbeat immediately if an unknown shard is requested (so we do not have to wait for the next heartbeat to process shard requests)
Why?
We need to support ephemeral shards for use cases where the set of shard is changing (e.g. Cadence Matching)

How did you test it?

Potential risks

Release notes

Documentation Changes
diff --git a/service/sharddistributor/handler/handler.go b/service/sharddistributor/handler/handler.go
@@ -26,6 +26,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"slices"
 	"sync"
 
@@ -92,6 +93,10 @@ func (h *handlerImpl) GetShardOwner(ctx context.Context, request *types.GetShard
 
 	executorID, err := h.storage.GetShardOwner(ctx, request.Namespace, request.ShardKey)
 	if errors.Is(err, store.ErrShardNotFound) {
+		if h.shardDistributionCfg.Namespaces[namespaceIdx].Type == config.NamespaceTypeEphemeral {
+			return h.assignEphemeralShard(ctx, request.Namespace, request.ShardKey)
+		}
+
 		return nil, &types.ShardNotFoundError{
 			Namespace: request.Namespace,
 			ShardKey:  request.ShardKey,
@@ -108,3 +113,33 @@ func (h *handlerImpl) GetShardOwner(ctx context.Context, request *types.GetShard
 
 	return resp, nil
 }
+
+func (h *handlerImpl) assignEphemeralShard(ctx context.Context, namespace string, shardID string) (*types.GetShardOwnerResponse, error) {
+
+	// Get the current state of the namespace and find the executor with the least assigned shards
+	state, err := h.storage.GetState(ctx, namespace)
+	if err != nil {
+		return nil, fmt.Errorf("get state: %w", err)
+	}
+
+	var executor string
+	minAssignedShards := math.MaxInt
+
+	for assignedExecutor, assignment := range state.ShardAssignments {
+		if len(assignment.AssignedShards) < minAssignedShards {
+			minAssignedShards = len(assignment.AssignedShards)
+			executor = assignedExecutor
+		}
+	}
+
+	// Assign the shard to the executor with the least assigned shards
+	err = h.storage.AssignShard(ctx, namespace, shardID, executor)
+	if err != nil {
+		return nil, fmt.Errorf("assign ephemeral shard: %w", err)
+	}
+
+	return &types.GetShardOwnerResponse{
+		Owner:     executor,
+		Namespace: namespace,
+	}, nil
+}
diff --git a/service/sharddistributor/handler/handler_test.go b/service/sharddistributor/handler/handler_test.go
@@ -37,15 +37,22 @@ import (
 )
 
 const (
-	_testNamespace = "test-matching"
+	_testNamespaceFixed     = "test-fixed"
+	_testNamespaceEphemeral = "test-ephemeral"
 )
 
 func TestGetShardOwner(t *testing.T) {
 	cfg := config.ShardDistribution{
 		Enabled: true,
 		Namespaces: []config.Namespace{
 			{
-				Name: _testNamespace,
+				Name:     _testNamespaceFixed,
+				Type:     config.NamespaceTypeFixed,
+				ShardNum: 32,
+			},
+			{
+				Name: _testNamespaceEphemeral,
+				Type: config.NamespaceTypeEphemeral,
 			},
 		},
 	}
@@ -59,50 +66,119 @@ func TestGetShardOwner(t *testing.T) {
 		expectedErrMsg string
 	}{
 		{
-			name: "Existing_Success",
+			name: "InvalidNamespace",
+			request: &types.GetShardOwnerRequest{
+				Namespace: "namespace not found invalidNamespace",
+				ShardKey:  "1",
+			},
+			expectedError:  true,
+			expectedErrMsg: "namespace not found",
+		},
+		{
+			name: "LookupError",
+			request: &types.GetShardOwnerRequest{
+				Namespace: _testNamespaceFixed,
+				ShardKey:  "1",
+			},
+			setupMocks: func(mockStore *store.MockStore) {
+				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespaceFixed, "1").Return("", errors.New("lookup error"))
+			},
+			expectedError:  true,
+			expectedErrMsg: "lookup error",
+		},
+		{
+			name: "Existing_Success_Fixed",
 			request: &types.GetShardOwnerRequest{
-				Namespace: _testNamespace,
+				Namespace: _testNamespaceFixed,
 				ShardKey:  "123",
 			},
 			setupMocks: func(mockStore *store.MockStore) {
-				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespace, "123").Return("owner1", nil)
+				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespaceFixed, "123").Return("owner1", nil)
 			},
 			expectedOwner: "owner1",
 			expectedError: false,
 		},
-
 		{
-			name: "InvalidNamespace",
+			name: "ShardNotFound_Fixed",
 			request: &types.GetShardOwnerRequest{
-				Namespace: "namespace not found invalidNamespace",
-				ShardKey:  "1",
+				Namespace: _testNamespaceFixed,
+				ShardKey:  "NON-EXISTING-SHARD",
+			},
+			setupMocks: func(mockStore *store.MockStore) {
+				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespaceFixed, "NON-EXISTING-SHARD").Return("", store.ErrShardNotFound)
 			},
 			expectedError:  true,
-			expectedErrMsg: "namespace not found",
+			expectedErrMsg: "shard not found",
 		},
 		{
-			name: "LookupError",
+			name: "Existing_Success_Ephemeral",
 			request: &types.GetShardOwnerRequest{
-				Namespace: _testNamespace,
-				ShardKey:  "1",
+				Namespace: _testNamespaceEphemeral,
+				ShardKey:  "123",
+			},
+			setupMocks: func(mockStore *store.MockStore) {
+				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespaceEphemeral, "123").Return("owner1", nil)
+			},
+			expectedOwner: "owner1",
+			expectedError: false,
+		},
+		{
+			name: "ShardNotFound_Ephemeral",
+			request: &types.GetShardOwnerRequest{
+				Namespace: _testNamespaceEphemeral,
+				ShardKey:  "NON-EXISTING-SHARD",
 			},
 			setupMocks: func(mockStore *store.MockStore) {
-				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespace, "1").Return("", errors.New("lookup error"))
+				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespaceEphemeral, "NON-EXISTING-SHARD").Return("", store.ErrShardNotFound)
+				mockStore.EXPECT().GetState(gomock.Any(), _testNamespaceEphemeral).Return(&store.NamespaceState{
+					ShardAssignments: map[string]store.AssignedState{
+						"owner1": {
+							AssignedShards: map[string]*types.ShardAssignment{
+								"shard1": {Status: types.AssignmentStatusREADY},
+								"shard2": {Status: types.AssignmentStatusREADY},
+								"shard3": {Status: types.AssignmentStatusREADY},
+							},
+						},
+						"owner2": {
+							AssignedShards: map[string]*types.ShardAssignment{
+								"shard4": {Status: types.AssignmentStatusREADY},
+							},
+						},
+					},
+				}, nil)
+				// owner2 has the fewest shards assigned, so we assign the shard to it
+				mockStore.EXPECT().AssignShard(gomock.Any(), _testNamespaceEphemeral, "NON-EXISTING-SHARD", "owner2").Return(nil)
+			},
+			expectedOwner: "owner2",
+			expectedError: false,
+		},
+		{
+			name: "ShardNotFound_Ephemeral_GetStateFailure",
+			request: &types.GetShardOwnerRequest{
+				Namespace: _testNamespaceEphemeral,
+				ShardKey:  "NON-EXISTING-SHARD",
+			},
+			setupMocks: func(mockStore *store.MockStore) {
+				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespaceEphemeral, "NON-EXISTING-SHARD").Return("", store.ErrShardNotFound)
+				mockStore.EXPECT().GetState(gomock.Any(), _testNamespaceEphemeral).Return(nil, errors.New("get state failure"))
 			},
 			expectedError:  true,
-			expectedErrMsg: "lookup error",
+			expectedErrMsg: "get state failure",
 		},
 		{
-			name: "ShardNotFound",
+			name: "ShardNotFound_Ephemeral_AssignShardFailure",
 			request: &types.GetShardOwnerRequest{
-				Namespace: _testNamespace,
+				Namespace: _testNamespaceEphemeral,
 				ShardKey:  "NON-EXISTING-SHARD",
 			},
 			setupMocks: func(mockStore *store.MockStore) {
-				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespace, "NON-EXISTING-SHARD").Return("", store.ErrShardNotFound)
+				mockStore.EXPECT().GetShardOwner(gomock.Any(), _testNamespaceEphemeral, "NON-EXISTING-SHARD").Return("", store.ErrShardNotFound)
+				mockStore.EXPECT().GetState(gomock.Any(), _testNamespaceEphemeral).Return(&store.NamespaceState{
+					ShardAssignments: map[string]store.AssignedState{"owner1": {AssignedShards: map[string]*types.ShardAssignment{}}}}, nil)
+				mockStore.EXPECT().AssignShard(gomock.Any(), _testNamespaceEphemeral, "NON-EXISTING-SHARD", "owner1").Return(errors.New("assign shard failure"))
 			},
 			expectedError:  true,
-			expectedErrMsg: "shard not found",
+			expectedErrMsg: "assign shard failure",
 		},
 	}
 
diff --git a/service/sharddistributor/leader/process/processor.go b/service/sharddistributor/leader/process/processor.go
@@ -316,8 +316,8 @@ func (p *namespaceProcessor) rebalanceShardsImpl(ctx context.Context, metricsLoo
 
 func (p *namespaceProcessor) findShardsToReassign(activeExecutors []string, namespaceState *store.NamespaceState) ([]string, map[string][]string) {
 	allShards := make(map[string]struct{})
-	for _, shardID := range getShards(p.namespaceCfg) {
-		allShards[strconv.FormatInt(shardID, 10)] = struct{}{}
+	for _, shardID := range getShards(p.namespaceCfg, namespaceState) {
+		allShards[shardID] = struct{}{}
 	}
 
 	shardsToReassign := make([]string, 0)
@@ -449,17 +449,23 @@ func assignShardsToEmptyExecutors(currentAssignments map[string][]string) bool {
 	return true
 }
 
-func getShards(cfg config.Namespace) []int64 {
+func getShards(cfg config.Namespace, namespaceState *store.NamespaceState) []string {
 	if cfg.Type == config.NamespaceTypeFixed {
-		return makeRange(0, cfg.ShardNum-1)
+		return makeShards(cfg.ShardNum)
+	} else if cfg.Type == config.NamespaceTypeEphemeral {
+		shards := make([]string, 0)
+		for shardID := range namespaceState.Shards {
+			shards = append(shards, shardID)
+		}
+		return shards
 	}
 	return nil
 }
 
-func makeRange(min, max int64) []int64 {
-	a := make([]int64, max-min+1)
-	for i := range a {
-		a[i] = min + int64(i)
+func makeShards(num int64) []string {
+	shards := make([]string, num)
+	for i := range num {
+		shards[i] = strconv.FormatInt(i, 10)
 	}
-	return a
+	return shards
 }
diff --git a/service/sharddistributor/leader/process/processor_test.go b/service/sharddistributor/leader/process/processor_test.go
@@ -3,6 +3,7 @@ package process
 import (
 	"context"
 	"errors"
+	"slices"
 	"sync"
 	"testing"
 	"time"
@@ -386,15 +387,34 @@ func TestRebalanceShards_WithUnassignedShards(t *testing.T) {
 }
 
 func TestGetShards_Utility(t *testing.T) {
-	// Fixed type
-	cfg := config.Namespace{Type: config.NamespaceTypeFixed, ShardNum: 5}
-	shards := getShards(cfg)
-	assert.Equal(t, []int64{0, 1, 2, 3, 4}, shards)
-
-	// Other type
-	cfg = config.Namespace{Type: "other"}
-	shards = getShards(cfg)
-	assert.Nil(t, shards)
+	t.Run("Fixed type", func(t *testing.T) {
+		cfg := config.Namespace{Type: config.NamespaceTypeFixed, ShardNum: 5}
+		shards := getShards(cfg, nil)
+		assert.Equal(t, []string{"0", "1", "2", "3", "4"}, shards)
+	})
+
+	t.Run("Ephemeral type", func(t *testing.T) {
+		cfg := config.Namespace{Type: config.NamespaceTypeEphemeral}
+		nsState := &store.NamespaceState{
+			Shards: map[string]store.ShardState{
+				"s0": {ExecutorID: "exec-1"},
+				"s1": {ExecutorID: "exec-1"},
+				"s2": {ExecutorID: "exec-1"},
+				"s3": {ExecutorID: "exec-1"},
+				"s4": {ExecutorID: "exec-1"},
+			},
+		}
+		shards := getShards(cfg, nsState)
+		slices.Sort(shards)
+		assert.Equal(t, []string{"s0", "s1", "s2", "s3", "s4"}, shards)
+	})
+
+	// Unknown type
+	t.Run("Other type", func(t *testing.T) {
+		cfg := config.Namespace{Type: "other"}
+		shards := getShards(cfg, nil)
+		assert.Nil(t, shards)
+	})
 }
 
 func TestAssignShardsToEmptyExecutors(t *testing.T) {
diff --git a/service/sharddistributor/store/etcd/etcdstore.go b/service/sharddistributor/store/etcd/etcdstore.go
@@ -465,44 +465,15 @@ func (s *Store) AssignShard(ctx context.Context, namespace, shardID, executorID
 	}
 }
 
+// DeleteExecutors deletes the given executors from the store. It does not delete the shards owned by the executors, this
+// should be handled by the namespace processor loop as we want to reassign, not delete the shards.
 func (s *Store) DeleteExecutors(ctx context.Context, namespace string, executorIDs []string, guard store.GuardFunc) error {
 	if len(executorIDs) == 0 {
 		return nil
 	}
 	var ops []clientv3.Op
 
-	// For each executor, find all shards it owns and create conditional deletion operations.
 	for _, executorID := range executorIDs {
-		// 1. Find the shards currently assigned to this executor.
-		assignedStateKey := s.buildExecutorKey(namespace, executorID, executorAssignedStateKey)
-		resp, err := s.client.Get(ctx, assignedStateKey)
-		if err != nil {
-			return fmt.Errorf("get assigned state for executor %s: %w", executorID, err)
-		}
-
-		if len(resp.Kvs) > 0 {
-			var state store.AssignedState
-			if err := json.Unmarshal(resp.Kvs[0].Value, &state); err != nil {
-				return fmt.Errorf("unmarshal assigned state for executor %s: %w", executorID, err)
-			}
-
-			// 2. For each shard, create a nested transaction to conditionally delete it.
-			for shardID := range state.AssignedShards {
-				shardOwnerKey := s.buildShardKey(namespace, shardID, shardAssignedKey)
-
-				// This is an atomic check-and-delete operation for a single shard.
-				// IF the shard owner is still the executor being deleted,
-				// THEN delete the shard ownership key.
-				conditionalDelete := clientv3.OpTxn(
-					[]clientv3.Cmp{clientv3.Compare(clientv3.Value(shardOwnerKey), "=", executorID)}, // IF condition
-					[]clientv3.Op{clientv3.OpDelete(shardOwnerKey)},                                  // THEN operations
-					nil, // ELSE do nothing
-				)
-				ops = append(ops, conditionalDelete)
-			}
-		}
-
-		// 3. Add the unconditional deletion for the executor's entire record.
 		executorPrefix := s.buildExecutorKey(namespace, executorID, "")
 		ops = append(ops, clientv3.OpDelete(executorPrefix, clientv3.WithPrefix()))
 	}
diff --git a/service/sharddistributor/store/etcd/etcdstore_test.go b/service/sharddistributor/store/etcd/etcdstore_test.go
