From c1c0c77ddcc7ba8ad003ab3ac8f72c5a82aa755e Mon Sep 17 00:00:00 2001
From: Jakub Zych <jakub@keios.eu>
Date: Fri, 23 Oct 2026 15:56:57 +0200
Subject: [PATCH] fix(pier): reserve session slots atomically and build servers
 outside router locks

- an unexported reserver (Reserve, Release) next to updater: the memory
  store reserves under its mutex, the Redis store in one Lua script
  (prune, SCARD, SET PX plus SADD); serveNew reserves instead of Count
  then Put and releases on a failed connect or initialize (A-WR-03)
- serveNew, lookup's rehydration and the stateless getServer build the
  SDK server with no router-wide lock held and insert or adopt under the
  lock afterwards, so a slow guard or database check no longer stalls
  the instance (A-WR-07)
---
 modules/pier/session.go      | 99 +++++++++++++++++++++++++-----------
 modules/pier/stateless.go    | 39 +++++++++++---
 modules/pier/store.go        | 11 ++++
 modules/pier/store_memory.go | 26 ++++++++++
 modules/pier/store_redis.go  | 84 ++++++++++++++++++++++++++++++
 5 files changed, 223 insertions(+), 36 deletions(-)

diff --git a/modules/pier/session.go b/modules/pier/session.go
index a153200c..5e2289b5 100644
--- a/modules/pier/session.go
+++ b/modules/pier/session.go
@@ -136,43 +136,76 @@ func (sr *sessionRouter) serveNew(w http.ResponseWriter, r *http.Request, id Ide
 		return
 	}
 	ctx := r.Context()
-	n, err := sr.p.store.Count(ctx, id.Binding)
-	if err != nil {
-		sr.p.storeUnavailable(w, err)
-		return
-	}
-	if n >= sr.p.cfg.MaxPerBinding {
-		writeJSON(w, http.StatusTooManyRequests, map[string]any{"error": "too_many_sessions", "max_per_binding": sr.p.cfg.MaxPerBinding})
-		return
-	}
 	sid := rand.Text()
