GRPC: Address potential segfault in dedicated connection pooling (#15751

) Signed-off-by: Manan Gupta <[email protected]> Signed-off-by: Shlomi Noach <[email protected]> Signed-off-by: Matt Lord <[email protected]> Co-authored-by: Shlomi Noach <[email protected]> Co-authored-by: Matt Lord <[email protected]>
vitessio · Apr 18, 2024 · ff6ca02 · ff6ca02
1 parent 008e02a
commit ff6ca02
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/go/vt/vttablet/grpctmclient/client.go b/go/vt/vttablet/grpctmclient/client.go
@@ -235,8 +235,12 @@ func (client *grpcClient) dialDedicatedPool(ctx context.Context, dialPoolGroup D
 	invalidator := func() {
 		client.mu.Lock()
 		defer client.mu.Unlock()
-		m[addr].cc.Close()
-		delete(m, addr)
+		if tm, ok := m[addr]; ok {
+			if tm != nil && tm.cc != nil {
+				tm.cc.Close()
+			}
+			delete(m, addr)
+		}
 	}
 	return m[addr].client, invalidator, nil
 }

diff --git a/go/vt/vttablet/grpctmclient/client_test.go b/go/vt/vttablet/grpctmclient/client_test.go
@@ -48,6 +48,13 @@ func TestDialDedicatedPool(t *testing.T) {
 		assert.NoError(t, err)
 		assert.NotNil(t, invalidator)
 		assert.NotNil(t, cli)
+		_, invalidatorTwo, err := poolDialer.dialDedicatedPool(ctx, dialPoolGroupThrottler, tablet)
+		assert.NoError(t, err)
+		// Ensure that running both the invalidators doesn't cause any issues.
+		invalidator()
+		invalidatorTwo()
+		_, _, err = poolDialer.dialDedicatedPool(ctx, dialPoolGroupThrottler, tablet)
+		assert.NoError(t, err)
 	})
 
 	var cachedTmc *tmc