Skip to content

Commit 2d52f8a

Browse files
hktitofYuan325
hktitof
and
Yuan325
authored
fix(source/cockroachdb,source/redis): release the handle when a connect attempt fails (#3933)
## Description #3921 made every source release the handle it built when verification fails, but two sources slip through that sweep bcs the handle is built inside a helper instead of inline in Initialize cockroachdb: the retry loop in initCockroachDBConnectionPoolWithRetry creates a fresh pgxpool.Pool per attempt and drops the previous one unclosed on ping failure, so one failed Initialize leaks up to maxRetries+1 pools (default is 5 retries = 6 pools), each keeping its health check goroutine alive. dynamic reload makes it accumulate, one generation per reload attempt redis: the cluster path returns on a failed ForEachShard ping and the standalone path on a failed Ping without closing the client, and both are built with MinIdleConns: 1, so the abandoned pool keeps dialing in the background forever now every failed attempt closes what it built before the next retry or the error return, same shape as #3921, no control flow or error message changes. regression tests cover both, verified they fail without the fix (3+ leaked goroutines on cockroachdb retries, redis dialers that never stop) ## PR Checklist - [ ] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/mcp-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/mcp-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure you have manually reviewed the entire diff before requesting a review - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involve a breaking change Verified locally: ``` go build ./... go test -race ./cmd/... ./internal/... # 405 packages, all pass golangci-lint run ./internal/... ./cmd/... # 0 issues ``` Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
1 parent 98a11f6 commit 2d52f8a

4 files changed

Lines changed: 65 additions & 0 deletions

File tree

‎internal/sources/cockroachdb/cockroachdb.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -507,6 +507,10 @@ func initCockroachDBConnectionPoolWithRetry(ctx context.Context, tracer trace.Tr
507507
return pool, nil
508508
}
509509

510+
if pool != nil {
511+
pool.Close()
512+
}
513+
510514
if attempt < maxRetries {
511515
backoff := baseDelay * time.Duration(math.Pow(2, float64(attempt)))
512516
time.Sleep(backoff)

‎internal/sources/cockroachdb/cockroachdb_test.go‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,13 @@ package cockroachdb
1616

1717
import (
1818
"context"
19+
"runtime"
1920
"strings"
2021
"testing"
22+
"time"
2123

2224
"github.com/goccy/go-yaml"
25+
"go.opentelemetry.io/otel/trace/noop"
2326
)
2427

2528
func TestCockroachDBSourceConfig(t *testing.T) {
@@ -222,3 +225,29 @@ func TestConvertParamMapToRawQuery(t *testing.T) {
222225
func contains(s, substr string) bool {
223226
return strings.Contains(s, substr)
224227
}
228+
229+
// pgxPoolGoroutines counts live goroutines whose stacks still sit inside
230+
// pgx (e.g. a pool's background checks for a client nobody closed).
231+
func pgxPoolGoroutines() int {
232+
buf := make([]byte, 2<<20)
233+
n := runtime.Stack(buf, true)
234+
return strings.Count(string(buf[:n]), "github.com/jackc/pgx/v5")
235+
}
236+
237+
func TestInitCockroachDBConnectionPoolWithRetryReleasesFailedPools(t *testing.T) {
238+
_, err := initCockroachDBConnectionPoolWithRetry(
239+
context.Background(),
240+
noop.NewTracerProvider().Tracer("cockroachdb-test"),
241+
"test-source", "127.0.0.1", "1", "u", "p", "db", nil, 2, time.Millisecond,
242+
)
243+
if err == nil {
244+
t.Fatal("expected the connection to fail against a dead endpoint")
245+
}
246+
// Give any surviving background goroutines a moment to surface.
247+
time.Sleep(500 * time.Millisecond)
248+
if n := pgxPoolGoroutines(); n > 0 {
249+
b := make([]byte, 2<<20)
250+
runtime.Stack(b, true)
251+
t.Fatalf("failed Initialize left %d pgx goroutine(s) running in the background:\n%s", n, b[:2<<20])
252+
}
253+
}

‎internal/sources/redis/redis.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ func initRedisClient(ctx context.Context, r Config) (RedisClient, error) {
124124
return shard.Ping(ctx).Err()
125125
})
126126
if err != nil {
127+
clusterClient.Close()
127128
return nil, fmt.Errorf("unable to connect to redis cluster: %s", err)
128129
}
129130
client = clusterClient
@@ -144,6 +145,7 @@ func initRedisClient(ctx context.Context, r Config) (RedisClient, error) {
144145
})
145146
_, err = standaloneClient.Ping(ctx).Result()
146147
if err != nil {
148+
standaloneClient.Close()
147149
return nil, fmt.Errorf("unable to connect to redis: %s", err)
148150
}
149151
client = standaloneClient

‎internal/sources/redis/redis_test.go‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,10 @@ package redis_test
1616

1717
import (
1818
"context"
19+
"runtime"
1920
"strings"
2021
"testing"
22+
"time"
2123

2224
"github.com/google/go-cmp/cmp"
2325
"github.com/googleapis/mcp-toolbox/internal/server"
@@ -154,3 +156,31 @@ func TestFailParseFromYaml(t *testing.T) {
154156
})
155157
}
156158
}
159+
160+
// goRedisPoolGoroutines counts live goroutines whose stacks still sit inside
161+
// go-redis (e.g. the MinIdleConns dialer of a client nobody closed).
162+
func goRedisPoolGoroutines() int {
163+
buf := make([]byte, 2<<20)
164+
n := runtime.Stack(buf, true)
165+
return strings.Count(string(buf[:n]), "redis/go-redis/v9")
166+
}
167+
168+
func TestInitializeRedisReleasesClientOnFailedPing(t *testing.T) {
169+
cfg := redis.Config{
170+
Name: "test-source",
171+
Type: redis.SourceType,
172+
Address: []string{"127.0.0.1:1"},
173+
}
174+
if _, err := cfg.Initialize(context.Background(), nil); err == nil {
175+
t.Fatal("expected the connection to fail against a dead endpoint")
176+
}
177+
178+
// Give a surviving dialer a moment to park itself, then make sure no
179+
// go-redis goroutine is still trying to connect for the abandoned client.
180+
time.Sleep(1500 * time.Millisecond)
181+
if n := goRedisPoolGoroutines(); n > 0 {
182+
b := make([]byte, 2<<20)
183+
runtime.Stack(b, true)
184+
t.Fatalf("failed Initialize left %d go-redis goroutine(s) dialing in the background:\n%s", n, b[:2<<20])
185+
}
186+
}

0 commit comments

Comments
 (0)