From 0a9ba3687346c4cc58d3933fef3c4da049c1d893 Mon Sep 17 00:00:00 2001
From: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>
Date: Sun, 24 Mar 2024 15:59:35 +0200
Subject: [PATCH] specialized dialPool

Signed-off-by: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>
---
 go/vt/vttablet/grpctmclient/client.go | 44 ++++++++++++++++-----------
 1 file changed, 27 insertions(+), 17 deletions(-)

diff --git a/go/vt/vttablet/grpctmclient/client.go b/go/vt/vttablet/grpctmclient/client.go
index 45ba031b23c..79abfbc9ce3 100644
--- a/go/vt/vttablet/grpctmclient/client.go
+++ b/go/vt/vttablet/grpctmclient/client.go
@@ -104,7 +104,7 @@ type tmc struct {
 }
 
 // rpcClientMap maps an address to a tmc
-type rpcClientMap map[string]*tmc
+type rpcClientMap map[string](chan *tmc)
 
 // grpcClient implements both dialer and poolDialer.
 type grpcClient struct {
@@ -190,24 +190,31 @@ func (client *grpcClient) dialPool(ctx context.Context, dialPoolGroup DialPoolGr
 		}
 		client.mu.Lock()
 		defer client.mu.Unlock()
-		m[addr].cc.Close()
+		for tm := range m[addr] {
+			tm.cc.Close()
+		}
+		close(m[addr])
 		delete(m, addr)
 	}
 
-	if tm, ok := m[addr]; ok {
-		return tm.client, invalidator, nil
-	}
-
-	cc, err := grpcclient.Dial(addr, grpcclient.FailFast(false), opt)
-	if err != nil {
-		return nil, nil, err
-	}
-	tm := &tmc{
-		cc:     cc,
-		client: tabletmanagerservicepb.NewTabletManagerClient(cc),
+	if _, ok := m[addr]; !ok {
+		c := make(chan *tmc, concurrency)
+		for range cap(m[addr]) {
+			cc, err := grpcclient.Dial(addr, grpcclient.FailFast(false), opt)
+			if err != nil {
+				return nil, nil, err
+			}
+			tm := &tmc{
+				cc:     cc,
+				client: tabletmanagerservicepb.NewTabletManagerClient(cc),
+			}
+			c <- tm
+		}
+		m[addr] = c
 	}
-	m[addr] = tm
-
+	c := m[addr]
+	tm := <-c
+	c <- tm
 	return tm.client, invalidator, nil
 }
 
@@ -216,8 +223,11 @@ func (client *grpcClient) Close() {
 	client.mu.Lock()
 	defer client.mu.Unlock()
 	for _, m := range client.rpcClientMaps {
-		for _, tm := range m {
-			tm.cc.Close()
+		for _, c := range m {
+			for tm := range c {
+				tm.cc.Close()
+			}
+			close(c)
 		}
 	}
 	client.rpcClientMaps = nil