Add a component to assign polling shards (#3775)

* Add a component to assign polling shards

Author: yux0

common/cluster/metadata.go

@@ -104,6 +104,7 @@ type (
 		InitialFailoverVersion int64 `yaml:"initialFailoverVersion"`
 		// Address indicate the remote service address(Host:Port). Host can be DNS name.
 		RPCAddress string `yaml:"rpcAddress"`
+		ShardCount int32
 		// private field to track cluster information updates
 		version int64
 	}

service/history/replication/poller_manager.go

@@ -0,0 +1,82 @@
+// The MIT License
+//
+// Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
+//
+// Copyright (c) 2020 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package replication
+
+import (
+	"fmt"
+
+	"go.temporal.io/server/common/cluster"
+)
+
+type (
+	pollerManagerImpl struct {
+		currentShardId  int32
+		clusterMetadata cluster.Metadata
+	}
+)
+
+func newPollerManager(
+	currentShardId int32,
+	clusterMetadata cluster.Metadata,
+) *pollerManagerImpl {
+	return &pollerManagerImpl{
+		currentShardId:  currentShardId,
+		clusterMetadata: clusterMetadata,
+	}
+}
+
+func (p pollerManagerImpl) getPollingShardIDs(remoteClusterName string) []int32 {
+	currentCluster := p.clusterMetadata.GetCurrentClusterName()
+	allClusters := p.clusterMetadata.GetAllClusterInfo()
+	currentClusterInfo, ok := allClusters[currentCluster]
+	if !ok {
+		panic("Cannot get current cluster info from cluster metadata cache")
+	}
+	remoteClusterInfo, ok := allClusters[remoteClusterName]
+	if !ok {
+		panic(fmt.Sprintf("Cannot get remote cluster %s info from cluster metadata cache", remoteClusterName))
+	}
+	return generatePollingShardIDs(p.currentShardId, currentClusterInfo.ShardCount, remoteClusterInfo.ShardCount)
+}
+
+func generatePollingShardIDs(localShardId int32, localShardCount int32, remoteShardCount int32) []int32 {
+	var pollingShards []int32
+	if remoteShardCount <= localShardCount {
+		if localShardId <= remoteShardCount {
+			pollingShards = append(pollingShards, localShardId)
+		}
+		return pollingShards
+	}
+
+	// remoteShardCount > localShardCount, replication poller will poll from multiple remote shard.
+	// The remote shard count and local shard count must be multiples.
+	if remoteShardCount%localShardCount != 0 {
+		panic(fmt.Sprintf("Remote shard count %d and local shard count %d are not multiples.", remoteShardCount, localShardCount))
+	}
+	for i := localShardId; i <= remoteShardCount; i += localShardCount {
+		pollingShards = append(pollingShards, i)
+	}
+	return pollingShards
+}

service/history/replication/poller_manager_test.go

@@ -0,0 +1,97 @@
+// The MIT License
+//
+// Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
+//
+// Copyright (c) 2020 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package replication
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetPollingShardIds(t *testing.T) {
+	testCases := []struct {
+		shardID          int32
+		remoteShardCount int32
+		localShardCount  int32
+		expectedPanic    bool
+		expectedShardIDs []int32
+	}{
+		{
+			1,
+			4,
+			4,
+			false,
+			[]int32{1},
+		},
+		{
+			1,
+			2,
+			4,
+			false,
+			[]int32{1},
+		},
+		{
+			3,
+			2,
+			4,
+			false,
+			[]int32{},
+		},
+		{
+			1,
+			16,
+			4,
+			false,
+			[]int32{1, 5, 9, 13},
+		},
+		{
+			4,
+			16,
+			4,
+			false,
+			[]int32{4, 8, 12, 16},
+		},
+		{
+			4,
+			17,
+			4,
+			true,
+			[]int32{},
+		},
+	}
+	for idx, tt := range testCases {
+		t.Run(fmt.Sprintf("Testcase %d", idx), func(t *testing.T) {
+			t.Parallel()
+			defer func() {
+				if r := recover(); tt.expectedPanic && r == nil {
+					t.Errorf("The code did not panic")
+				}
+			}()
+			shardIDs := generatePollingShardIDs(tt.shardID, tt.localShardCount, tt.remoteShardCount)
+			assert.Equal(t, tt.expectedShardIDs, shardIDs)
+		})
+	}
+}