Update shard controller getters to use atomic (#6864)

Author: Shaddoll

service/history/shard/controller.go

@@ -75,6 +75,11 @@ type (
 		ShardIDs() []int32
 	}
 
+	shardIDSnapshot struct {
+		shardIDs  []int32
+		numShards int
+	}
+
 	controller struct {
 		resource.Resource
 
@@ -90,7 +95,8 @@ type (
 		metricsScope       metrics.Scope
 
 		sync.RWMutex
-		historyShards map[int]*historyShardsItem
+		historyShards   map[int]*historyShardsItem
+		shardIDSnapshot atomic.Pointer[shardIDSnapshot]
 	}
 
 	historyShardsItemStatus int
@@ -223,22 +229,19 @@ func (c *controller) Status() int32 {
 }
 
 func (c *controller) NumShards() int {
-	nShards := 0
-	c.RLock()
-	nShards = len(c.historyShards)
-	c.RUnlock()
-	return nShards
+	s := c.shardIDSnapshot.Load()
+	if s == nil {
+		return 0
+	}
+	return s.numShards
 }
 
 func (c *controller) ShardIDs() []int32 {
-	c.RLock()
-	ids := []int32{}
-	for id := range c.historyShards {
-		id32 := int32(id)
-		ids = append(ids, id32)
+	s := c.shardIDSnapshot.Load()
+	if s == nil {
+		return []int32{}
 	}
-	c.RUnlock()
-	return ids
+	return s.shardIDs
 }
 
 func (c *controller) removeEngineForShard(shardID int, shardItem *historyShardsItem) {
@@ -330,6 +333,7 @@ func (c *controller) getOrCreateHistoryShardItem(shardID int) (*historyShardsIte
 			return nil, err
 		}
 		c.historyShards[shardID] = shardItem
+		c.updateShardIDSnapshotLocked()
 		c.metricsScope.IncCounter(metrics.ShardItemCreatedCounter)
 
 		shardItem.logger.Info("Shard item state changed", tag.LifeCycleStarted, tag.ComponentShardItem)
@@ -340,6 +344,18 @@ func (c *controller) getOrCreateHistoryShardItem(shardID int) (*historyShardsIte
 	return nil, CreateShardOwnershipLostError(c.GetHostInfo(), info)
 }
 
+func (c *controller) updateShardIDSnapshotLocked() {
+	shardIDs := make([]int32, 0, len(c.historyShards))
+	for shardID := range c.historyShards {
+		shardIDs = append(shardIDs, int32(shardID))
+	}
+	snapshot := &shardIDSnapshot{
+		shardIDs:  shardIDs,
+		numShards: len(shardIDs),
+	}
+	c.shardIDSnapshot.Store(snapshot)
+}
+
 func (c *controller) removeHistoryShardItem(shardID int, shardItem *historyShardsItem) (*historyShardsItem, error) {
 	c.Lock()
 	defer c.Unlock()
@@ -355,6 +371,7 @@ func (c *controller) removeHistoryShardItem(shardID int, shardItem *historyShard
 	}
 
 	delete(c.historyShards, shardID)
+	c.updateShardIDSnapshotLocked()
 	c.metricsScope.IncCounter(metrics.ShardItemRemovedCounter)
 	currentShardItem.logger.Info("Shard item state changed", tag.LifeCycleStopped, tag.ComponentShardItem, tag.Number(int64(len(c.historyShards))))
 	return currentShardItem, nil
@@ -450,6 +467,7 @@ func (c *controller) doShutdown() {
 		item.stopEngine()
 	}
 	c.historyShards = nil
+	c.updateShardIDSnapshotLocked()
 }
 
 func (c *controller) isShuttingDown() bool {

service/history/shard/controller_benchmark_test.go

@@ -0,0 +1,62 @@
+// The MIT License (MIT)
+
+// Copyright (c) 2017-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 shard
+
+import (
+	"sync"
+	"testing"
+)
+
+// BenchmarkController_ShardIDs-96         52588224               293.9 ns/op            66 B/op          0 allocs/op
+// go test -bench=. --benchtime=10s --benchmem
+// goos: linux
+// goarch: amd64
+// pkg: github.com/uber/cadence/service/history/shard
+// cpu: AMD EPYC 7B13
+// With the old approach, the benchmark result is:
+// BenchmarkController_ShardIDs-96            39314            324629 ns/op          272333 B/op         19 allocs/op
+func BenchmarkController_ShardIDs(b *testing.B) {
+	numShards := 16384
+	historyShards := make(map[int]*historyShardsItem)
+	for i := 0; i < numShards; i++ {
+		historyShards[i] = &historyShardsItem{shardID: i}
+	}
+	shardController := &controller{
+		historyShards: historyShards,
+	}
+	var wg sync.WaitGroup
+	for i := 0; i < b.N; i++ {
+		if i%1000 == 0 { // update is much much less frequent than read
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				shardController.Lock()
+				shardController.updateShardIDSnapshotLocked()
+				shardController.Unlock()
+			}()
+		}
+		shardController.ShardIDs()
+		shardController.NumShards()
+	}
+	wg.Wait()
+}

service/history/shard/controller_test.go

@@ -175,6 +175,8 @@ func (s *controllerSuite) TestAcquireShardSuccess() {
 		count++
 	}
 	s.Equal(3, count)
+	s.Equal(3, s.shardController.NumShards())
+	s.ElementsMatch([]int32{0, 4, 8}, s.shardController.ShardIDs())
 }
 
 func (s *controllerSuite) TestAcquireShardsConcurrently() {
@@ -259,6 +261,8 @@ func (s *controllerSuite) TestAcquireShardsConcurrently() {
 		count++
 	}
 	s.Equal(3, count)
+	s.Equal(3, s.shardController.NumShards())
+	s.ElementsMatch([]int32{0, 4, 8}, s.shardController.ShardIDs())
 }
 
 func (s *controllerSuite) TestAcquireShardLookupFailure() {
@@ -273,6 +277,8 @@ func (s *controllerSuite) TestAcquireShardLookupFailure() {
 		s.mockMembershipResolver.EXPECT().Lookup(service.History, string(rune(shardID))).Return(membership.HostInfo{}, errors.New("ring failure")).Times(1)
 		s.Nil(s.shardController.GetEngineForShard(shardID))
 	}
+	s.Equal(0, s.shardController.NumShards())
+	s.Empty(s.shardController.ShardIDs())
 }
 
 func (s *controllerSuite) TestAcquireShardRenewSuccess() {
@@ -341,11 +347,17 @@ func (s *controllerSuite) TestAcquireShardRenewSuccess() {
 
 	s.shardController.acquireShards()
 
+	s.Equal(2, s.shardController.NumShards())
+	s.ElementsMatch([]int32{0, 1}, s.shardController.ShardIDs())
+
 	for shardID := 0; shardID < numShards; shardID++ {
 		s.mockMembershipResolver.EXPECT().Lookup(service.History, string(rune(shardID))).Return(s.hostInfo, nil).Times(1)
 	}
 	s.shardController.acquireShards()
 
+	s.Equal(2, s.shardController.NumShards())
+	s.ElementsMatch([]int32{0, 1}, s.shardController.ShardIDs())
+
 	for shardID := 0; shardID < numShards; shardID++ {
 		s.NotNil(s.shardController.GetEngineForShard(shardID))
 	}
@@ -510,6 +522,9 @@ func (s *controllerSuite) TestHistoryEngineClosed() {
 		s.mockMembershipResolver.EXPECT().Lookup(service.History, string(rune(shardID))).Return(s.hostInfo, nil).AnyTimes()
 	}
 	s.shardController.Stop()
+
+	s.Equal(0, s.shardController.NumShards())
+	s.Empty(s.shardController.ShardIDs())
 }
 
 func (s *controllerSuite) TestShardControllerClosed() {
@@ -555,6 +570,9 @@ func (s *controllerSuite) TestShardControllerClosed() {
 	}
 	s.shardController.Stop()
 	workerWG.Wait()
+
+	s.Equal(0, s.shardController.NumShards())
+	s.Empty(s.shardController.ShardIDs())
 }
 
 func (s *controllerSuite) TestGetOrCreateHistoryShardItem_InvalidShardID_Error() {