[gRPC-TMClient]Optimizes the gRPC tablet manager client by splitting a shared mutex into two independent locks (#18933)

Signed-off-by: siddharth16396 <[email protected]>
Signed-off-by: Arthur Schreiber <[email protected]>
Signed-off-by: siddharth16396 <[email protected]>
Co-authored-by: Arthur Schreiber <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

go/vt/vttablet/grpctmclient/client.go

@@ -103,17 +103,25 @@ type tmc struct {
 	client tabletmanagerservicepb.TabletManagerClient
 }
 
-type addrTmcMap map[string]*tmc
+type tmcEntry struct {
+	once sync.Once
+	tmc  *tmc
+	err  error
+}
+
+type addrTmcMap map[string]*tmcEntry
 
 // grpcClient implements both dialer and poolDialer.
 type grpcClient struct {
 	// This cache of connections is to maximize QPS for ExecuteFetchAs{Dba,App},
 	// CheckThrottler and FullStatus. Note we'll keep the clients open and close them upon Close() only.
 	// But that's OK because usually the tasks that use them are one-purpose only.
-	// The map is protected by the mutex.
-	mu             sync.Mutex
+	// rpcClientMapMu protects rpcClientMap.
+	rpcClientMapMu sync.Mutex
 	rpcClientMap   map[string]chan *tmc
-	rpcDialPoolMap map[DialPoolGroup]addrTmcMap
+	// rpcDialPoolMapMu protects rpcDialPoolMap.
+	rpcDialPoolMapMu sync.Mutex
+	rpcDialPoolMap   map[DialPoolGroup]addrTmcMap
 }
 
 type dialer interface {
@@ -185,25 +193,32 @@ func (client *grpcClient) dialPool(ctx context.Context, tablet *topodatapb.Table
 		return nil, vterrors.FromGRPC(err)
 	}
 
-	client.mu.Lock()
-	if client.rpcClientMap == nil {
-		client.rpcClientMap = make(map[string]chan *tmc)
-	}
-	c, ok := client.rpcClientMap[addr]
-	if !ok {
+	c, isEmpty := func() (chan *tmc, bool) {
+		client.rpcClientMapMu.Lock()
+		defer client.rpcClientMapMu.Unlock()
+
+		if client.rpcClientMap == nil {
+			client.rpcClientMap = make(map[string]chan *tmc)
+		}
+		c, ok := client.rpcClientMap[addr]
+		if ok {
+			return c, false
+		}
+
 		c = make(chan *tmc, concurrency)
 		client.rpcClientMap[addr] = c
-		client.mu.Unlock()
+		return c, true
+	}()
 
+	// If the channel is empty, populate it with connections.
+	if isEmpty {
 		for i := 0; i < cap(c); i++ {
 			tm, err := client.createTmc(ctx, addr, opt)
 			if err != nil {
 				return nil, vterrors.FromGRPC(err)
 			}
 			c <- tm
 		}
-	} else {
-		client.mu.Unlock()
 	}
 
 	result := <-c
@@ -218,44 +233,81 @@ func (client *grpcClient) dialDedicatedPool(ctx context.Context, dialPoolGroup D
 		return nil, nil, err
 	}
 
-	client.mu.Lock()
-	defer client.mu.Unlock()
-	if client.rpcDialPoolMap == nil {
-		client.rpcDialPoolMap = make(map[DialPoolGroup]addrTmcMap)
-	}
-	if _, ok := client.rpcDialPoolMap[dialPoolGroup]; !ok {
-		client.rpcDialPoolMap[dialPoolGroup] = make(addrTmcMap)
-	}
-	m := client.rpcDialPoolMap[dialPoolGroup]
-	if _, ok := m[addr]; !ok {
-		tm, err := client.createTmc(ctx, addr, opt)
-		if err != nil {
-			return nil, nil, err
+	entry := func() *tmcEntry {
+		client.rpcDialPoolMapMu.Lock()
+		defer client.rpcDialPoolMapMu.Unlock()
+
+		if client.rpcDialPoolMap == nil {
+			client.rpcDialPoolMap = make(map[DialPoolGroup]addrTmcMap)
+		}
+		if _, ok := client.rpcDialPoolMap[dialPoolGroup]; !ok {
+			client.rpcDialPoolMap[dialPoolGroup] = make(addrTmcMap)
+		}
+
+		poolEntries := client.rpcDialPoolMap[dialPoolGroup]
+		entry, ok := poolEntries[addr]
+		if ok {
+			return entry
 		}
-		m[addr] = tm
+
+		entry = &tmcEntry{}
+		poolEntries[addr] = entry
+		return entry
+	}()
+
+	// Initialize connection exactly once, without holding the mutex
+	entry.once.Do(func() {
+		entry.tmc, entry.err = client.createTmc(ctx, addr, opt)
+	})
+
+	if entry.err != nil {
+		return nil, nil, entry.err
 	}
+
 	invalidator := func() {
-		client.mu.Lock()
-		defer client.mu.Unlock()
-		if tm := m[addr]; tm != nil && tm.cc != nil {
-			tm.cc.Close()
+		client.rpcDialPoolMapMu.Lock()
+		defer client.rpcDialPoolMapMu.Unlock()
+
+		if entry.tmc != nil && entry.tmc.cc != nil {
+			entry.tmc.cc.Close()
+		}
+
+		if poolEntries, ok := client.rpcDialPoolMap[dialPoolGroup]; ok {
+			delete(poolEntries, addr)
 		}
-		delete(m, addr)
 	}
-	return m[addr].client, invalidator, nil
+	return entry.tmc.client, invalidator, nil
 }
 
 // Close is part of the tmclient.TabletManagerClient interface.
 func (client *grpcClient) Close() {
-	client.mu.Lock()
-	defer client.mu.Unlock()
-	for _, c := range client.rpcClientMap {
-		close(c)
-		for ch := range c {
-			ch.cc.Close()
+	func() {
+		client.rpcClientMapMu.Lock()
+		defer client.rpcClientMapMu.Unlock()
+
+		for _, c := range client.rpcClientMap {
+			close(c)
+			for ch := range c {
+				ch.cc.Close()
+			}
 		}
-	}
-	client.rpcClientMap = nil
+		client.rpcClientMap = nil
+	}()
+
+	// Close dedicated pools
+	func() {
+		client.rpcDialPoolMapMu.Lock()
+		defer client.rpcDialPoolMapMu.Unlock()
+
+		for _, addrMap := range client.rpcDialPoolMap {
+			for _, tm := range addrMap {
+				if tm != nil && tm.tmc != nil && tm.tmc.cc != nil {
+					tm.tmc.cc.Close()
+				}
+			}
+		}
+		client.rpcDialPoolMap = nil
+	}()
 }
 
 //

go/vt/vttablet/grpctmclient/client_test.go

@@ -65,11 +65,12 @@ func TestDialDedicatedPool(t *testing.T) {
 		assert.NotEmpty(t, rpcClient.rpcDialPoolMap[dialPoolGroupThrottler])
 		assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupVTOrc])
 
-		c := rpcClient.rpcDialPoolMap[dialPoolGroupThrottler][addr]
-		assert.NotNil(t, c)
-		assert.Contains(t, []connectivity.State{connectivity.Connecting, connectivity.TransientFailure}, c.cc.GetState())
+		entry := rpcClient.rpcDialPoolMap[dialPoolGroupThrottler][addr]
+		assert.NotNil(t, entry)
+		assert.NotNil(t, entry.tmc)
+		assert.Contains(t, []connectivity.State{connectivity.Connecting, connectivity.TransientFailure}, entry.tmc.cc.GetState())
 
-		cachedTmc = c
+		cachedTmc = entry.tmc
 	})
 
 	t.Run("CheckThrottler", func(t *testing.T) {
@@ -145,22 +146,31 @@ func TestDialPool(t *testing.T) {
 		_, err := client.CheckThrottler(ctx, tablet, req)
 		assert.Error(t, err)
 	})
+
 	t.Run("post throttler maps", func(t *testing.T) {
 		rpcClient, ok := client.dialer.(*grpcClient)
 		require.True(t, ok)
 
-		rpcClient.mu.Lock()
-		defer rpcClient.mu.Unlock()
+		func() {
+			rpcClient.rpcDialPoolMapMu.Lock()
+			defer rpcClient.rpcDialPoolMapMu.Unlock()
 
-		assert.NotEmpty(t, rpcClient.rpcDialPoolMap)
-		assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupThrottler])
-		assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupVTOrc])
+			assert.NotEmpty(t, rpcClient.rpcDialPoolMap)
+			assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupThrottler])
+			assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupVTOrc])
+		}()
 
