Router: get rid of nameToReplicasetMutex, use atomic instead

Author: nurzhan-saktaganov

CHANGELOG.md

@@ -4,6 +4,7 @@ CHANGES:
 * Slog provider moved to providers directory.
 * More strict check of vshard.storage.call response.
 * Bump go-tarantool from v2.3.0 to v2.3.1.
+* Get rid of nameToReplicasetMutex, use atomic instead.
 
 ## v2.0.5
 

api_test.go

@@ -9,14 +9,6 @@ import (
 	"github.com/vmihailenco/msgpack/v5/msgpcode"
 )
 
-var emptyRouter = &Router{
-	cfg: Config{
-		TotalBucketCount: uint64(10),
-		Loggerf:          emptyLogfProvider,
-		Metrics:          emptyMetricsProvider,
-	},
-}
-
 func TestVshardMode_String_NotEmpty(t *testing.T) {
 	t.Parallel()
 	require.NotEmpty(t, ReadMode.String())
@@ -25,6 +17,14 @@ func TestVshardMode_String_NotEmpty(t *testing.T) {
 
 func TestRouter_RouterRouteAll(t *testing.T) {
 	t.Parallel()
+	var emptyRouter = &Router{
+		cfg: Config{
+			TotalBucketCount: uint64(10),
+			Loggerf:          emptyLogfProvider,
+			Metrics:          emptyMetricsProvider,
+		},
+	}
+	emptyRouter.setEmptyNameToReplicaset()
 	m := emptyRouter.RouteAll()
 	require.Empty(t, m)
 }

topology.go

@@ -11,6 +11,8 @@ import (
 var (
 	ErrReplicasetExists    = fmt.Errorf("replicaset already exists")
 	ErrReplicasetNotExists = fmt.Errorf("replicaset not exists")
+
+	ErrConcurrentTopologyChangeDetected = fmt.Errorf("concurrent topology change detected")
 )
 
 // TopologyController is an entity that allows you to interact with the topology.
@@ -33,18 +35,21 @@ func copyMap[K comparable, V any](m map[K]V) map[K]V {
 	return copy
 }
 
-func (r *Router) getNameToReplicaset() map[string]*Replicaset {
-	r.nameToReplicasetMutex.RLock()
-	nameToReplicasetRef := r.nameToReplicaset
-	r.nameToReplicasetMutex.RUnlock()
+func (r *Router) setEmptyNameToReplicaset() {
+	var nameToReplicasetRef map[string]*Replicaset
+	_ = r.swapNameToReplicaset(nil, &nameToReplicasetRef)
+}
 
-	return nameToReplicasetRef
+func (r *Router) swapNameToReplicaset(old, new *map[string]*Replicaset) error {
+	if swapped := r.nameToReplicaset.CompareAndSwap(old, new); !swapped {
+		return ErrConcurrentTopologyChangeDetected
+	}
+	return nil
 }
 
-func (r *Router) setNameToReplicaset(nameToReplicasetNew map[string]*Replicaset) {
-	r.nameToReplicasetMutex.Lock()
-	r.nameToReplicaset = nameToReplicasetNew
-	r.nameToReplicasetMutex.Unlock()
+func (r *Router) getNameToReplicaset() map[string]*Replicaset {
+	ptr := r.nameToReplicaset.Load()
+	return *ptr
 }
 
 func (r *Router) Topology() TopologyController {
@@ -121,9 +126,8 @@ func (r *Router) AddReplicaset(ctx context.Context, rsInfo ReplicasetInfo, insta
 		return err
 	}
 
-	nameToReplicasetOld := r.getNameToReplicaset()
-
-	if _, ok := nameToReplicasetOld[rsInfo.Name]; ok {
+	nameToReplicasetOldPtr := r.nameToReplicaset.Load()
+	if _, ok := (*nameToReplicasetOldPtr)[rsInfo.Name]; ok {
 		return ErrReplicasetExists
 	}
 
@@ -168,14 +172,14 @@ func (r *Router) AddReplicaset(ctx context.Context, rsInfo ReplicasetInfo, insta
 	}
 
 	// Create an entirely new map object
-	nameToReplicasetNew := copyMap(nameToReplicasetOld)
+	nameToReplicasetNew := copyMap(*nameToReplicasetOldPtr)
 	nameToReplicasetNew[rsInfo.Name] = replicaset // add when conn is ready
 
-	// We could detect concurrent access to the TopologyController interface
-	// by comparing references to r.idToReplicaset and idToReplicasetOld.
-	// But it requires reflection which I prefer to avoid.
-	// See: https://stackoverflow.com/questions/58636694/how-to-know-if-2-go-maps-reference-the-same-data.
-	r.setNameToReplicaset(nameToReplicasetNew)
+	if err = r.swapNameToReplicaset(nameToReplicasetOldPtr, &nameToReplicasetNew); err != nil {
+		// replicaset has not added, so just close it
+		_ = replicaset.conn.Close()
+		return err
+	}
 
 	return nil
 }
@@ -198,18 +202,19 @@ func (r *Router) AddReplicasets(ctx context.Context, replicasets map[ReplicasetI
 func (r *Router) RemoveReplicaset(ctx context.Context, rsName string) []error {
 	r.log().Debugf(ctx, "Trying to remove replicaset %s from router topology", rsName)
 
-	nameToReplicasetOld := r.getNameToReplicaset()
-
-	rs := nameToReplicasetOld[rsName]
+	nameToReplicasetOldPtr := r.nameToReplicaset.Load()
+	rs := (*nameToReplicasetOldPtr)[rsName]
 	if rs == nil {
 		return []error{ErrReplicasetNotExists}
 	}
 
 	// Create an entirely new map object
-	nameToReplicasetNew := copyMap(nameToReplicasetOld)
+	nameToReplicasetNew := copyMap(*nameToReplicasetOldPtr)
 	delete(nameToReplicasetNew, rsName)
 
-	r.setNameToReplicaset(nameToReplicasetNew)
+	if err := r.swapNameToReplicaset(nameToReplicasetOldPtr, &nameToReplicasetNew); err != nil {
+		return []error{err}
+	}
 
 	return rs.conn.CloseGraceful()
 }

topology_test.go

@@ -23,11 +23,11 @@ func TestController_AddInstance(t *testing.T) {
 
 	t.Run("no such replicaset", func(t *testing.T) {
 		router := Router{
-			nameToReplicaset: map[string]*Replicaset{},
 			cfg: Config{
 				Loggerf: emptyLogfProvider,
 			},
 		}
+		router.setEmptyNameToReplicaset()
 
 		err := router.Topology().AddInstance(ctx, uuid.New().String(), InstanceInfo{
 			Addr: "127.0.0.1:8060",
@@ -38,11 +38,11 @@ func TestController_AddInstance(t *testing.T) {
 
 	t.Run("invalid instance info", func(t *testing.T) {
 		router := Router{
-			nameToReplicaset: map[string]*Replicaset{},
 			cfg: Config{
 				Loggerf: emptyLogfProvider,
 			},
 		}
+		router.setEmptyNameToReplicaset()
 
 		err := router.Topology().AddInstance(ctx, uuid.New().String(), InstanceInfo{})
 		require.True(t, errors.Is(err, ErrInvalidInstanceInfo))
@@ -54,11 +54,11 @@ func TestController_RemoveInstance_NoSuchReplicaset(t *testing.T) {
 	ctx := context.Background()
 
 	router := Router{
-		nameToReplicaset: map[string]*Replicaset{},
 		cfg: Config{
 			Loggerf: emptyLogfProvider,
 		},
 	}
+	router.setEmptyNameToReplicaset()
 
 	err := router.Topology().RemoveInstance(ctx, uuid.New().String(), "")
 	require.True(t, errors.Is(err, ErrReplicasetNotExists))
@@ -81,15 +81,15 @@ func TestController_RemoveInstance_NoReplicasetNameProvided(t *testing.T) {
 	mp.On("Remove", mock.Anything).Return(nil)
 
 	router := Router{
-		nameToReplicaset: map[string]*Replicaset{
-			"replicaset_1": {
-				conn: mp,
-			},
-		},
 		cfg: Config{
 			Loggerf: emptyLogfProvider,
 		},
 	}
+	_ = router.swapNameToReplicaset(nil, &map[string]*Replicaset{
+		"replicaset_1": {
+			conn: mp,
+		},
+	})
 
 	err := router.Topology().RemoveInstance(ctx, "", instanceName)
 	require.NoError(t, err)
@@ -106,13 +106,13 @@ func TestController_RemoveReplicaset(t *testing.T) {
 	mPool.On("CloseGraceful").Return(nil)
 
 	router := Router{
-		nameToReplicaset: map[string]*Replicaset{
-			uuidToRemove.String(): {conn: mPool},
-		},
 		cfg: Config{
 			Loggerf: emptyLogfProvider,
 		},
 	}
+	_ = router.swapNameToReplicaset(nil, &map[string]*Replicaset{
+		uuidToRemove.String(): {conn: mPool},
+	})
 
 	t.Run("no such replicaset", func(t *testing.T) {
 		t.Parallel()
@@ -132,13 +132,13 @@ func TestRouter_AddReplicaset_AlreadyExists(t *testing.T) {
 	alreadyExistingRsName := uuid.New().String()
 
 	router := Router{
-		nameToReplicaset: map[string]*Replicaset{
-			alreadyExistingRsName: {conn: nil},
-		},
 		cfg: Config{
 			Loggerf: emptyLogfProvider,
 		},
 	}
+	_ = router.swapNameToReplicaset(nil, &map[string]*Replicaset{
+		alreadyExistingRsName: {conn: nil},
+	})
 
 	// Test that such replicaset already exists
 	err := router.AddReplicaset(ctx, ReplicasetInfo{Name: alreadyExistingRsName}, []InstanceInfo{})
@@ -153,13 +153,13 @@ func TestRouter_AddReplicaset_InvalidReplicaset(t *testing.T) {
 	alreadyExistingRsName := uuid.New().String()
 
 	router := Router{
-		nameToReplicaset: map[string]*Replicaset{
-			alreadyExistingRsName: {conn: nil},
-		},
 		cfg: Config{
 			Loggerf: emptyLogfProvider,
 		},
 	}
+	_ = router.swapNameToReplicaset(nil, &map[string]*Replicaset{
+		alreadyExistingRsName: {conn: nil},
+	})
 
 	// Test that such replicaset already exists
 	rsInfo := ReplicasetInfo{}

vshard.go

@@ -3,7 +3,6 @@ package vshard_router //nolint:revive
 import (
 	"context"
 	"fmt"
-	"sync"
 	"sync/atomic"
 	"time"
 
@@ -30,15 +29,13 @@ type routeMap = []atomic.Pointer[Replicaset]
 type Router struct {
 	cfg Config
 
-	// nameToReplicasetMutex guards not the map itself, but the variable idToReplicaset.
-	// nameToReplicaset is an immutable object by our convention.
+	// nameToReplicaset is an atomic pointer, that points to an immutable object by our convention.
 	// Whenever we add or remove a replicaset, we create a new map object.
-	// nameToReplicaset can be modified only by TopologyController methods.
+	// Object under nameToReplicaset can be replaced to a new one only by TopologyController methods.
 	// Assuming that we rarely add or remove some replicaset,
 	// it should be the simplest and most efficient way of handling concurrent access.
 	// Additionally, we can safely iterate over a map because it never changes.
-	nameToReplicasetMutex sync.RWMutex
-	nameToReplicaset      map[string]*Replicaset
+	nameToReplicaset atomic.Pointer[map[string]*Replicaset]
 
 	routeMap atomic.Pointer[routeMap]
 
@@ -204,10 +201,10 @@ func NewRouter(ctx context.Context, cfg Config) (*Router, error) {
 	}
 
 	router := &Router{
-		cfg:              cfg,
-		nameToReplicaset: make(map[string]*Replicaset),
+		cfg: cfg,
 	}
 	router.setEmptyRouteMap()
+	router.setEmptyNameToReplicaset()
 
 	err = cfg.TopologyProvider.Init(router.Topology())
 	if err != nil {