Rework server-side per-key coordination with a non-blocking scheduler

[diegosalvi]

Motivation

The coordinator serialized every per-key operation (fetch/get/invalidate/load/put/lock/unlock) with a per-key StampedLock (KeyedLockManager). The write lock was held by a channelsHandlers pool thread for the entire network round-trip of the broadcast (released only in the async reply callback). Under contention on a hot key, pool threads pile up in writeLock()/readLock() until the broadcast completes → pool exhaustion → stalls that spread to unrelated keys. This is the root cause of #188.

This PR replaces that locking with an event-driven, non-blocking per-key scheduler: no pool thread is ever blocked on a lock across I/O, work on a key is serialized by an ordered queue (not a held lock), the design is structurally deadlock-free, and redundant work is coalesced.

The scheduler

New blazingcache.server.KeyedScheduler: a ConcurrentHashMap<RawString, KeySlot> where each key owns a FIFO queue drained asynchronously — submitting an op enqueues it and returns; when the running op completes (via its callback, typically fired by the network layer) the next starts.

• No blocking across I/O. The only lock is a tiny per-slot mutex, held for O(1) bookkeeping, never during I/O and never nested → deadlock-free by construction.
• Readers–writer discipline with a writer barrier. Fetches are SHARED (run concurrently); put/load/invalidate/lock are EXCLUSIVE; FIFO ordering means readers arriving while a writer is queued line up behind it, so writers are never starved.
• Iterative, non-reentrant drain — a run of synchronously-completing ops cannot overflow the stack.
• Aggressive coalescing via submit-time folding under INVALIDATE > PUT > LOAD: a dominant write absorbs the broadcast of contiguous, dominated, not-yet-started predecessors while always preserving their CacheStatus registration, and never folds across a fetch barrier. Separating the droppable broadcast from the always-kept registration is what makes the merge provably equivalent (e.g. load(B) then put(A) keeps B registered so A's PUT still pushes it the value). Plain invalidations arriving while one is in flight attach to it (idempotent) instead of queueing another broadcast.
• Every completion path is idempotent and always resumes the drain. Reply callbacks can fire more than once (send failure then reply timeout, see NettyChannel.sendMessageWithAsyncReply); each op completes at most once (AtomicBoolean) and the drain is resumed in a finally, so a failed/duplicate reply cannot corrupt reader/writer accounting, strand coalesced callers, or wedge a key.

KeyedLockManager, PendingInvalidationsManager, LockID and KeyedLockManagerLockIdTest are removed; the scheduler is the single source of truth.

Application locks — owned by the scheduler, tied to the connection

Locks live entirely in the scheduler (heldToken + owning connection id, published atomically at dequeue). Ownership is keyed on the connection, not just the client id, so both sides are reconnect-safe:

• Grant: granted only if the token is still held and the owning connection is still the current one for that client id; the liveness check is wrapped in try/catch so a failure releases the token instead of leaving the key locked and wedging the drain.
• Release: releaseLocksForConnection(connectionId) (from clientDisconnected) releases the held lock and cancels queued lock requests only for that connection — a late cleanup of a dead connection can no longer release a lock a reconnected sibling legitimately acquired.

This closes the acquire-vs-disconnect race in both directions (grant and release) when a client dies and reconnects with the same id on a new connection.

invalidateByPrefix — single broadcast, no server-side pruning

A single INVALIDATE_BY_PREFIX message per client (fan-out stays O(clients); each client drops its matching keys locally), and the server does not prune the holder registrations. The legacy bulk removal raced an in-flight putEntry on a matching key and could leave a client holding an entry the server had forgotten (stale). Not pruning makes that dangerous direction impossible; only the benign "server still thinks a client holds a key it dropped" remains, which self-heals on the next invalidate/put/fetch of the key, on expiry, or on disconnect. (Routing prefix removal per-key through the scheduler was considered and rejected: it can self-deadlock behind a held, timeout-less application lock on a matching key, and mishandle coalescing.)

A read-path self-heal (prune-on-fetch-miss: when a fetch is routed to a holder that replies it no longer has the entry, prune that holder) was also prototyped and rejected. A fetch runs as a SHARED op, and SHARED is mutually exclusive only with the EXCLUSIVE writers, not with other concurrent fetches on the same key. So the prune can race a second fetch whose fetcher happens to be the very holder being pruned: that fetch legitimately re-registers the holder (and the client caches the value) on its own reply, and the two reply callbacks commit under the CacheStatus lock in arbitrary order. If the prune commits after the re-registration it erases a valid registration, resurrecting the forbidden "client holds the key / server forgot it" state (silent stale reads, non-self-healing for immortal entries). Since the prune acts on already-stale information (the holder answered "absent" earlier, but re-acquired the key meanwhile), there is no safe single-signal fix short of per-registration versioning, so the read path deliberately does not contribute to pruning. This confirms that the server-side over-knowing left by invalidateByPrefix can only be healed by the write path (invalidate/put), expiry, or disconnect.

Client-listener cleanup — connection-identity aware on reconnect

Holder registrations now record the owning connection (connectionForClient, advanced forward-only via merge(Math::max) since connection ids are monotonic). removeClientListeners(clientId, connectionId) removes a disconnected connection's registrations only if that connection still owns them, checked atomically with a concurrent registerKeyForClient under the CacheStatus lock. This closes the reconnect TOCTOU where a dead connection's late cleanup wiped the registrations a reconnected connection had just made. CacheServerEndpoint.connectionClosed likewise unmaps the client only if it still points to the closing connection.

Wire protocol / compatibility

No protocol change — fully backward compatible. No Message type or parameter changed and the client is untouched. The only observable difference is that the lockId returned by lock() is now a small monotonic token instead of a StampedLock stamp, but it travels as the same opaque String in KeyLock and is echoed back verbatim, so it is indistinguishable on the wire. Locks live on the leader only and connections close on leadership change, so no token crosses server versions.

Preserved behavior (invariants covered by existing tests)

Hot-key fetch/invalidate storm without stalls (FetchAndInvalidateStormTest); writer not starved (WriterStarvationTest); invalidations never lost / no resurrection (FetchAndInvalidateHammerTest, ConcurrentFetchAndInvalidationTest); put-global vs load-local visibility and load/put conflicts (LoadAndPutEntryTest, LoadConcurrencyTest); per-key lock counts and LockedEntries JMX (LostFetchMessage*, ManagementStatusMXBeanTest); explicit KeyLock incl. release on disconnect (LockBasicTest, LockLostTest); stuck/errored client doesn't deadlock the coordinator (ApparentlyStuckClientDueToServerSideErrorTest, ErrorOnFetchTest); fetch-source priority (FetchPriorityTest); prefix invalidation (InvalidateByPrefixTest, SimpleEvictMaxMemoryTest).

Testing

• New KeyedSchedulerTest — RW fairness/writer barrier, reader batching, coalescing dominance, in-flight invalidation attach, token bypass + malformed/invalid token, per-op idempotency (shared/exclusive/bypass), lock↔disconnect (release-held, cancel-queued, not-granted-to-disconnected/gone, bypass barrier after unlock, liveness-check-throws), token-0 unlock guard, grant-reply-failure release.
• New CacheStatusReconnectTest — connection-identity listener cleanup and forward-only ownership.
• New InvalidateByPrefixResurrectionTest — 300 rounds of concurrent put + prefix invalidation asserting the safe direction (never "client holds it / server forgot").
• Full blazingcache-core suite (46 test classes) green; reactor build (core + jcache + services) green; quality gates green (apache-rat, spotbugs).

Known, intentionally-accepted trade-offs

• invalidateByPrefix leaves the server over-knowing holder registrations for dropped keys (benign, self-healing) — the deliberate cost of the single-broadcast, deadlock-free design.
• On a reconnect race, a dead connection's now-stale registrations are left as benign phantoms rather than risk wiping the live connection's fresh ones.
• releaseLocksForConnection scans the active-key slots on disconnect (O(active keys)) instead of keeping a per-client index, to avoid reintroducing cross-structure non-atomicity.
• foldDominatedTail walks the dominated tail in more than one pass (correctness-neutral micro-opt).

History

The design converged over several review rounds; the commit history keeps the intermediate steps:

1. Initial rework — replaced the per-key StampedLock with the non-blocking KeyedScheduler (readers-writer, coalescing, app locks).
2. Idempotency & robustness — reply callbacks can double-fire, so every completion path was made at-most-once and the drain always-resumed (try/finally); the fetch reply registration was moved inside its guard.
3. Lock ↔ disconnect (grant side) — added a grant-time liveness check so a lock is not handed to a client that disconnected while its request was queued.
4. Lock ↔ disconnect (release side) + reconnect — moved lock ownership to connection identity on both grant and release, so a dead connection's late cleanup can't release a reconnected sibling's lock.
5. invalidateByPrefix — first routed per-key through the scheduler, then a two-phase variant; both were reverted after review surfaced a self-deadlock behind held locks, early-ACK via coalescing, and a fresh-client wipe. The final design is the single-broadcast, no-server-pruning approach above.
6. Listener cleanup on reconnect — made connection-identity aware and race-free (atomic supersession check, forward-only ownership), and hardened the lock liveness check against a throwing supplier.
7. Read-path pruning evaluated and rejected — a prune-on-fetch-miss self-heal was prototyped, then dropped after an adversarial review confirmed it races a concurrent same-key fetch that re-registers the pruned holder, reintroducing the forbidden "client holds it / server forgot" state (see the invalidateByPrefix section). No code from this step ships.

Each step was validated by the full test suite, the quality gates, and an adversarial code-review pass; the final state has no outstanding review findings.

[NiccoMlt]

Left a couple of comments.

I don't know the codebase well enough to be 100% sure about them, I had some help from Claude, so they might be wrong, resolve the conversations if they are not the case

[NiccoMlt]

@diegosalvi How does this work if the client disconnects "at the wrong time" and then reconnects?

I mean:

1. Client A on connection C1 requests lockKey, so a LOCK operation is enqueued in the slot of the key. Note that the reply for this request will be sent on C1, because CacheServerSideConnection#messageReceived > case Message.TYPE_LOCK_ENTRY captures it in the _channel variable.
2. Connection C1 dies. The server executes CacheServerSideConnection#channelClosed, which removes the connection from the map and then, via releaseLocksForClient(<client A>), releases the locks held by A and cancels A's queued LOCK operations.
3. This cleanup is a one-shot scan, and I don't think it's guaranteed to catch the LOCK from step 1, because message handling and disconnect handling run on different threads with no ordering between them:

• the lockKey message may still be in flight in messageReceived while the cleanup runs, so submitLock enqueues the operation after the scan has already finished, and nobody will ever look for it again;
• even if submitLock runs during the scan, slots.forEach on a ConcurrentHashMap is weakly consistent: a slot created concurrently with the iteration may simply not be visited.

4. Suppose the orphaned LOCK survives and A quickly reconnects with a new connection C2 (same clientId). When the operation reaches the head of the queue, the liveness check at grant time is acceptor.getActualConnectionFromClient(sourceClientId) != null, so it finds C2 and passes.
5. The reply carrying the token is sent on C1, which is dead. NettyChannel.sendReplyMessage on a closed channel discards the message silently without throwing, so the granted && !replied guard in runLock never fires.
6. End state: heldToken is set, but the client never received the token, so it can't unlockKey; and it's connected, so releaseLocksForClient won't run either. The key stalls until A disconnects again.

Does it make sense to you, or am I tripping?

If yes, maybe we could try by making the liveness check compare connection identity instead of mere presence, in CacheServer#lockKey, doing like:

acceptor.getActualConnectionFromClient(sourceClientId) == connAtSubmit

[diegosalvi]

Confirmed, good catch. Fixed in e78da97: the grant-time check now compares connection identity instead of mere presence, as you suggested. lockKey passes the submitting connection and the scheduler grants the lock only if acceptor.getActualConnectionFromClient(sourceClientId) == connection. If the original connection died and the client reconnected with a new one (same clientId), the check fails, the token published at dequeue is released, and the (harmless) error reply goes to the dead channel, so the key can no longer stall. releaseLocksForClient still backstops the case where the held token is already published when the disconnect cleanup scans.

[NiccoMlt]

Thanks, the identity check closes the grant side. But the same asymmetry now exists on the release side: grant compares connection identity, while releaseLocksForClient still releases by clientId. Mirror sequence of the one you just fixed:

1. A's connection C1 dies
2. before C1's channelClosed cleanup is processed, A reconnects as C2
3. A legitimately acquires the lock on K via C2 (the identity check passes, because the request comes from C2, which is the current connection)
4. C1's late cleanup runs releaseLocksForClient("A")
5. it finds lockOwnerClientId == "A" on K and releases it (it has no way to tell C2's lock from a C1 leftover)
6. another client acquires K while A still believes it holds it (the release is silent server-side); A only finds out later that mutual exclusion was broken when its operations with that lockId get rejected and its unlock fails with "no lock held"

Also, could you add the reconnect tests while you're at it?

[diegosalvi]

Done. The application lock is now owned by a specific connection, not just a client id, so grant and release are symmetric:

• submitLock records the owning connection id; the grant-time check verifies it is still the current connection for that client id;
• releaseLocksForConnection(connectionId) (from clientDisconnected) releases/cancels only that connection's locks, so a dead connection's late cleanup cannot release a lock a reconnected sibling legitimately holds.

Reconnect/robustness coverage added in KeyedSchedulerTest (releaseLocksForConnectionIgnoresOtherConnections, lockForAlreadyDisconnectedOwnerIsNotGranted, lockOwnerLivenessCheckThrowingReleasesAndDoesNotWedge). Commits dd1fca8 / 202fc1b.

[NiccoMlt]

@diegosalvi this doesn't go through the scheduler, do we actually like this?

A prefix invalidation can interleave with an in-flight putEntry on a matching key:

1. Client A receives the prefix invalidation first and drops the entry, then receives the putEntry and stores it again, resurrecting the entry
2. server-side, the prefix invalidation already removed A's registration for that key, so A now holds the entry while the server doesn't know it does, so future invalidations of that key are never sent to A, and it serves stale data until the entry expires

[diegosalvi]

Confirmed. It predates this PR (the old // LOCKS ??) but it is a real defect. Fixed across 31f9bb9/e78da97: invalidateByPrefix now runs in two phases.

• Phase 1: for each matching key, an exclusive scheduler op clears the holder registrations (CacheStatus.removeKey) with no network, serialized on that key's slot with any concurrent put/load/fetch.
• Phase 2, only after phase 1 completes: a single INVALIDATE_BY_PREFIX per client (client set captured before phase 1), so the fan-out stays O(clients), not O(matching keys) - important for large prefixes. Running phase 2 after phase 1 also orders any PUT_ENTRY from a concurrent put before this invalidation on the client's channel.

A put/load on a matching key landing between the two phases can leave a transient, self-healing divergence (the server briefly thinks a client still holds a key it dropped) - never the reverse 'client holds it, server does not know' direction that would serve stale data. Documented in the code with a concurrency note.

[NiccoMlt]

Well done, the two-phase design looks correct to me!

Do you think it is possible to add the resurrection reproducer as a test? It's the scenario that motivated the rework.

Also, I think that it is worth documenting in the PR that the legacy prefix invalidation ignored locks entirely, prefix invalidation now waits for in-flight broadcasts and held locks on matching keys (which have no timeout).

[diegosalvi]

This one evolved after more review. I first routed prefix invalidation per-key through the scheduler, but that introduced a self-deadlock behind held (timeout-less) application locks on matching keys, plus early-ACK-via-coalescing and a fresh-client wipe. Final design in 202fc1b: keep a single INVALIDATE_BY_PREFIX broadcast per client and do not prune the server-side holder registrations at all. That makes the dangerous 'client dropped it / server forgot' direction (the one you flagged) impossible - the server never removes on prefix invalidation - leaving only the benign 'server still thinks a client holds a key it dropped', which self-heals on the next invalidate/put/fetch of the key, on expiry, or on disconnect. (So the earlier 'waits for held locks' note no longer applies: it doesn't touch the scheduler.) Added InvalidateByPrefixResurrectionTest - 300 rounds of concurrent put + prefix invalidation asserting the safe direction.

[NiccoMlt]

Thanks, both fixes look right!

Still, I left another couple of comments mostly about finishing the job.

[NiccoMlt]

You never set op.resultKey here, so the SEVERE error logs for fetches print "error running fetch for key null" (L515, and L678 for the lock-bypass path).

I think adding op.resultKey = key; here should be enough.

[diegosalvi]

Fixed in dd1fca8: op.resultKey = key is set in submitFetch, so the fetch (and lock-bypass) error logs no longer print for key null. Thanks!

[NiccoMlt]

Please double-check this comment.

replied stays false only when the callback throws; when the reply fails to reach the client, i.e. when the channel is dead, it doesn't throw
(NettyChannel.sendReplyMessage silently discards on a closed socket), so this guard never fires for it. That case is actually covered by the connection-identity check at grant time plus the disconnect cleanup.

I think we should reword the comment to make this clear.

[diegosalvi]

Reworded in dd1fca8: the comment now makes clear the granted && !replied guard only covers a reply callback that throws; the reply silently dropped on an already-dead channel (NettyChannel discards without throwing) is handled instead by the connection-identity liveness check at grant time plus releaseLocksForConnection on disconnect. In 202fc1b the liveness check itself is also wrapped in try/catch so a throw there can't escape drain() and wedge the key.

[NiccoMlt]

Thanks, the identity check closes the grant side. But the same asymmetry now exists on the release side: grant compares connection identity, while releaseLocksForClient still releases by clientId. Mirror sequence of the one you just fixed:

1. A's connection C1 dies
2. before C1's channelClosed cleanup is processed, A reconnects as C2
3. A legitimately acquires the lock on K via C2 (the identity check passes, because the request comes from C2, which is the current connection)
4. C1's late cleanup runs releaseLocksForClient("A")
5. it finds lockOwnerClientId == "A" on K and releases it (it has no way to tell C2's lock from a C1 leftover)
6. another client acquires K while A still believes it holds it (the release is silent server-side); A only finds out later that mutual exclusion was broken when its operations with that lockId get rejected and its unlock fails with "no lock held"

Also, could you add the reconnect tests while you're at it?

[NiccoMlt]

Well done, the two-phase design looks correct to me!

Do you think it is possible to add the resurrection reproducer as a test? It's the scenario that motivated the rework.

Also, I think that it is worth documenting in the PR that the legacy prefix invalidation ignored locks entirely, prefix invalidation now waits for in-flight broadcasts and held locks on matching keys (which have no timeout).

[NiccoMlt]

Great work, LGTM!