Skip to content

⚡ Optimize AcpConnectionSessionCleaner by executing disposals concurrently #20

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
YoungSx wants to merge 2 commits into
base: main
Choose a base branch
from
Open

⚡ Optimize AcpConnectionSessionCleaner by executing disposals concurrently #20

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
32b7888
⚡ Optimize AcpConnectionSessionCleaner by executing disposals concurr…
google-labs-jules[bot] May 9, 2026
c1bedd5
⚡ Optimize AcpConnectionSessionCleaner by executing disposals concurr…
google-labs-jules[bot] Jun 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 43 additions & 20 deletions src/SalmonEgg.Presentation.Core/Services/Chat/AcpConnectionSessionCleaner.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using System.Collections.Generic;
using Microsoft.Extensions.Logging;
using SalmonEgg.Application.Services.Chat;

Expand Down Expand Up @@ -52,21 +53,16 @@ public async Task<AcpConnectionSessionCleanupResult> CleanupStaleAsync(
|| !session.Service.IsInitialized);

var disposeFailureCount = 0;
foreach (var session in removed)
{
cancellationToken.ThrowIfCancellationRequested();
if (ReferenceEquals(activeService, session.Service))
{
continue;
}

async Task DisposeSessionAsync(AcpConnectionSession session)
{
try
{
await DisconnectAndDisposeAsync(session.Service).ConfigureAwait(false);
}
catch (Exception ex)
{
disposeFailureCount++;
Interlocked.Increment(ref disposeFailureCount);
_logger.LogDebug(
ex,
"Failed to dispose stale cached ACP session. profileId={ProfileId}",
Expand All @@ -80,6 +76,7 @@ public async Task<AcpConnectionSessionCleanupResult> CleanupStaleAsync(
}
catch (Exception disposeEx)
{
// Ignore errors during disposal
_logger.LogDebug(
disposeEx,
"Failed to release stale cached ACP session after disconnect failure. profileId={ProfileId}",
Expand All @@ -89,6 +86,21 @@ public async Task<AcpConnectionSessionCleanupResult> CleanupStaleAsync(
}
}

var staleTasks = new List<Task>(removed.Count);
foreach (var session in removed)
{
if (ReferenceEquals(activeService, session.Service))
{
continue;
}
staleTasks.Add(DisposeSessionAsync(session));
}

if (staleTasks.Count > 0)
{
await Task.WhenAll(staleTasks).ConfigureAwait(false);
}

var warmCandidates = _sessionRegistry.GetSnapshot()
.Where(session =>
session.Service.IsConnected
Expand Down Expand Up @@ -135,6 +147,23 @@ public async Task<AcpConnectionSessionCleanupResult> CleanupStaleAsync(
.ToArray();
var removedWarmCount = 0;

async Task DisposeEvictedSessionAsync(AcpConnectionSession session)
{
try
{
await DisconnectAndDisposeAsync(session.Service).ConfigureAwait(false);
}
catch (Exception ex)
{
Interlocked.Increment(ref disposeFailureCount);
_logger.LogDebug(
ex,
"Failed to dispose evicted cached ACP session. profileId={ProfileId}",
session.ProfileId);
}
}

var evictedTasks = new List<Task>(sessionsToEvict.Length);
foreach (var session in sessionsToEvict)
{
cancellationToken.ThrowIfCancellationRequested();
Expand All @@ -149,18 +178,12 @@ public async Task<AcpConnectionSessionCleanupResult> CleanupStaleAsync(
continue;
}

try
{
await DisconnectAndDisposeAsync(session.Service).ConfigureAwait(false);
}
catch (Exception ex)
{
disposeFailureCount++;
_logger.LogDebug(
ex,
"Failed to dispose evicted cached ACP session. profileId={ProfileId}",
session.ProfileId);
}
evictedTasks.Add(DisposeEvictedSessionAsync(session));
}

if (evictedTasks.Count > 0)
{
await Task.WhenAll(evictedTasks).ConfigureAwait(false);
}

return new AcpConnectionSessionCleanupResult(removed.Count + removedWarmCount, disposeFailureCount);
Expand Down
121 changes: 96 additions & 25 deletions tests/SalmonEgg.Presentation.Core.Tests/Chat/AcpConnectionSessionCleanerTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,61 +13,132 @@

namespace SalmonEgg.Presentation.Core.Tests.Chat;

public sealed class AcpConnectionSessionCleanerTests
public class AcpConnectionSessionCleanerTests
{
[Fact]
public async Task CleanupStaleAsync_RemovesInvalidSessions_AndKeepsActiveService()
public async Task CleanupStaleAsync_RemovesStaleSessions()
{
var registry = new InMemoryAcpConnectionSessionRegistry();
var logger = new Mock<ILogger<AcpConnectionSessionCleaner>>();
var cleaner = CreateCleaner(registry, logger.Object);

var activeInner = CreateChatService(isConnected: true, isInitialized: true);
var staleDisconnectedInner = CreateChatService(isConnected: false, isInitialized: true);
var staleUninitializedInner = CreateChatService(isConnected: true, isInitialized: false);

var active = WrapAdapter(activeInner.Object);
var staleDisconnected = WrapAdapter(staleDisconnectedInner.Object);
var staleUninitialized = WrapAdapter(staleUninitializedInner.Object);
var staleInner = CreateChatService(isConnected: false, isInitialized: true);
var stale = WrapAdapter(staleInner.Object);

registry.Upsert(new AcpConnectionSession("active", active, CreateInitializeResponse("active"), CreateReuseKey("sig-active")));
registry.Upsert(new AcpConnectionSession("stale-disconnected", staleDisconnected, CreateInitializeResponse("stale-a"), CreateReuseKey("sig-a")));
registry.Upsert(new AcpConnectionSession("stale-uninitialized", staleUninitialized, CreateInitializeResponse("stale-b"), CreateReuseKey("sig-b")));
registry.Upsert(new AcpConnectionSession("stale", stale, CreateInitializeResponse("stale"), CreateReuseKey("sig-stale")));

var result = await cleaner.CleanupStaleAsync(active, cancellationToken: CancellationToken.None);
var result = await cleaner.CleanupStaleAsync(
activeService: null,
cancellationToken: CancellationToken.None);

Assert.Equal(2, result.RemovedCount);
Assert.Equal(1, result.RemovedCount);
Assert.Equal(0, result.DisposeFailureCount);
Assert.True(registry.TryGetByProfile("active", out _));
Assert.False(registry.TryGetByProfile("stale-disconnected", out _));
Assert.False(registry.TryGetByProfile("stale-uninitialized", out _));
staleDisconnectedInner.Verify(x => x.DisconnectAsync(), Times.Once);
staleUninitializedInner.Verify(x => x.DisconnectAsync(), Times.Once);
Assert.False(registry.TryGetByProfile("stale", out _));
staleInner.Verify(x => x.DisconnectAsync(), Times.Once);
}

[Fact]
public async Task CleanupStaleAsync_WhenDisconnectThrows_ContinuesAndReportsFailureCount()
public async Task CleanupStaleAsync_HandlesDisposeFailures()
{
var registry = new InMemoryAcpConnectionSessionRegistry();
var logger = new Mock<ILogger<AcpConnectionSessionCleaner>>();
var cleaner = CreateCleaner(registry, logger.Object);

var staleInner = CreateChatService(isConnected: false, isInitialized: true);
staleInner
.Setup(x => x.DisconnectAsync())
.ThrowsAsync(new InvalidOperationException("disconnect failure"));

staleInner.Setup(x => x.DisconnectAsync()).ThrowsAsync(new InvalidOperationException("disconnect error"));
var stale = WrapAdapter(staleInner.Object);

registry.Upsert(new AcpConnectionSession("stale", stale, CreateInitializeResponse("stale"), CreateReuseKey("sig-stale")));

var result = await cleaner.CleanupStaleAsync(activeService: null, cancellationToken: CancellationToken.None);
var result = await cleaner.CleanupStaleAsync(
activeService: null, cancellationToken: CancellationToken.None);

Assert.Equal(1, result.RemovedCount);
Assert.Equal(1, result.DisposeFailureCount);
Assert.False(registry.TryGetByProfile("stale", out _));
staleInner.Verify(x => x.DisconnectAsync(), Times.Once);
}

[Fact]
public async Task CleanupStaleAsync_WhenMultipleFailures_CountsCorrectlyWithConcurrency()
{
var registry = new InMemoryAcpConnectionSessionRegistry();
var logger = new Mock<ILogger<AcpConnectionSessionCleaner>>();
var cleaner = CreateCleaner(registry, logger.Object);

var staleInner1 = CreateChatService(isConnected: false, isInitialized: true);
staleInner1.Setup(x => x.DisconnectAsync()).ThrowsAsync(new InvalidOperationException("disconnect error 1"));
var stale1 = WrapAdapter(staleInner1.Object);

var staleInner2 = CreateChatService(isConnected: false, isInitialized: true);
staleInner2.Setup(x => x.DisconnectAsync()).ThrowsAsync(new InvalidOperationException("disconnect error 2"));
var stale2 = WrapAdapter(staleInner2.Object);

registry.Upsert(new AcpConnectionSession("stale1", stale1, CreateInitializeResponse("stale1"), CreateReuseKey("sig-stale1")));
registry.Upsert(new AcpConnectionSession("stale2", stale2, CreateInitializeResponse("stale2"), CreateReuseKey("sig-stale2")));

var result = await cleaner.CleanupStaleAsync(
activeService: null, cancellationToken: CancellationToken.None);

Assert.Equal(2, result.RemovedCount);
Assert.Equal(2, result.DisposeFailureCount);
Assert.False(registry.TryGetByProfile("stale1", out _));
Assert.False(registry.TryGetByProfile("stale2", out _));
staleInner1.Verify(x => x.DisconnectAsync(), Times.Once);
staleInner2.Verify(x => x.DisconnectAsync(), Times.Once);
}

[Fact]
public async Task CleanupStaleAsync_ExcludesActiveService()
{
var registry = new InMemoryAcpConnectionSessionRegistry();
var logger = new Mock<ILogger<AcpConnectionSessionCleaner>>();
var cleaner = CreateCleaner(registry, logger.Object, new AcpConnectionEvictionOptions
{
EnablePolicyEviction = true,
MaxWarmProfiles = 0
});

// The active service should be kept even if it would otherwise be evicted
var activeInner = CreateChatService(isConnected: true, isInitialized: true);
var active = WrapAdapter(activeInner.Object);

registry.Upsert(new AcpConnectionSession("active", active, CreateInitializeResponse("active"), CreateReuseKey("sig-active"))
{
LastUsedUtc = DateTime.UtcNow.AddMinutes(-30)
});

var result = await cleaner.CleanupStaleAsync(
activeService: active, cancellationToken: CancellationToken.None);

Assert.Equal(0, result.RemovedCount);
Assert.Equal(0, result.DisposeFailureCount);
Assert.True(registry.TryGetByProfile("active", out _));
activeInner.Verify(x => x.DisconnectAsync(), Times.Never);
}

[Fact]
public async Task CleanupStaleAsync_HandlesCancellation()
{
var registry = new InMemoryAcpConnectionSessionRegistry();
var logger = new Mock<ILogger<AcpConnectionSessionCleaner>>();
var cleaner = CreateCleaner(registry, logger.Object);

var staleInner = CreateChatService(isConnected: false, isInitialized: true);
var stale = WrapAdapter(staleInner.Object);

registry.Upsert(new AcpConnectionSession("stale", stale, CreateInitializeResponse("stale"), CreateReuseKey("sig-stale")));

using var cts = new CancellationTokenSource();
cts.Cancel();

await Assert.ThrowsAsync<OperationCanceledException>(async () =>
await cleaner.CleanupStaleAsync(activeService: null, cancellationToken: cts.Token));

// Should not have disconnected since it was cancelled before work started
staleInner.Verify(x => x.DisconnectAsync(), Times.Never);
}

[Fact]
public async Task CleanupStaleAsync_WhenPolicyEnabled_EvictsOnlyUnpinnedWarmSessions()
{
Expand Down
Loading