From ae2a3a6d5595cef956e899aeca4e508d4de58c1d Mon Sep 17 00:00:00 2001 From: Jihyeon Gim Date: Sun, 6 Sep 2026 23:57:12 +0900 Subject: [PATCH] rdma: acquire session admission credit atomically Checking the publication backlog and taking the admission credit were separate steps, so concurrent registrations could each observe the same headroom and overshoot the session quota together. Both now share one critical section, and a concurrent test pins the behavior: sixteen registrations against a limit of eight with one record pending admit exactly seven. An admission refusal now publishes the same SlowDown error the wire response carries, so operational accounting matches what the client saw, and unregister releases the credit an unfinalized registration was holding so the admission budget cannot leak. --- rdma/rcroutes/ops_linux.go | 34 +++++++++++++++++++++----- rdma/rcroutes/ops_linux_test.go | 42 +++++++++++++++++++++++++++++++++ rdma/rcroutes/routes_linux.go | 8 +++++-- 3 files changed, 76 insertions(+), 8 deletions(-) diff --git a/rdma/rcroutes/ops_linux.go b/rdma/rcroutes/ops_linux.go index 11aa5695..eb8370f0 100644 --- a/rdma/rcroutes/ops_linux.go +++ b/rdma/rcroutes/ops_linux.go @@ -518,11 +518,26 @@ func (t *opsTracker) register(sessionID string, acct auth.Account, // but teardown notifications fire before the audit records // land, so session turnover can queue more records than the // quota bounds. Refusing new sessions while the unpublished - // backlog exceeds the quota turns a stalled sink into + // backlog reaches the quota turns a stalled sink into // latency (the client retries) instead of unbounded memory. - // Unbounded when sessionLimit is unset (tests). - if t.sessionLimit > 0 && t.pubPending.Load() >= int64(t.sessionLimit) { - return errPubBacklog + // The check and the credit acquisition share the session + // mutex so concurrent registrations cannot each observe the + // same headroom and overshoot together. Unbounded when + // sessionLimit is unset (tests). + if t.sessionLimit > 0 { + t.mu.Lock() + full := t.pubPending.Load() >= int64(t.sessionLimit) + if !full { + t.pubPending.Add(1) + } + t.mu.Unlock() + if full { + return errPubBacklog + } + } else { + t.mu.Lock() + t.pubPending.Add(1) + t.mu.Unlock() } acct.Access = strings.Clone(acct.Access) emit := &opsEmitter{ @@ -538,7 +553,6 @@ func (t *opsTracker) register(sessionID string, acct auth.Account, t.mu.Lock() defer t.mu.Unlock() - t.pubPending.Add(1) t.sessions[sessionID] = &sessionRecord{emit: emit} return nil } @@ -554,7 +568,14 @@ func (t *opsTracker) unregister(sessionID string) { } t.mu.Lock() defer t.mu.Unlock() - delete(t.sessions, sessionID) + if _, ok := t.sessions[sessionID]; ok { + delete(t.sessions, sessionID) + // Release the admission credit the registration took: + // no callback will ever publish for this entry, so + // leaving the credit held would permanently shrink the + // admission budget. + t.pubPending.Add(-1) + } } // failOutcome publishes a failed finalization exactly once: when @@ -589,6 +610,7 @@ type reservation struct { emit *opsEmitter gen uint64 } + // reserve marks a session record as owned by its request path: the // teardown callback skips a reserved record because the request // path publishes the real outcome itself. Returns the reservation diff --git a/rdma/rcroutes/ops_linux_test.go b/rdma/rcroutes/ops_linux_test.go index 312622eb..ae1d6e65 100644 --- a/rdma/rcroutes/ops_linux_test.go +++ b/rdma/rcroutes/ops_linux_test.go @@ -19,7 +19,9 @@ package rcroutes import ( "errors" + "fmt" "sync" + "sync/atomic" "testing" "time" @@ -365,3 +367,43 @@ func TestOpsTrackerPublishesExactlyOncePerSession(t *testing.T) { } } } + +func TestOpsTrackerAdmissionAtomicUnderConcurrency(t *testing.T) { + tr := newOpsTracker(8) + // One predecessor record is already pending. + if err := tr.register("s-seed", auth.Account{Access: "a"}, + "r", "b", "k", false, time.Now()); err != nil { + t.Fatalf("seed registration: %v", err) + } + + const rounds = 16 + var wg sync.WaitGroup + var admitted atomic.Int64 + var refused atomic.Int64 + for i := 0; i < rounds; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + err := tr.register(fmt.Sprintf("s-%d", i), auth.Account{Access: "a"}, + "r", "b", "k", false, time.Now()) + if err == nil { + admitted.Add(1) + } else if errors.Is(err, errPubBacklog) { + refused.Add(1) + } else { + t.Errorf("registration %d: unexpected error %v", i, err) + } + }(i) + } + wg.Wait() + + if got := admitted.Load(); got != 7 { + t.Fatalf("admitted %d registrations, want exactly 7 (limit 8, 1 pending)", got) + } + if got := refused.Load(); got != rounds-7 { + t.Fatalf("refused %d registrations, want %d", got, rounds-7) + } + if got := tr.pubPending.Load(); got != 8 { + t.Fatalf("pending credits = %d, want 8", got) + } +} diff --git a/rdma/rcroutes/routes_linux.go b/rdma/rcroutes/routes_linux.go index 382beefe..394d9893 100644 --- a/rdma/rcroutes/routes_linux.go +++ b/rdma/rcroutes/routes_linux.go @@ -255,8 +255,12 @@ func (h *Handler) prepareCore(ctx fiber.Ctx) error { if err := h.ops.register(resp.SessionID, acct, regionFromCtx(ctx), bucket, key, isPut, time.Now()); err != nil { _ = h.svc.FinishPrepare(resp.SessionID, false) - h.ops.publishRequest(ctx, acct, err, bucket, key, isPut) - return s3err.GetAPIError(s3err.ErrSlowDown) + // The audit record carries the same SlowDown the wire + // shows, so operator-side accounting matches what the + // client saw. + apiErr := s3err.GetAPIError(s3err.ErrSlowDown) + h.ops.publishRequest(ctx, acct, apiErr, bucket, key, isPut) + return apiErr } if err := h.svc.FinishPrepare(resp.SessionID, true); err != nil { // The finalization failed. Exactly one publication