diff --git a/xray/xray.go b/xray/xray.go index 18929ef..9700a3e 100644 --- a/xray/xray.go +++ b/xray/xray.go @@ -183,18 +183,21 @@ func tryCloseAndDelete(url string, srv *Server) { } } -// Close marks the server as draining. The sweeper goroutine will actually close -// the xray instance after DrainTimeout elapses, giving in-flight operations a -// chance to finish cleanly and preventing premature close from leaking goroutines. +// Close synchronously closes the xray instance and removes it from the servers +// map. This prevents goroutine and memory leaks when testing high volumes of +// proxies where the previous delayed-close behavior caused resource accumulation. func Close(proxyURL string) { - startSweeper() mu.Lock() defer mu.Unlock() - i, ok := servers[proxyURL] - if ok && i.DrainedAt.IsZero() { - i.DrainedAt = time.Now() + srv, ok := servers[proxyURL] + if !ok { + return + } + if srv.Instance != nil { + srv.Instance.Close() //nolint: errcheck } + delete(servers, proxyURL) } // CloseAll marks all servers as draining immediately. The sweeper will close diff --git a/xray/xray_test.go b/xray/xray_test.go index 341f7f3..c721a0f 100644 --- a/xray/xray_test.go +++ b/xray/xray_test.go @@ -57,39 +57,29 @@ func TestGetNonExistent(t *testing.T) { } } -func TestCloseRevivesServer(t *testing.T) { +func TestCloseRemovesServer(t *testing.T) { ResetForTest() - DrainTimeout = 50 * time.Millisecond - SweepInterval = 10 * time.Millisecond // Set up an active server. mu.Lock() servers["socks5://127.0.0.1:1080"] = &Server{SocksPort: 1080, DrainedAt: time.Time{}} mu.Unlock() - // Close it — marks DrainedAt non-zero. + // Close it — should immediately remove from map. Close("socks5://127.0.0.1:1080") - // Verify DrainedAt is non-zero (read through map under lock). + // Verify server is removed from map. mu.Lock() - wasZero := servers["socks5://127.0.0.1:1080"].DrainedAt.IsZero() + _, ok := servers["socks5://127.0.0.1:1080"] mu.Unlock() - if wasZero { - t.Error("expected DrainedAt to be non-zero after Close()") + if ok { + t.Error("expected server to be removed from map after Close()") } - // getServer should revive it (reset DrainedAt to zero). + // getServer should return nil since server was removed. got := getServer("socks5://127.0.0.1:1080") - if got == nil { - t.Fatal("expected server after getServer") - } - - // Verify DrainedAt is now zero — read through the map under lock. - mu.Lock() - stillZero := servers["socks5://127.0.0.1:1080"].DrainedAt.IsZero() - mu.Unlock() - if !stillZero { - t.Error("expected DrainedAt to be reset to zero after getServer (revive)") + if got != nil { + t.Fatal("expected nil after getServer on removed server") } } @@ -104,8 +94,8 @@ func TestCloseIdempotent(t *testing.T) { mu.Lock() defer mu.Unlock() - if servers["socks5://127.0.0.1:1080"].DrainedAt.IsZero() { - t.Error("expected DrainedAt non-zero") + if _, ok := servers["socks5://127.0.0.1:1080"]; ok { + t.Error("expected server to be removed from map") } }