-	sr.mu.Lock()
-	ls, err := sr.connect(ctx, sid, id, nil)
-	if err == nil {
-		sr.local[sid] = ls
+	rec := Record{Binding: id.Binding, Created: time.Now().UTC()}
+	// The slot is reserved atomically before the server is built
+	// (A-WR-03); a store without Reserve falls back to Count then Put.
+	res, reserves := sr.p.store.(reserver)
+	if reserves {
+		ok, err := res.Reserve(ctx, id.Binding, sid, rec, sr.p.cfg.MaxPerBinding, sr.p.cfg.IdleTTL)
+		if err != nil {
+			sr.p.storeUnavailable(w, err)
+			return
+		}
+		if !ok {
+			sr.tooMany(w)
+			return
+		}
+	} else {
+		n, err := sr.p.store.Count(ctx, id.Binding)
+		if err != nil {
+			sr.p.storeUnavailable(w, err)
+			return
+		}
+		if n >= sr.p.cfg.MaxPerBinding {
+			sr.tooMany(w)
+			return
+		}
 	}
-	sr.mu.Unlock()
+	release := func() {
+		if reserves {
+			_ = res.Release(context.WithoutCancel(ctx), id.Binding, sid)
+		} else {
+			_ = sr.p.store.Delete(context.WithoutCancel(ctx), sid)
+		}
+	}
+	// The server is built with no router lock held: its guard checks may
+	// query a database (A-WR-07). The id is new, so nothing else holds it.
+	ls, err := sr.connect(ctx, sid, id, nil)
 	if err != nil {
+		release()
 		logger().Error("pier: session connect failed", "error", err)
 		http.Error(w, "failed connection", http.StatusInternalServerError)
 		return
 	}
-	if err := sr.p.store.Put(ctx, sid, Record{Binding: id.Binding, Created: ls.created}, sr.p.cfg.IdleTTL); err != nil {
-		sr.closeLocal(sid)
-		sr.p.storeUnavailable(w, err)
-		return
+	ls.created = rec.Created
+	sr.mu.Lock()
+	sr.local[sid] = ls
+	sr.mu.Unlock()
+	if !reserves {
+		if err := sr.p.store.Put(ctx, sid, rec, sr.p.cfg.IdleTTL); err != nil {
+			sr.closeLocal(sid)
+			sr.p.storeUnavailable(w, err)
+			return
+		}
 	}
 	logSession(logger(), "open", sid, id.Binding)
 	ls.transport.ServeHTTP(w, r)
 	if ls.ss.InitializeParams() == nil {
-		// initialize failed: nothing to keep (the SDK cleans up the same way).
-		_ = sr.p.store.Delete(ctx, sid)
+		// initialize failed: nothing to keep (the SDK cleans up the same
+		// way), and the reserved slot frees.
+		release()
 		sr.closeLocal(sid)
 		return
 	}
 	sr.persist(ctx, ls)
 }
 
+// tooMany answers 429 for an initialize past mcp.sessions.max_per_binding.
+func (sr *sessionRouter) tooMany(w http.ResponseWriter) {
+	writeJSON(w, http.StatusTooManyRequests, map[string]any{"error": "too_many_sessions", "max_per_binding": sr.p.cfg.MaxPerBinding})
+}
+
 // serveEphemeral answers one request outside any session with the state
 // the SDK synthesises for a pre-initialize request.
 func (sr *sessionRouter) serveEphemeral(w http.ResponseWriter, r *http.Request, id Identity) {
@@ -259,33 +292,41 @@ func (sr *sessionRouter) lookup(w http.ResponseWriter, r *http.Request, id Ident
 	}
 	sr.mu.Lock()
 	ls := sr.local[sid]
+	sr.mu.Unlock()
 	if ls == nil {
 		if len(rec.State) == 0 {
-			sr.mu.Unlock()
 			http.Error(w, "session not found", http.StatusNotFound)
 			return nil, false
 		}
 		var state mcp.ServerSessionState
 		if err := json.Unmarshal(rec.State, &state); err != nil || state.InitializeParams == nil {
-			sr.mu.Unlock()
 			http.Error(w, "session not found", http.StatusNotFound)
 			return nil, false
 		}
-		ls, err = sr.connect(ctx, sid, id, &state)
+		// Rehydrate with no router lock held (A-WR-07), then insert unless
+		// a concurrent request of the same session got there first.
+		built, err := sr.connect(ctx, sid, id, &state)
 		if err != nil {
-			sr.mu.Unlock()
 			logger().Error("pier: session rehydrate failed", "error", err)
 			http.Error(w, "failed connection", http.StatusInternalServerError)
 			return nil, false
 		}
-		ls.created = rec.Created
+		built.created = rec.Created
 		if state.InitializedParams != nil {
-			ls.persisted.Store(true)
+			built.persisted.Store(true)
+		}
+		sr.mu.Lock()
+		if cur := sr.local[sid]; cur != nil {
+			sr.mu.Unlock()
+			_ = built.ss.Close()
+			ls = cur
+		} else {
+			sr.local[sid] = built
+			sr.mu.Unlock()
+			ls = built
+			logSession(logger(), "rehydrate", sid, id.Binding)
 		}
-		sr.local[sid] = ls
-		logSession(logger(), "rehydrate", sid, id.Binding)
 	}
-	sr.mu.Unlock()
 	ls.mu.Lock()
 	ls.identity = id
 	ls.mu.Unlock()
diff --git a/modules/pier/stateless.go b/modules/pier/stateless.go
index c020ad65..dc294463 100644
--- a/modules/pier/stateless.go
+++ b/modules/pier/stateless.go
@@ -82,17 +82,22 @@ func (sp *statelessPath) getServer(r *http.Request) *mcp.Server {
 	if !ok {
 		return nil
 	}
-	sp.mu.Lock()
-	defer sp.mu.Unlock()
-	if e, ok := sp.entries[id.Binding]; ok {
-		sp.order.MoveToFront(e.elem)
-		e.mu.Lock()
-		e.identity = id
-		e.mu.Unlock()
+	if e := sp.hit(id.Binding); e != nil {
+		sp.refresh(e, id)
 		return e.server
 	}
+	// Built with no lock held: the guard checks may query a database
+	// (A-WR-07). A concurrent request of the same binding may have
+	// inserted an entry meanwhile; that one is adopted.
 	srv, tools, prompts := sp.p.newServer(r.Context(), id, nil)
 	e := &statelessEntry{binding: id.Binding, server: srv, identity: id, tools: tools, prompts: prompts}
+	sp.mu.Lock()
+	if cur, ok := sp.entries[id.Binding]; ok {
+		sp.order.MoveToFront(cur.elem)
+		sp.mu.Unlock()
+		sp.refresh(cur, id)
+		return cur.server
+	}
 	e.elem = sp.order.PushFront(e)
 	sp.entries[id.Binding] = e
 	for sp.order.Len() > sp.capacity {
@@ -104,9 +109,29 @@ func (sp *statelessPath) getServer(r *http.Request) *mcp.Server {
 		sp.order.Remove(last)
 		delete(sp.entries, old.binding)
 	}
+	sp.mu.Unlock()
 	return srv
 }
 
+// hit answers the entry of a binding, most recently used now, or nil.
+func (sp *statelessPath) hit(binding string) *statelessEntry {
+	sp.mu.Lock()
+	defer sp.mu.Unlock()
+	e, ok := sp.entries[binding]
+	if !ok {
+		return nil
+	}
+	sp.order.MoveToFront(e.elem)
+	return e
+}
+
+// refresh stores the request's identity on an entry.
+func (sp *statelessPath) refresh(e *statelessEntry, id Identity) {
+	e.mu.Lock()
+	e.identity = id
+	e.mu.Unlock()
+}
+
 // matching lists the entries whose identity includes the subject.
 func (sp *statelessPath) matching(guard string, userID uint) []*statelessEntry {
 	sp.mu.Lock()
diff --git a/modules/pier/store.go b/modules/pier/store.go
index 136d4e3a..39092292 100644
--- a/modules/pier/store.go
+++ b/modules/pier/store.go
@@ -52,3 +52,14 @@ type Store interface {
 type updater interface {
 	Update(ctx context.Context, id string, rec Record, ttl time.Duration) (bool, error)
 }
+
+// reserver is implemented by the stores pier ships: Reserve is the atomic
+// form of Count then Put (A-WR-03). It writes rec for id with the TTL and
+// adds id to the binding's sessions only when the binding holds fewer than
+// max live sessions, answering false (nothing written) at the cap, so
+// concurrent initializes of one binding cannot exceed it. Release drops a
+// reservation whose session never opened (a failed connect or initialize).
+type reserver interface {
+	Reserve(ctx context.Context, binding, id string, rec Record, max int, ttl time.Duration) (bool, error)
+	Release(ctx context.Context, binding, id string) error
+}
diff --git a/modules/pier/store_memory.go b/modules/pier/store_memory.go
index fcf0ae6a..1a582897 100644
--- a/modules/pier/store_memory.go
+++ b/modules/pier/store_memory.go
@@ -98,6 +98,32 @@ func (m *memoryStore) Count(_ context.Context, binding string) (int, error) {
 	return n, nil
 }
 
+// Reserve counts the binding's live records and writes rec under one lock.
+func (m *memoryStore) Reserve(_ context.Context, binding, id string, rec Record, max int, ttl time.Duration) (bool, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.sweepLocked()
+	n := 0
+	for _, r := range m.records {
+		if r.rec.Binding == binding {
+			n++
+		}
+	}
+	if n >= max {
+		return false, nil
+	}
+	m.records[id] = memoryRecord{rec: rec, expires: m.now().Add(ttl)}
+	return true, nil
+}
+
+// Release drops a reservation.
+func (m *memoryStore) Release(_ context.Context, _ string, id string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.records, id)
+	return nil
+}
+
 func (m *memoryStore) Publish(_ context.Context, ev Event) error {
 	m.mu.Lock()
 	subs := append([]func(Event){}, m.subs...)
diff --git a/modules/pier/store_redis.go b/modules/pier/store_redis.go
index eec5a1b1..8d9cf7ec 100644
--- a/modules/pier/store_redis.go
+++ b/modules/pier/store_redis.go
@@ -136,6 +136,60 @@ func (s *redisStore) Count(ctx context.Context, binding string) (int, error) {
 	return n, nil
 }
 
+// reserveScript is the atomic per-binding reservation: it drops the binding
+// set's members whose record has expired, refuses at max (0) and otherwise
+// writes the record with its PX TTL, adds the id to the set and slides the
+// set's TTL (1). KEYS: the binding set, the record. ARGV: the record key
+// prefix, max, the record JSON, the TTL in milliseconds, the session id.
+var reserveScript = redis.NewScript(`
+local members = redis.call('SMEMBERS', KEYS[1])
+local n = 0
+for _, m in ipairs(members) do
+	if redis.call('EXISTS', ARGV[1] .. m) == 1 then
+		n = n + 1
+	else
+		redis.call('SREM', KEYS[1], m)
+	end
+end
+if n >= tonumber(ARGV[2]) then
+	return 0
+end
+redis.call('SET', KEYS[2], ARGV[3], 'PX', ARGV[4])
+redis.call('SADD', KEYS[1], ARGV[5])
+redis.call('PEXPIRE', KEYS[1], ARGV[4])
+return 1
+`)
+
+// Reserve runs reserveScript, so the count and the write are one step on
+// the server.
+func (s *redisStore) Reserve(ctx context.Context, binding, id string, rec Record, max int, ttl time.Duration) (bool, error) {
+	raw, err := json.Marshal(rec)
+	if err != nil {
+		return false, err
+	}
+	ms := ttl.Milliseconds()
+	if ms < 1 {
+		ms = 1
+	}
+	n, err := reserveScript.Run(ctx, s.client,
+		[]string{s.bindingKey(binding), s.sessionKey(id)},
+		s.prefix+redisSessionKey, max, string(raw), ms, id,
+	).Int()
+	if err != nil {
+		return false, err
+	}
+	return n == 1, nil
+}
+
+// Release drops a reservation: the record and its binding set member.
+func (s *redisStore) Release(ctx context.Context, binding, id string) error {
+	pipe := s.client.TxPipeline()
+	pipe.Del(ctx, s.sessionKey(id))
+	pipe.SRem(ctx, s.bindingKey(binding), id)
+	_, err := pipe.Exec(ctx)
+	return err
+}
+
 func (s *redisStore) Publish(ctx context.Context, ev Event) error {
 	raw, err := json.Marshal(ev)
 	if err != nil {
@@ -231,6 +285,36 @@ func (l *lazyStore) Update(ctx context.Context, id string, rec Record, ttl time.
 	return true, s.Put(ctx, id, rec, ttl)
 }
 
+// Reserve forwards to the opened store's Reserve (Count then Put for a
+// store without one).
+func (l *lazyStore) Reserve(ctx context.Context, binding, id string, rec Record, max int, ttl time.Duration) (bool, error) {
+	s, err := l.get(ctx)
+	if err != nil {
+		return false, err
+	}
+	if r, ok := s.(reserver); ok {
+		return r.Reserve(ctx, binding, id, rec, max, ttl)
+	}
+	n, err := s.Count(ctx, binding)
+	if err != nil || n >= max {
+		return false, err
+	}
+	return true, s.Put(ctx, id, rec, ttl)
+}
+
+// Release forwards to the opened store's Release (Delete for a store
+// without one).
+func (l *lazyStore) Release(ctx context.Context, binding, id string) error {
+	s, err := l.get(ctx)
+	if err != nil {
+		return err
+	}
+	if r, ok := s.(reserver); ok {
+		return r.Release(ctx, binding, id)
+	}
+	return s.Delete(ctx, id)
+}
+
 func (l *lazyStore) Touch(ctx context.Context, id string, ttl time.Duration) error {
 	s, err := l.get(ctx)
 	if err != nil {