fix(management): single-flight transport bundle commands + drop redundant export copy (arch-review S4)

Claude-Session: https://claude.ai/code/session_01MtdgwpEeCUn6cUA5f1LMPj
This commit is contained in:
Joseph Doherty
2026-07-10 06:04:22 -04:00
parent 433c6db4ec
commit d6bad1738e
2 changed files with 215 additions and 86 deletions
@@ -3115,6 +3115,14 @@ public class ManagementActor : ReceiveActor
// Transport bundle operations
// ========================================================================
/// <summary>Single-flight gate for transport bundle commands: a 200 MB import holds
/// ~4 concurrent whole-artifact copies (base64 string, byte[], MemoryStream, JSON envelope),
/// so two concurrent imports on one central node are a realistic OOM path (arch-review S4).
/// Commands queue here instead of running concurrently on the thread pool. Note the actor
/// pipes these handlers to the sender (HandleEnvelope), so without this gate two envelopes
/// received back-to-back run their handlers concurrently.</summary>
private static readonly SemaphoreSlim TransportGate = new(1, 1);
/// <summary>
/// Resolves the per-type name lists in <paramref name="cmd"/> against the
/// repositories, builds an <see cref="ExportSelection"/>, exports the bundle,
@@ -3123,6 +3131,9 @@ public class ManagementActor : ReceiveActor
private static async Task<object?> HandleExportBundle(
IServiceProvider sp, ExportBundleCommand cmd, string username)
{
await TransportGate.WaitAsync();
try
{
var templateRepo = sp.GetRequiredService<ITemplateEngineRepository>();
var externalRepo = sp.GetRequiredService<IExternalSystemRepository>();
var notifRepo = sp.GetRequiredService<INotificationRepository>();
@@ -3213,10 +3224,19 @@ public class ManagementActor : ReceiveActor
selection, user: username, sourceEnvironment: cmd.SourceEnvironment,
passphrase: cmd.Passphrase);
using var ms = new MemoryStream();
await stream.CopyToAsync(ms);
var bytes = ms.ToArray();
return new ExportBundleResult(Convert.ToBase64String(bytes), bytes.Length);
using var ms = new MemoryStream();
await stream.CopyToAsync(ms);
// Encode directly off the MemoryStream's backing buffer rather than
// ms.ToArray() — the latter allocates a second whole-artifact copy.
// The long-term shape is streaming multipart to the caller (as the
// audit export path does, AuditEndpoints.cs:183+); deferred for now.
var base64 = Convert.ToBase64String(ms.GetBuffer().AsSpan(0, (int)ms.Length));
return new ExportBundleResult(base64, (int)ms.Length);
}
finally
{
TransportGate.Release();
}
}
/// <summary>
@@ -3226,33 +3246,41 @@ public class ManagementActor : ReceiveActor
private static async Task<object?> HandlePreviewBundle(
IServiceProvider sp, PreviewBundleCommand cmd)
{
var importer = sp.GetRequiredService<IBundleImporter>();
var bytes = DecodeBundle(cmd.Base64Bundle);
BundleSession session;
await TransportGate.WaitAsync();
try
{
await using var stream = new MemoryStream(bytes);
session = await importer.LoadAsync(stream, cmd.Passphrase);
}
catch (ArgumentException ex)
{
throw new ManagementCommandException(ex.Message);
}
catch (System.Security.Cryptography.CryptographicException)
{
throw new ManagementCommandException("Wrong passphrase or bundle was tampered.");
}
var importer = sp.GetRequiredService<IBundleImporter>();
var bytes = DecodeBundle(cmd.Base64Bundle);
BundleSession session;
try
{
await using var stream = new MemoryStream(bytes);
session = await importer.LoadAsync(stream, cmd.Passphrase);
}
catch (ArgumentException ex)
{
throw new ManagementCommandException(ex.Message);
}
catch (System.Security.Cryptography.CryptographicException)
{
throw new ManagementCommandException("Wrong passphrase or bundle was tampered.");
}
var preview = await importer.PreviewAsync(session.SessionId);
var adds = preview.Items.Count(i => i.Kind == ConflictKind.New);
var mods = preview.Items.Count(i => i.Kind == ConflictKind.Modified);
var ids = preview.Items.Count(i => i.Kind == ConflictKind.Identical);
var blocks = preview.Items.Count(i => i.Kind == ConflictKind.Blocker);
// Surface the required site/connection mappings so an operator
// (CLI or UI) sees which references need resolving before applying.
return new PreviewBundleResult(
preview.Items, adds, mods, ids, blocks,
preview.RequiredSiteMappings, preview.RequiredConnectionMappings);
var preview = await importer.PreviewAsync(session.SessionId);
var adds = preview.Items.Count(i => i.Kind == ConflictKind.New);
var mods = preview.Items.Count(i => i.Kind == ConflictKind.Modified);
var ids = preview.Items.Count(i => i.Kind == ConflictKind.Identical);
var blocks = preview.Items.Count(i => i.Kind == ConflictKind.Blocker);
// Surface the required site/connection mappings so an operator
// (CLI or UI) sees which references need resolving before applying.
return new PreviewBundleResult(
preview.Items, adds, mods, ids, blocks,
preview.RequiredSiteMappings, preview.RequiredConnectionMappings);
}
finally
{
TransportGate.Release();
}
}
/// <summary>
@@ -3263,69 +3291,77 @@ public class ManagementActor : ReceiveActor
private static async Task<object?> HandleImportBundle(
IServiceProvider sp, ImportBundleCommand cmd, string username)
{
var policy = ParseConflictPolicy(cmd.DefaultConflictPolicy);
var importer = sp.GetRequiredService<IBundleImporter>();
var bytes = DecodeBundle(cmd.Base64Bundle);
BundleSession session;
await TransportGate.WaitAsync();
try
{
await using var stream = new MemoryStream(bytes);
session = await importer.LoadAsync(stream, cmd.Passphrase);
}
catch (ArgumentException ex)
{
throw new ManagementCommandException(ex.Message);
}
catch (System.Security.Cryptography.CryptographicException)
{
throw new ManagementCommandException("Wrong passphrase or bundle was tampered.");
}
var policy = ParseConflictPolicy(cmd.DefaultConflictPolicy);
var importer = sp.GetRequiredService<IBundleImporter>();
var bytes = DecodeBundle(cmd.Base64Bundle);
var preview = await importer.PreviewAsync(session.SessionId);
var blockers = preview.Items.Where(i => i.Kind == ConflictKind.Blocker).ToList();
if (blockers.Count > 0)
{
var details = string.Join("; ",
blockers.Select(b => $"{b.Name}: {b.BlockerReason}"));
throw new ManagementCommandException(
$"Bundle has {blockers.Count} blocker(s); import aborted. {details}");
}
// Resolve every required site/connection reference into a
// BundleNameMap. Precedence: an explicit operator spec wins; otherwise
// fall back to the preview's auto-match; otherwise (no match) create-new
// ONLY if the create-missing flag is set, else fail with a clear message.
var nameMap = BuildNameMap(cmd, preview);
// Dedupe by (EntityType, Name) -- the preview can emit multiple rows per
// entity (e.g. one row per modified member of a template), but ApplyAsync
// requires a unique resolution per key. Last-write-wins matches the
// Central UI's TransportImport.BuildDefaultResolutions behavior. For
// DataConnection rows the preview item Name is already site-qualified
// ({site}/{name}), so keying generically by (EntityType, Name) is
// automatically correct — no bare-connection-name special case needed.
var renameStamp = DateTime.UtcNow.ToString("yyyyMMdd-HHmmss");
var resolutionsMap = new Dictionary<(string, string), ImportResolution>();
foreach (var item in preview.Items)
{
var action = item.Kind switch
BundleSession session;
try
{
ConflictKind.New => ResolutionAction.Add,
ConflictKind.Identical => ResolutionAction.Skip,
ConflictKind.Modified => policy,
_ => ResolutionAction.Skip,
};
var renameTo = (item.Kind == ConflictKind.Modified && policy == ResolutionAction.Rename)
? $"{item.Name}-imported-{renameStamp}"
: null;
resolutionsMap[(item.EntityType, item.Name)] = new ImportResolution(
item.EntityType, item.Name, action, renameTo);
}
await using var stream = new MemoryStream(bytes);
session = await importer.LoadAsync(stream, cmd.Passphrase);
}
catch (ArgumentException ex)
{
throw new ManagementCommandException(ex.Message);
}
catch (System.Security.Cryptography.CryptographicException)
{
throw new ManagementCommandException("Wrong passphrase or bundle was tampered.");
}
return await importer.ApplyAsync(
session.SessionId, resolutionsMap.Values.ToList(), username, nameMap: nameMap);
var preview = await importer.PreviewAsync(session.SessionId);
var blockers = preview.Items.Where(i => i.Kind == ConflictKind.Blocker).ToList();
if (blockers.Count > 0)
{
var details = string.Join("; ",
blockers.Select(b => $"{b.Name}: {b.BlockerReason}"));
throw new ManagementCommandException(
$"Bundle has {blockers.Count} blocker(s); import aborted. {details}");
}
// Resolve every required site/connection reference into a
// BundleNameMap. Precedence: an explicit operator spec wins; otherwise
// fall back to the preview's auto-match; otherwise (no match) create-new
// ONLY if the create-missing flag is set, else fail with a clear message.
var nameMap = BuildNameMap(cmd, preview);
// Dedupe by (EntityType, Name) -- the preview can emit multiple rows per
// entity (e.g. one row per modified member of a template), but ApplyAsync
// requires a unique resolution per key. Last-write-wins matches the
// Central UI's TransportImport.BuildDefaultResolutions behavior. For
// DataConnection rows the preview item Name is already site-qualified
// ({site}/{name}), so keying generically by (EntityType, Name) is
// automatically correct — no bare-connection-name special case needed.
var renameStamp = DateTime.UtcNow.ToString("yyyyMMdd-HHmmss");
var resolutionsMap = new Dictionary<(string, string), ImportResolution>();
foreach (var item in preview.Items)
{
var action = item.Kind switch
{
ConflictKind.New => ResolutionAction.Add,
ConflictKind.Identical => ResolutionAction.Skip,
ConflictKind.Modified => policy,
_ => ResolutionAction.Skip,
};
var renameTo = (item.Kind == ConflictKind.Modified && policy == ResolutionAction.Rename)
? $"{item.Name}-imported-{renameStamp}"
: null;
resolutionsMap[(item.EntityType, item.Name)] = new ImportResolution(
item.EntityType, item.Name, action, renameTo);
}
return await importer.ApplyAsync(
session.SessionId, resolutionsMap.Values.ToList(), username, nameMap: nameMap);
}
finally
{
TransportGate.Release();
}
}
/// <summary>
@@ -0,0 +1,93 @@
using Akka.Actor;
using Akka.TestKit.Xunit2;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging.Abstractions;
using NSubstitute;
using ZB.MOM.WW.ScadaBridge.Commons.Interfaces.Transport;
using ZB.MOM.WW.ScadaBridge.Commons.Messages.Management;
using ZB.MOM.WW.ScadaBridge.Commons.Types.Transport;
namespace ZB.MOM.WW.ScadaBridge.ManagementService.Tests;
/// <summary>
/// Verifies the static <c>TransportGate</c> single-flight guard (arch-review S4):
/// transport bundle commands (export/preview/import) queue rather than run
/// concurrently, because each import holds ~4 concurrent whole-artifact copies of
/// the payload and two concurrent large imports on one central node are a
/// realistic OOM path. The gate is exercised through the real import handler so
/// the test tracks the actual concurrency shape of <see cref="IBundleImporter"/>.
/// </summary>
public class TransportGateTests : TestKit, IDisposable
{
private readonly IBundleImporter _importer;
private readonly ServiceCollection _services;
public TransportGateTests()
{
_importer = Substitute.For<IBundleImporter>();
_services = new ServiceCollection();
_services.AddScoped(_ => _importer);
}
private IActorRef CreateActor()
{
var sp = _services.BuildServiceProvider();
return Sys.ActorOf(Props.Create(() => new ManagementActor(
sp, NullLogger<ManagementActor>.Instance)));
}
private static ManagementEnvelope Envelope(object command) =>
new(new AuthenticatedUser("admin", "Admin", new[] { "Administrator" }, Array.Empty<string>()),
command, Guid.NewGuid().ToString("N"));
void IDisposable.Dispose() => Shutdown();
[Fact]
public async Task ImportBundle_ConcurrentCommands_AreSingleFlighted()
{
// LoadAsync blocks on this TCS so the first import stays inside the gate.
var release = new TaskCompletionSource<BundleSession>(
TaskCreationOptions.RunContinuationsAsynchronously);
var loadStarted = new TaskCompletionSource(
TaskCreationOptions.RunContinuationsAsynchronously);
var loadCalls = 0;
var session = new BundleSession { SessionId = Guid.NewGuid() };
_importer.LoadAsync(Arg.Any<Stream>(), Arg.Any<string?>(), Arg.Any<CancellationToken>())
.Returns(_ =>
{
Interlocked.Increment(ref loadCalls);
loadStarted.TrySetResult();
return release.Task;
});
_importer.PreviewAsync(Arg.Any<Guid>(), Arg.Any<CancellationToken>())
.Returns(new ImportPreview(session.SessionId, Array.Empty<ImportPreviewItem>()));
_importer.ApplyAsync(
Arg.Any<Guid>(), Arg.Any<IReadOnlyList<ImportResolution>>(),
Arg.Any<string>(), Arg.Any<CancellationToken>(), Arg.Any<BundleNameMap?>())
.Returns(new ImportResult(
Guid.NewGuid(), 0, 0, 0, 0, Array.Empty<int>(), "corr"));
var actor = CreateActor();
// A trivially-valid base64 payload; the substitute importer never parses it.
var bundle = Convert.ToBase64String(new byte[] { 1, 2, 3, 4 });
actor.Tell(Envelope(new ImportBundleCommand(bundle, null, "skip")));
actor.Tell(Envelope(new ImportBundleCommand(bundle, null, "skip")));
// The first import must reach LoadAsync...
await loadStarted.Task.WaitAsync(TimeSpan.FromSeconds(5));
// ...and the second must be parked at the gate: give it ample time to
// (wrongly) proceed, then assert it did not.
await Task.Delay(750);
Assert.Equal(1, Volatile.Read(ref loadCalls));
// Release the first; the second should now proceed and both should answer.
release.SetResult(session);
ExpectMsg<ManagementSuccess>(TimeSpan.FromSeconds(5));
ExpectMsg<ManagementSuccess>(TimeSpan.FromSeconds(5));
AwaitAssert(() => Assert.Equal(2, Volatile.Read(ref loadCalls)),
TimeSpan.FromSeconds(5));
}
}