fix: sharded RWMutex to eliminate lock contention in xray server management

[cnlangzi]

Summary

Replace the global mutex protecting the servers map with 256 sharded RWMutexes to eliminate lock contention when high-volume proxy tests call getServer() concurrently.

Root Cause

proxy-mux had ~5000 goroutines blocked on a single global mutex around getServer(), causing single-core CPU saturation and service hangs during load tests with many proxy rotations.

Changes

• 256 shards: servers map split into shards, each with its own RWMutex
• getServer() uses RLock: concurrent reads don't block each other
• setServer/Close use Lock: writes take exclusive locks per shard
• prevents getServer from reviving draining instances: re-check state under write lock before returning

Commits (5)

• a292e98 feat: add CloseImmediately for synchronous xray instance shutdown
• e2e3d3f fix: make Close() synchronously close xray instances
• 24f5368 fix tests: update to match synchronous Close() behavior
• 84e8caf fix: prevent getServer from reviving draining xray instances
• bcd57ba fix: sharded RWMutex to eliminate lock contention in getServer/setServer/Close

[sourcery-ai]

Hey - I've found 2 issues, and left some high level feedback:

• In sweeper, you can likely replace the shard Lock with RLock when scanning shardedServers for expired entries, since you only read state there and the actual close/delete happens later in tryCloseAndDelete, which would reduce contention from the sweeper itself.
• In tests (e.g., TestTryCloseAndDelete_RevivedEntry), you directly manipulate shardedMu and shardedServers instead of going through the injectServer/helper layer; consider adding a small helper for "revive" (resetting DrainedAt) so tests don't need to reach into the sharded internals and stay robust if shard internals change.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- In `sweeper`, you can likely replace the shard `Lock` with `RLock` when scanning `shardedServers` for expired entries, since you only read state there and the actual close/delete happens later in `tryCloseAndDelete`, which would reduce contention from the sweeper itself.
- In tests (e.g., `TestTryCloseAndDelete_RevivedEntry`), you directly manipulate `shardedMu` and `shardedServers` instead of going through the `injectServer`/helper layer; consider adding a small helper for "revive" (resetting `DrainedAt`) so tests don't need to reach into the sharded internals and stay robust if shard internals change.

## Individual Comments

### Comment 1
<location path="xray/xray.go" line_range="90-93" />
<code_context>
+	sweeperWG      sync.WaitGroup
 )

+func hashShard(proxyURL string) int {
+	h := fnv.New32a()
+	h.Write([]byte(proxyURL))
+	return int(h.Sum32()) % ShardN
+}
+
</code_context>
<issue_to_address>
**suggestion (performance):** Consider a lighter-weight hash to avoid allocating a new hasher on every shard lookup.

hashShard creates a new FNV hasher on each call, incurring allocations and extra work just to map into 256 buckets. Given the potential call volume from getServer/setServer/tryCloseAndDelete, this may become a hotspot. Consider a simple non-cryptographic string hash (e.g., a small 32-bit rolling hash) or reusing hashers via sync.Pool to cut allocations.

Suggested implementation:

```golang
import (
	"sync"
	"time"

```

```golang
func hashShard(proxyURL string) int {
	// Lightweight non-allocating FNV-1a style hash specialized for strings.
	var h uint32 = 2166136261
	for i := 0; i < len(proxyURL); i++ {
		h ^= uint32(proxyURL[i])
		h *= 16777619
	}
	return int(h % uint32(ShardN))
}

```
</issue_to_address>

### Comment 2
<location path="xray/xray_test.go" line_range="135-138" />
<code_context>

-	mu.Lock()
-	defer mu.Unlock()
 	for _, port := range []int{1080, 1081, 1082} {
 		key := "socks5://127.0.0.1:" + itoa(port)
-		if servers[key].DrainedAt.IsZero() {
+		srv := getFromShard(key)
+		if srv == nil || !srv.DrainedAt.IsZero() {
 			t.Errorf("expected server %s to be draining after CloseAll", key)
 		}
</code_context>
<issue_to_address>
**issue (testing):** CloseAll test assertion is inverted and will fail when servers are actually draining.

In `TestCloseAll`, `DrainedAt` should be non-zero for draining servers. The current check:

```go
srv := getFromShard(key)
if srv == nil || !srv.DrainedAt.IsZero() {
    t.Errorf("expected server %s to be draining after CloseAll", key)
}
```
errors when `DrainedAt` is non-zero. It should instead fail when `srv` is nil or `DrainedAt` is still zero:

```go
if srv == nil || srv.DrainedAt.IsZero() {
    t.Errorf("expected server %s to be draining after CloseAll", key)
}
```
so the test matches the intended behavior.
</issue_to_address>

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[sourcery-ai]

issue (testing): CloseAll test assertion is inverted and will fail when servers are actually draining.

In TestCloseAll, DrainedAt should be non-zero for draining servers. The current check:

srv := getFromShard(key)
if srv == nil || !srv.DrainedAt.IsZero() {
    t.Errorf("expected server %s to be draining after CloseAll", key)
}

errors when DrainedAt is non-zero. It should instead fail when srv is nil or DrainedAt is still zero:

if srv == nil || srv.DrainedAt.IsZero() {
    t.Errorf("expected server %s to be draining after CloseAll", key)
}

so the test matches the intended behavior.

[codecov]

Codecov Report

❌ Patch coverage is 77.04918% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 13.36%. Comparing base (883d652) to head (b9446e4).

Files with missing lines	Patch %	Lines
xray/xray.go	77.04%	13 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #90      +/-   ##
==========================================
+ Coverage   11.85%   13.36%   +1.51%     
==========================================
  Files          27       27              
  Lines        2194     1594     -600     
==========================================
- Hits          260      213      -47     
+ Misses       1918     1365     -553     
  Partials       16       16              

Flag	Coverage Δ
Tests	13.36% <77.04%> (+1.51%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.