-		assert.NotEmpty(t, rpcClient.rpcClientMap)
-		assert.NotEmpty(t, rpcClient.rpcClientMap[addr])
+		func() {
+			rpcClient.rpcClientMapMu.Lock()
+			defer rpcClient.rpcClientMapMu.Unlock()
+
+			assert.NotEmpty(t, rpcClient.rpcClientMap)
+			assert.NotEmpty(t, rpcClient.rpcClientMap[addr])
+		}()
 
 		assert.Contains(t, []connectivity.State{connectivity.Connecting, connectivity.TransientFailure}, cachedTmc.cc.GetState())
 	})
+
 	t.Run("ExecuteFetchAsDba", func(t *testing.T) {
 		ctx, cancel := context.WithTimeout(ctx, time.Second)
 		defer cancel()
@@ -174,16 +184,22 @@ func TestDialPool(t *testing.T) {
 		rpcClient, ok := client.dialer.(*grpcClient)
 		require.True(t, ok)
 
-		rpcClient.mu.Lock()
-		defer rpcClient.mu.Unlock()
-
-		assert.NotEmpty(t, rpcClient.rpcDialPoolMap)
-		assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupThrottler])
-		assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupVTOrc])
-
-		// The default pools are unaffected. Invalidator does not run, connections are not closed.
-		assert.NotEmpty(t, rpcClient.rpcClientMap)
-		assert.NotEmpty(t, rpcClient.rpcClientMap[addr])
+		func() {
+			rpcClient.rpcDialPoolMapMu.Lock()
+			defer rpcClient.rpcDialPoolMapMu.Unlock()
+
+			assert.NotEmpty(t, rpcClient.rpcDialPoolMap)
+			assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupThrottler])
+			assert.Empty(t, rpcClient.rpcDialPoolMap[dialPoolGroupVTOrc])
+		}()
+
+		func() {
+			rpcClient.rpcClientMapMu.Lock()
+			defer rpcClient.rpcClientMapMu.Unlock()
+			// The default pools are unaffected. Invalidator does not run, connections are not closed.
+			assert.NotEmpty(t, rpcClient.rpcClientMap)
+			assert.NotEmpty(t, rpcClient.rpcClientMap[addr])
+		}()
 
 		assert.NotNil(t, cachedTmc)
 		assert.Contains(t, []connectivity.State{connectivity.Connecting, connectivity.TransientFailure}, cachedTmc.cc.GetState())