feat: add a test and fix the segmentation fault

Signed-off-by: Manan Gupta <[email protected]>
vitessio · Apr 18, 2024 · 632b3d9 · 632b3d9
1 parent 64d9037
commit 632b3d9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/go/vt/vttablet/grpctmclient/client.go b/go/vt/vttablet/grpctmclient/client.go
@@ -235,8 +235,10 @@ 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 _, exists := m[addr]; exists {
+			m[addr].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