fix(grpcconn): match full and prefixed keys in CloseWorkerConn

[MeiSiristhebest]

Summary

Fixes #1584

In CubeMaster/pkg/cubelet/grpcconn/worker_conn.go, GetWorkerConn stores connection pools under ua + "+" + addr. However, CloseWorkerConn(addr) previously performed a direct lookup with the bare addr string (connPool.cache.Load(addr)).

Because the keys did not match, CloseWorkerConn was a silent no-op. Calls from node_cache (node_cache.go:220/240) upon node down / deletion failed to immediately close active gRPC connections for the removed node.

Changes

1. Updated CloseWorkerConn(addr string) to scan connPool.cache using Range and match entries where key == addr or strings.HasSuffix(key, "+"+addr).
2. Delete matched keys from connPool.cache and close the corresponding connection pools (Pool.Close()).
3. Added if connPool == nil { return } nil guard.
4. Added comprehensive unit tests in CubeMaster/pkg/cubelet/grpcconn/worker_conn_test.go covering matching, nil-pool handling, and uninitialized pool errors (100% target function test coverage).

Verification

• go test -v -cover ./pkg/cubelet/grpcconn/... passes cleanly.
• gofmt -w verified.

[cubesandboxbot]

Concurrency caveat (informational, not blocking): Range deletes matching keys, but a concurrent GetWorkerConn(addr) for the same address whose LoadOrStore lands after this Range passes that key will re-cache a brand-new pool that this call never sees or closes — so "close on node down" is best-effort under in-flight traffic, not a hard guarantee. It is self-healing, since the periodic checkWorkerConn sweep reclaims idle pools (ref == 0 && expired), so a leaked pool is eventually cleaned up. If a strict guarantee were wanted, the close path would need to coordinate with GetWorkerConn (e.g. a per-address mutex or a tombstone check after Range).

[cubesandboxbot]

Minor: these tests mutate the package-global connPool and never restore it, so correctness depends on execution order within this file (Go runs tests in file order and this is currently the only test file in the package, so it works today). It's fragile: if another _test.go file in this package is added later and assumes an initialized pool (e.g. via Init), ordering breaks. Consider saving the prior connPool value and restoring it with t.Cleanup, or initializing the pool in a TestMain. Also note the mock Get() returns a nil grpcpool.Conn, so TestGetWorkerConnCached exercises the cache-hit branch but doesn't validate a usable connection — the "100% coverage" claim in the description overstates the GetWorkerConn create-new-pool path, which is untested.

[cubesandboxbot]

Review: fix(grpcconn): match full and prefixed keys in CloseWorkerConn (#1585)

AI-generated review — no human approval implied.

Verdict

Approve with minor comments. The change correctly fixes a real bug: GetWorkerConn stores pools under UA + "+" + addr, while the old CloseWorkerConn did a bare cache.Load(addr) lookup, so it was always a silent no-op and node-down/node-deletion never closed the worker gRPC connections. The new sync.Map.Range scan that closes every pool whose key is addr or ends with "+addr" addresses the root cause, and the added connPool == nil guard is a genuine improvement (the old code would nil-pointer-panic if called before Init).

I verified against the base tree: all GetWorkerConn callers pass cubelet.GetCubeletAddr(...) = hostIP:grpcPort, which matches the addr passed to CloseWorkerConn from node_cache.go:220/240, so the suffix convention is consistent. I could not find a high-severity correctness issue.

Findings

1. Low — suffix-string matching is fragile and the k == addr branch is dead code (worker_conn.go:96). The match depends entirely on the + delimiter and exact address formatting. Today it cannot collide for valid ip:port keys, but the exact-match branch is unreachable (GetWorkerConn always stores a prefixed key, even +addr for an empty UA), and any future key-format change would silently over- or under-close pools. Consider keying the cache by address (UA in the value) rather than a string-suffix convention.

2. Low (question) — force-close aborts in-flight RPCs (worker_conn.go:99). Pool.Close() immediately resets all underlying gRPC connections. For node-down this is right, but CloseWorkerConn is also called from node deletion, where GracefulStop(maxWaitTime) would let in-flight operations drain. Consistent with existing checkWorkerConn behavior, so not a regression — flagging so the choice is deliberate.

3. Low — test gap: no near-miss negative case (worker_conn_test.go:96). The tests cover exact, prefixed, and different-IP cases but not the dangerous class where one address is a proper prefix/suffix of another stored key's address (e.g. closing 1.2.3.4:1234 must not evict a pool keyed ua+1.2.3.4:12345). A negative test would lock in the delimiter behavior.

Notes

• Closing a pool inside Range while GetWorkerConn may concurrently LoadOrStore leaves a narrow window where a stale pool can survive or a concurrent Get() returns ErrClosed; this is inherent to the existing lock-free design and pre-existing.
• Tests are in-package (can access connPool), restore it via t.Cleanup, and never call Init, so no cleanup goroutine leaks. mockPool satisfies the full grpcpool.Pool interface.

[cubesandboxbot]

The whole match now rests on strings.HasSuffix(k, "+"+addr) — i.e., on the + separator and on the exact address string format. For the current callers this is safe: every key is UA + "+" + addr where addr comes from GetCubeletAddr(HostIP()) (ip:port), and no two valid ip:port strings are proper suffixes of one another (the + delimiter anchors the match). Two notes:

1. The k == addr exact-match branch is actually unreachable today — GetWorkerConn always stores UA + "+" + addr (even with an empty UA the key is +addr), never a bare addr. It's harmless defensive code, but it's misleading: it suggests bare-addr keys exist when they don't.
2. Because the identity lives in the key as a string-suffix convention, any future change to the key format (e.g., a UA containing +, or a hostname address that is a suffix of another) would silently over- or under-close pools. Since the address is the stable entity and the UA is just per-caller, keying the cache by addr alone (holding the UA in the value) — or storing the address in the value and scanning that — would make closure robust to format changes rather than depending on a delimiter.

[cubesandboxbot]

Question (not a regression): Pool.Close() immediately sets current = 0 and resets/closes all underlying grpc.ClientConns, so any in-flight RPC on these connections is aborted. CloseWorkerConn is invoked from both downNodeCache (node down — force-close is clearly right, the node is unreachable) and delNodeCache (node removal), where in-flight operations might still be able to complete. grpcpool.Pool exposes GracefulStop(maxWaitTime) for exactly this drain-then-close case. Note the existing checkWorkerConn cleanup path also uses Close(), so this is consistent with current behavior — flagging only so the choice is deliberate. A stale Get() racing with this close will surface as ErrClosed/ErrorCode_ConnHostFailed; that window is inherent to the lock-free design and pre-existing.

[cubesandboxbot]

The negative case here uses a different IP with the same port — it does not exercise the dangerous class of near-miss keys, i.e. one address being a proper suffix of another stored key's address. The entire fix rests on strings.HasSuffix, so it's worth pinning down that the + delimiter prevents over-closing. Suggest adding a case where the closed addr is a prefix of another cached addr with the same IP but a longer/different port — e.g. cache "uaX+" + "1.2.3.4:12345", call CloseWorkerConn("1.2.3.4:1234"), and assert the ...:12345 pool is untouched (HasSuffix correctly returns false because the key ends in 5). That locks in the delimiter behavior against future key-format changes.