Tunning and memory control
This commit is contained in:
@@ -0,0 +1,506 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
func TestNativeClientMaxConnectionsAndBatchedActiveDelta(t *testing.T) {
|
||||
oldStore := statsStore
|
||||
statsStore = &Store{}
|
||||
defer func() { statsStore = oldStore }()
|
||||
|
||||
const uuid = "11111111-1111-1111-1111-111111111111"
|
||||
m := &XrayManager{
|
||||
nativeQuotaByUUID: map[string]*xrayNativeQuotaState{
|
||||
uuid: {maxConns: 1, generation: 1},
|
||||
},
|
||||
}
|
||||
state := m.nativeQuotaState(uuid)
|
||||
|
||||
release, acquiredState, ok := m.acquireNativeClientConnection(uuid, "user@example")
|
||||
if !ok || release == nil {
|
||||
t.Fatal("first native connection was rejected")
|
||||
}
|
||||
if acquiredState != state {
|
||||
t.Fatal("connection lease did not retain the authenticated policy state")
|
||||
}
|
||||
if _, _, ok := m.acquireNativeClientConnection(uuid, "user@example"); ok {
|
||||
t.Fatal("connection above max_conns was accepted")
|
||||
}
|
||||
|
||||
m.nativeDBMu.Lock()
|
||||
pending := m.nativeActivePending[uuid]
|
||||
m.nativeDBMu.Unlock()
|
||||
if pending.Delta != 1 || !pending.Connected || pending.State != state {
|
||||
t.Fatalf("connect was not queued for batch persistence: %+v", pending)
|
||||
}
|
||||
|
||||
release()
|
||||
release() // idempotent release must not underflow counters.
|
||||
|
||||
m.nativeDBMu.Lock()
|
||||
pending = m.nativeActivePending[uuid]
|
||||
m.nativeDBMu.Unlock()
|
||||
if pending.Delta != 0 || !pending.Connected {
|
||||
t.Fatalf("connect/disconnect batch should net to zero and retain last-active: %+v", pending)
|
||||
}
|
||||
|
||||
state.mu.Lock()
|
||||
active := state.activeConns
|
||||
state.mu.Unlock()
|
||||
if active != 0 {
|
||||
t.Fatalf("active connection count = %d, want 0", active)
|
||||
}
|
||||
|
||||
release2, _, ok := m.acquireNativeClientConnection(uuid, "user@example")
|
||||
if !ok {
|
||||
t.Fatal("slot was not reusable after release")
|
||||
}
|
||||
release2()
|
||||
}
|
||||
|
||||
func TestRemoveNativeQuotaPolicyPrunesPendingMaps(t *testing.T) {
|
||||
m := &XrayManager{
|
||||
nativeQuotaByUUID: map[string]*xrayNativeQuotaState{
|
||||
"gone": {generation: 1},
|
||||
},
|
||||
nativeTrafficPending: map[string]xrayPendingTraffic{
|
||||
"gone": {Uplink: 10},
|
||||
},
|
||||
nativeActivePending: map[string]xrayPendingActive{
|
||||
"gone": {Delta: 1},
|
||||
},
|
||||
}
|
||||
m.removeNativeQuotaPolicy("gone")
|
||||
if m.nativeQuotaState("gone") != nil {
|
||||
t.Fatal("quota policy was not removed")
|
||||
}
|
||||
m.nativeDBMu.Lock()
|
||||
_, trafficExists := m.nativeTrafficPending["gone"]
|
||||
_, activeExists := m.nativeActivePending["gone"]
|
||||
m.nativeDBMu.Unlock()
|
||||
if trafficExists || activeExists {
|
||||
t.Fatal("deleted UUID remained in a pending persistence map")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeCounterIsBoundedAndReleaseIsIdempotent(t *testing.T) {
|
||||
var active atomicInt64ForTest
|
||||
release1, ok := acquireNativeCounter(&active.Int64, 2)
|
||||
if !ok {
|
||||
t.Fatal("first slot rejected")
|
||||
}
|
||||
release2, ok := acquireNativeCounter(&active.Int64, 2)
|
||||
if !ok {
|
||||
t.Fatal("second slot rejected")
|
||||
}
|
||||
if _, ok := acquireNativeCounter(&active.Int64, 2); ok {
|
||||
t.Fatal("slot above limit accepted")
|
||||
}
|
||||
release1()
|
||||
release1()
|
||||
if got := active.Load(); got != 1 {
|
||||
t.Fatalf("active after double release = %d, want 1", got)
|
||||
}
|
||||
release2()
|
||||
if got := active.Load(); got != 0 {
|
||||
t.Fatalf("active after releases = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Embedding keeps the test declaration readable while still passing the exact
|
||||
// atomic.Int64 type required by acquireNativeCounter.
|
||||
type atomicInt64ForTest struct{ Int64 atomic.Int64 }
|
||||
|
||||
func (a *atomicInt64ForTest) Load() int64 { return a.Int64.Load() }
|
||||
|
||||
func TestTrackedNativeConnectionsAreClosedOnShutdown(t *testing.T) {
|
||||
oldAccepting := nativeTransportAccepting.Load()
|
||||
defer nativeTransportAccepting.Store(oldAccepting)
|
||||
|
||||
beginNativeTransportAccepting()
|
||||
before := nativeTransportConnections.Load()
|
||||
serverSide, clientSide := net.Pipe()
|
||||
defer clientSide.Close()
|
||||
wrapped, ok := wrapTrackedNativeTransportConn(serverSide)
|
||||
if !ok {
|
||||
t.Fatal("tracked connection was rejected")
|
||||
}
|
||||
if got := nativeTransportConnections.Load(); got != before+1 {
|
||||
t.Fatalf("transport count = %d, want %d", got, before+1)
|
||||
}
|
||||
|
||||
stopNativeTransportAccepting()
|
||||
closeAllNativeTransportConnections()
|
||||
_ = clientSide.SetReadDeadline(time.Now().Add(time.Second))
|
||||
if _, err := clientSide.Read(make([]byte, 1)); err == nil {
|
||||
t.Fatal("peer remained open after native shutdown")
|
||||
}
|
||||
if err := wrapped.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
|
||||
t.Fatalf("second close returned unexpected error: %v", err)
|
||||
}
|
||||
if got := nativeTransportConnections.Load(); got != before {
|
||||
t.Fatalf("transport count after shutdown = %d, want %d", got, before)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseAllXHTTPSessionsReleasesGlobalSlots(t *testing.T) {
|
||||
oldLimit := nativeTuneXHTTPMaxSessions.Load()
|
||||
nativeTuneXHTTPMaxSessions.Store(8)
|
||||
defer nativeTuneXHTTPMaxSessions.Store(oldLimit)
|
||||
|
||||
before := nativeXHTTPSessions.Load()
|
||||
ib := &nativeInbound{xhttpMaxBufferedPosts: 2}
|
||||
for _, id := range []string{"one", "two"} {
|
||||
if sess := ib.upsertXHTTPSession(httptest.NewRecorder(), id); sess == nil {
|
||||
t.Fatalf("session %q was rejected", id)
|
||||
}
|
||||
}
|
||||
if got := nativeXHTTPSessions.Load(); got != before+2 {
|
||||
t.Fatalf("global XHTTP sessions = %d, want %d", got, before+2)
|
||||
}
|
||||
ib.closeAllXHTTPSessions()
|
||||
if got := nativeXHTTPSessions.Load(); got != before {
|
||||
t.Fatalf("global XHTTP sessions after close = %d, want %d", got, before)
|
||||
}
|
||||
ib.xhttpMu.Lock()
|
||||
remaining := len(ib.xhttpSessions)
|
||||
ib.xhttpMu.Unlock()
|
||||
if remaining != 0 {
|
||||
t.Fatalf("inbound retained %d XHTTP sessions", remaining)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNegativeXHTTPSessionLimitMeansUnlimited(t *testing.T) {
|
||||
old := nativeTuneXHTTPMaxSessions.Load()
|
||||
nativeTuneXHTTPMaxSessions.Store(0)
|
||||
defer nativeTuneXHTTPMaxSessions.Store(old)
|
||||
if got := (&nativeInbound{}).xhttpMaxActiveSessions(); got != 0 {
|
||||
t.Fatalf("unlimited XHTTP session limit normalized to %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeProtocolGuardsRemainFinite(t *testing.T) {
|
||||
oldRequests := nativeTuneMaxXHTTPRequests.Load()
|
||||
defer nativeTuneMaxXHTTPRequests.Store(oldRequests)
|
||||
|
||||
// Zero is the internal representation of an explicitly disabled application
|
||||
// request counter. HTTP/2 must still retain a finite per-connection guard.
|
||||
nativeTuneMaxXHTTPRequests.Store(0)
|
||||
if got := nativeHTTP2MaxConcurrentStreams(); got != defaultNativeHTTP2MaxStreams {
|
||||
t.Fatalf("HTTP/2 stream guard = %d, want %d", got, defaultNativeHTTP2MaxStreams)
|
||||
}
|
||||
|
||||
nativeTuneMaxXHTTPRequests.Store(32)
|
||||
if got := nativeHTTP2MaxConcurrentStreams(); got != 32 {
|
||||
t.Fatalf("HTTP/2 stream guard did not honor lower request cap: %d", got)
|
||||
}
|
||||
if got := nativeMuxMaxSessionLimit(); got != 64 {
|
||||
t.Fatalf("per-transport Mux session guard = %d, want 64", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestXHTTPMetadataLengthIsBoundedBeforeSessionAllocation(t *testing.T) {
|
||||
ib := &nativeInbound{transport: "xhttp", path: "/"}
|
||||
req := httptest.NewRequest("GET", "/"+strings.Repeat("a", nativeXHTTPMaxSessionIDBytes+1), nil)
|
||||
rec := httptest.NewRecorder()
|
||||
ib.ServeHTTP(rec, req)
|
||||
if rec.Code != 400 {
|
||||
t.Fatalf("oversized XHTTP session id status = %d, want 400", rec.Code)
|
||||
}
|
||||
ib.xhttpMu.Lock()
|
||||
sessions := len(ib.xhttpSessions)
|
||||
ib.xhttpMu.Unlock()
|
||||
if sessions != 0 {
|
||||
t.Fatalf("oversized metadata allocated %d sessions", sessions)
|
||||
}
|
||||
}
|
||||
|
||||
func TestXHTTPUploadMemoryIsReleasedOnReadAndClose(t *testing.T) {
|
||||
before := nativeXHTTPBufferedBytes.Load()
|
||||
q := newNativeXHTTPUploadQueue(4, 8)
|
||||
|
||||
lease, ok := acquireNativeXHTTPMemory(8)
|
||||
if !ok {
|
||||
t.Fatal("failed to reserve XHTTP test memory")
|
||||
}
|
||||
lease.shrink(4)
|
||||
if err := q.push(context.Background(), nativeXHTTPPacket{Payload: []byte("test"), Seq: 0}, lease); err != nil {
|
||||
lease.release()
|
||||
t.Fatalf("queue push failed: %v", err)
|
||||
}
|
||||
lease.release() // transferred leases are a no-op for the producer.
|
||||
if got := nativeXHTTPBufferedBytes.Load(); got != before+4 {
|
||||
t.Fatalf("buffered bytes after push = %d, want %d", got, before+4)
|
||||
}
|
||||
|
||||
buf := make([]byte, 4)
|
||||
if n, err := q.Read(buf); err != nil || n != 4 || string(buf) != "test" {
|
||||
t.Fatalf("queue read = (%d, %v, %q), want (4, nil, test)", n, err, string(buf))
|
||||
}
|
||||
if got := nativeXHTTPBufferedBytes.Load(); got != before {
|
||||
t.Fatalf("buffered bytes after read = %d, want %d", got, before)
|
||||
}
|
||||
|
||||
lease, ok = acquireNativeXHTTPMemory(3)
|
||||
if !ok {
|
||||
t.Fatal("failed to reserve second XHTTP test memory")
|
||||
}
|
||||
if err := q.push(context.Background(), nativeXHTTPPacket{Payload: []byte("xyz"), Seq: 2}, lease); err != nil {
|
||||
lease.release()
|
||||
t.Fatalf("second queue push failed: %v", err)
|
||||
}
|
||||
lease.release()
|
||||
q.close()
|
||||
if got := nativeXHTTPBufferedBytes.Load(); got != before {
|
||||
t.Fatalf("buffered bytes after close = %d, want %d", got, before)
|
||||
}
|
||||
}
|
||||
|
||||
func TestXHTTPUploadQueueEnforcesPerSessionByteBudget(t *testing.T) {
|
||||
before := nativeXHTTPBufferedBytes.Load()
|
||||
q := newNativeXHTTPUploadQueue(4, 4)
|
||||
defer q.close()
|
||||
|
||||
lease, ok := acquireNativeXHTTPMemory(5)
|
||||
if !ok {
|
||||
t.Fatal("failed to reserve XHTTP test memory")
|
||||
}
|
||||
defer lease.release()
|
||||
err := q.push(context.Background(), nativeXHTTPPacket{Payload: make([]byte, 5)}, lease)
|
||||
if !errors.Is(err, errNativeXHTTPUploadBufferFull) {
|
||||
t.Fatalf("oversized queue push error = %v, want buffer limit", err)
|
||||
}
|
||||
lease.release()
|
||||
if got := nativeXHTTPBufferedBytes.Load(); got != before {
|
||||
t.Fatalf("rejected payload retained %d bytes, baseline %d", got, before)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeQuotaResetWaitsForInFlightTraffic(t *testing.T) {
|
||||
const uuid = "22222222-2222-2222-2222-222222222222"
|
||||
state := &xrayNativeQuotaState{usedBytes: 123, generation: 1}
|
||||
m := &XrayManager{nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: state}}
|
||||
|
||||
state.trafficMu.RLock()
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
m.resetNativeQuotaUsage(uuid)
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
state.trafficMu.RUnlock()
|
||||
t.Fatal("traffic reset crossed an in-flight writer boundary")
|
||||
case <-time.After(25 * time.Millisecond):
|
||||
}
|
||||
state.trafficMu.RUnlock()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("traffic reset did not complete after writer released")
|
||||
}
|
||||
state.mu.Lock()
|
||||
used, generation := state.usedBytes, state.generation
|
||||
state.mu.Unlock()
|
||||
if used != 0 || generation != 2 {
|
||||
t.Fatalf("reset state = used %d generation %d, want 0/2", used, generation)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeRateWaitCanBeCanceled(t *testing.T) {
|
||||
lim := rate.NewLimiter(1, 1)
|
||||
if !lim.AllowN(time.Now(), 1) {
|
||||
t.Fatal("failed to consume initial limiter token")
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
if err := waitNativeRate(ctx, lim, 1); !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("waitNativeRate error = %v, want context.Canceled", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSHDirtyQueueDoesNotScanInactiveUsers(t *testing.T) {
|
||||
sshTrafficDirtyMu.Lock()
|
||||
old := sshTrafficDirty
|
||||
sshTrafficDirty = make(map[string]*UserState)
|
||||
sshTrafficDirtyMu.Unlock()
|
||||
defer func() {
|
||||
sshTrafficDirtyMu.Lock()
|
||||
sshTrafficDirty = old
|
||||
sshTrafficDirtyMu.Unlock()
|
||||
}()
|
||||
|
||||
active := &UserState{Cfg: UserConfig{Username: "active"}}
|
||||
inactive := &UserState{Cfg: UserConfig{Username: "inactive"}}
|
||||
markSSHUserTrafficDirty(active)
|
||||
dirty := takeSSHUserTrafficDirty()
|
||||
if len(dirty) != 1 || dirty["active"] != active {
|
||||
t.Fatalf("dirty queue = %#v", dirty)
|
||||
}
|
||||
if _, found := dirty[inactive.Cfg.Username]; found {
|
||||
t.Fatal("inactive user appeared in dirty queue")
|
||||
}
|
||||
if next := takeSSHUserTrafficDirty(); len(next) != 0 {
|
||||
t.Fatalf("dirty queue was not drained: %#v", next)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOldNativeConnectionCannotDecrementReplacementAccount(t *testing.T) {
|
||||
oldStore := statsStore
|
||||
statsStore = &Store{}
|
||||
defer func() { statsStore = oldStore }()
|
||||
|
||||
const uuid = "replacement-active-user"
|
||||
oldState := &xrayNativeQuotaState{maxConns: 1, generation: 1}
|
||||
m := &XrayManager{nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: oldState}}
|
||||
|
||||
release, acquiredState, ok := m.acquireNativeClientConnection(uuid, "old@example")
|
||||
if !ok || acquiredState != oldState {
|
||||
t.Fatal("failed to acquire old account connection")
|
||||
}
|
||||
// Discard the old account's successful connect delta so this assertion only
|
||||
// measures what happens when that old connection later disconnects.
|
||||
m.nativeDBMu.Lock()
|
||||
m.nativeActivePending = nil
|
||||
m.nativeDBMu.Unlock()
|
||||
|
||||
newState := &xrayNativeQuotaState{maxConns: 1, generation: 1}
|
||||
m.nativeQuotaMu.Lock()
|
||||
m.nativeQuotaByUUID[uuid] = newState
|
||||
m.nativeQuotaMu.Unlock()
|
||||
|
||||
release()
|
||||
|
||||
m.nativeDBMu.Lock()
|
||||
pending := m.nativeActivePending[uuid]
|
||||
m.nativeDBMu.Unlock()
|
||||
if pending.Delta != 0 || pending.State != nil {
|
||||
t.Fatalf("old disconnect was queued against replacement account: %+v", pending)
|
||||
}
|
||||
newState.mu.Lock()
|
||||
active := newState.activeConns
|
||||
newState.mu.Unlock()
|
||||
if active != 0 {
|
||||
t.Fatalf("replacement account active count changed to %d", active)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOldNativeTrafficCannotAttachToReplacementAccount(t *testing.T) {
|
||||
oldStore := statsStore
|
||||
statsStore = &Store{}
|
||||
defer func() { statsStore = oldStore }()
|
||||
|
||||
const uuid = "replacement-traffic-user"
|
||||
oldState := &xrayNativeQuotaState{generation: 1}
|
||||
newState := &xrayNativeQuotaState{generation: 1}
|
||||
m := &XrayManager{nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: oldState}}
|
||||
|
||||
meter := newTrafficMeter(uuid, "old@example", true, oldState)
|
||||
meter.n = 1234
|
||||
|
||||
m.nativeQuotaMu.Lock()
|
||||
m.nativeQuotaByUUID[uuid] = newState
|
||||
m.nativeQuotaMu.Unlock()
|
||||
|
||||
oldMgr := xrayMgr
|
||||
xrayMgr = m
|
||||
defer func() { xrayMgr = oldMgr }()
|
||||
meter.flush()
|
||||
|
||||
m.nativeDBMu.Lock()
|
||||
pending := m.nativeTrafficPending[uuid]
|
||||
m.nativeDBMu.Unlock()
|
||||
if pending.Uplink != 0 || pending.Downlink != 0 || pending.State != nil {
|
||||
t.Fatalf("old traffic was queued against replacement account: %+v", pending)
|
||||
}
|
||||
m.statsMu.RLock()
|
||||
stat := m.statsByEmail["old@example"]
|
||||
m.statsMu.RUnlock()
|
||||
if stat.Uplink != 0 || stat.Downlink != 0 {
|
||||
t.Fatalf("old traffic resurfaced in runtime stats: %+v", stat)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeFlusherDropsMismatchedPolicyIdentity(t *testing.T) {
|
||||
oldStore := statsStore
|
||||
statsStore = &Store{}
|
||||
defer func() { statsStore = oldStore }()
|
||||
|
||||
const uuid = "identity-prune-user"
|
||||
oldState := &xrayNativeQuotaState{generation: 1}
|
||||
newState := &xrayNativeQuotaState{generation: 1}
|
||||
m := &XrayManager{
|
||||
nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: newState},
|
||||
nativeTrafficPending: map[string]xrayPendingTraffic{
|
||||
uuid: {Email: "old@example", Uplink: 99, State: oldState},
|
||||
},
|
||||
nativeActivePending: map[string]xrayPendingActive{
|
||||
uuid: {Email: "old@example", Delta: -1, State: oldState},
|
||||
},
|
||||
}
|
||||
m.flushNativeStatsToDB()
|
||||
|
||||
m.nativeDBMu.Lock()
|
||||
defer m.nativeDBMu.Unlock()
|
||||
if len(m.nativeTrafficPending) != 0 || len(m.nativeActivePending) != 0 {
|
||||
t.Fatalf("mismatched pending deltas survived prune: traffic=%v active=%v", m.nativeTrafficPending, m.nativeActivePending)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeMuxFinishRunsOnce(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
s := &nativeMuxSession{
|
||||
closed: make(chan struct{}),
|
||||
uplink: make(chan nativeMuxUplinkItem),
|
||||
onClose: func(*nativeMuxSession) {
|
||||
calls.Add(1)
|
||||
},
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 32; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
s.finish()
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if got := calls.Load(); got != 1 {
|
||||
t.Fatalf("mux onClose called %d times, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
type closeTrackingReader struct {
|
||||
closed atomic.Bool
|
||||
}
|
||||
|
||||
func (r *closeTrackingReader) Read([]byte) (int, error) { return 0, io.EOF }
|
||||
func (r *closeTrackingReader) Close() error {
|
||||
r.closed.Store(true)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestNativeXHTTPQueueCloseClosesQueuedStreamReader(t *testing.T) {
|
||||
q := newNativeXHTTPUploadQueue(1, 1024)
|
||||
r := &closeTrackingReader{}
|
||||
if err := q.push(context.Background(), nativeXHTTPPacket{Reader: r}, nil); err != nil {
|
||||
t.Fatalf("queue stream reader: %v", err)
|
||||
}
|
||||
q.close()
|
||||
if !r.closed.Load() {
|
||||
t.Fatal("queued stream reader was not closed during queue shutdown")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user