Clean up session locking in agent pool manager#19
Clean up session locking in agent pool manager#19MarkCiliaVincenti wants to merge 3 commits intoxyOz-dev:masterfrom
Conversation
📝 WalkthroughWalkthroughReplaced per-session SemaphoreSlims with AsyncKeyedLocker in AgentPoolManager, scoping locks with LockAsync, re-checking and updating existing agents inside the lock, reorganizing recycle/create paths, and updating Dispose. Added AsyncKeyedLock package (v8.0.0) to project dependencies. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
Pre-merge checks❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🔇 Additional comments (1)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
Core/Agents/AgentPoolManager.cs (1)
271-284: Remove the per-semaphore disposal loop;AsyncKeyedLocker<string>does not expose.Valuesand manages its own cleanupThe foreach loop:
foreach (var semaphore in _sessionLocks.Values) { semaphore.Dispose(); }will not compile.
AsyncKeyedLocker<T>does not expose a.Valuesproperty; the internal semaphores are managed privately by the locker. This is a type mismatch error that prevents compilation.Remove the loop entirely and let
_sessionLocks.Dispose()handle its own internal cleanup:- foreach (var semaphore in _sessionLocks.Values) - { - semaphore.Dispose(); - } - _agentPool.Clear(); _sessionAgentMapping.Clear(); _sessionLocks.Dispose();
🧹 Nitpick comments (1)
Core/Agents/AgentPoolManager.cs (1)
7-7: AsyncKeyedLocker integration correctly scopes per‑session critical sectionUsing
AsyncKeyedLocker<string>with:using var _ = await _sessionLocks.LockAsync(sessionId);around the entire
GetOrCreateAgentAsynccritical path is a good fit here: it removes the manual semaphore dictionary bookkeeping and eliminates the race around re‑adding per‑session semaphores, while ensuring only oneGetOrCreateAgentAsyncruns persessionId. The re‑check for an existing agent inside the lock plus the recycle/evict/create flow and event emission all look consistent.Optional: to make intent clearer and match common AsyncKeyedLock examples, you could use an explicit releaser and
await using:await using var sessionLock = await _sessionLocks.LockAsync(sessionId); // ...Functionally this is equivalent here, but slightly more self‑documenting.
Also applies to: 19-19, 46-86
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
Core/Agents/AgentPoolManager.cs(4 hunks)Saturn.csproj(1 hunks)
🔇 Additional comments (1)
Saturn.csproj (1)
31-39: AsyncKeyedLock dependency addition looks appropriate; just confirm version/compatibilityAdding
AsyncKeyedLock7.1.8 is consistent with the newAsyncKeyedLocker<string>usage inAgentPoolManagerand targetsnet8.0fine. Please just double‑check that this exact version is what you validated in CI (and is compatible with your hosting environment and licensing guidelines).
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
Core/Agents/AgentPoolManager.cs (1)
59-69: Recycling path still allows the same idle agent to be assigned to multiple sessions concurrentlyThe
TryRecycleAgent→AssignAgentToSessionpath is still vulnerable to the race called out in the earlier review:
- Two concurrent
GetOrCreateAgentAsynccalls for differentsessionIds both callTryRecycleAgent(configuration)(no cross‑session coordination).- Both can select the same
AgentPoolEntryfrom_agentPool.Valuesbecause it’sIdle/SessionId == nullat the time of enumeration.- Both then call
AssignAgentToSession(recycledAgent.Id, sessionId), leading to the sameAgentinstance being reused for multiple sessions and its state/history being shared.The new per‑session
AsyncKeyedLockerdoesn’t prevent this because it only serializes operations persessionId, not across sessions. The risk is the same as before but is worth fixing while you’re tightening locking semantics.One straightforward fix is to serialize recycling itself so that a given idle agent can only be selected once, e.g.:
+ private readonly object _recycleLock = new(); ... private Agent? TryRecycleAgent(AgentConfiguration targetConfig) { if (!_configuration.EnableRecycling) return null; - - var recyclableAgent = _agentPool.Values - .Where(e => e.State == AgentState.Idle && - e.SessionId == null && - IsConfigurationCompatible(e.Configuration, targetConfig)) - .OrderBy(e => e.LastAccessTime) - .FirstOrDefault(); - - if (recyclableAgent != null) - { - recyclableAgent.Agent.ClearHistory(); - recyclableAgent.LastAccessTime = DateTime.UtcNow; - recyclableAgent.State = AgentState.Active; - recyclableAgent.RecycleCount++; - return recyclableAgent.Agent; - } - - return null; + lock (_recycleLock) + { + var recyclableAgent = _agentPool.Values + .Where(e => e.State == AgentState.Idle && + e.SessionId == null && + IsConfigurationCompatible(e.Configuration, targetConfig)) + .OrderBy(e => e.LastAccessTime) + .FirstOrDefault(); + + if (recyclableAgent == null) + return null; + + // Mark as active while still under the lock so it won't be picked again. + recyclableAgent.Agent.ClearHistory(); + recyclableAgent.LastAccessTime = DateTime.UtcNow; + recyclableAgent.State = AgentState.Active; + recyclableAgent.RecycleCount++; + + return recyclableAgent.Agent; + } }Because
TryRecycleAgentis synchronous and does noawait, using a simplelockhere is safe and avoids widening the existing per‑session async lock. This should eliminate the multi‑session double‑assignment of a recycled agent.
🧹 Nitpick comments (1)
Core/Agents/AgentPoolManager.cs (1)
46-57: Per-session lock cleanup is solid; consider also locking ReleaseAgent to avoid subtle racesThe new per‑session lock:
using var _ = await _sessionLocks.LockAsync(sessionId);plus the in‑lock re‑check of
_sessionAgentMappingnicely fixes the prior session‑semaphore race and ensures only oneGetOrCreateAgentAsyncruns persessionId. The reordering (existing-agent check → recycle → optional eviction → create) is clean and keeps behavior intuitive.One remaining concurrency edge to be aware of:
ReleaseAgentand other methods that touch_sessionAgentMappingfor a givensessionIddo so without taking the same keyed lock. If, in practice,ReleaseAgent(sessionId)can run concurrently withGetOrCreateAgentAsync(sessionId, …), you can get inconsistent state (e.g.,ReleaseAgentmarking an agent idle while a freshGetOrCreatehas just decided to reuse it).If that concurrency is possible, consider also guarding
ReleaseAgent(and any other per-session mutators) with the same keyed lock, using the synchronous API to avoid making it async, e.g.:public void ReleaseAgent(string sessionId) { using (_sessionLocks.Lock(sessionId)) { if (_sessionAgentMapping.TryRemove(sessionId, out var agentId) && _agentPool.TryGetValue(agentId, out var entry)) { entry.SessionId = null; entry.State = AgentState.Idle; entry.LastAccessTime = DateTime.UtcNow; } } }If your higher‑level protocol already guarantees that
ReleaseAgentandGetOrCreateAgentAsyncwon’t overlap for the samesessionId, you can treat this as a future hardening rather than a blocker.Also applies to: 71-76, 79-85
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
Core/Agents/AgentPoolManager.cs(4 hunks)
🔇 Additional comments (1)
Core/Agents/AgentPoolManager.cs (1)
7-7: AsyncKeyedLocker lifecycle and usage look good; just confirm language/version assumptionsUsing
AsyncKeyedLocker<string>as a long-lived field, acquiring locks viaawait _sessionLocks.LockAsync(sessionId)and disposing it inDispose()matches the library’s documented patterns. The zero-arg constructor and directDispose()call are appropriate for a singleton-style manager.Only follow‑ups I’d suggest:
- Ensure the project’s C# language version is high enough that
using var _ = await _sessionLocks.LockAsync(sessionId);is allowed on all target TFMs.- If you later see contention or allocation pressure, you may want to consider wiring
AsyncKeyedLocker<string>via DI with tunedAsyncKeyedLockOptions(pool size, etc.), but that’s not required for correctness here.Also applies to: 19-19, 278-278
|
Are you still interested in this PR @xyOz-dev? |
1 participant
Re-adding the semaphore to the dictionary was risky since it introduced a race condition.
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.