From 53c422eb6f4583585f53345e40e8ef9787168041 Mon Sep 17 00:00:00 2001 From: "vitess-bot[bot]" <108069721+vitess-bot[bot]@users.noreply.github.com> Date: Thu, 18 Apr 2024 22:41:52 +0530 Subject: [PATCH] GRPC: Address potential segfault in dedicated connection pooling (#15751) Signed-off-by: Manan Gupta Signed-off-by: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com> Signed-off-by: Matt Lord Co-authored-by: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com> Co-authored-by: Matt Lord --- go/vt/vttablet/grpctmclient/client.go | 8 ++++++-- go/vt/vttablet/grpctmclient/client_test.go | 7 +++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/go/vt/vttablet/grpctmclient/client.go b/go/vt/vttablet/grpctmclient/client.go index fb15b8dbf26..0c651268cda 100644 --- 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 index f842b216d8c..1487303163d 100644 --- 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