CBDDC/tests/ZB.MOM.WW.CDBBC.E2E.Benchmark.Tests/BenchmarkPeerNode.cs

using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using ZB.MOM.WW.CBDDC.Core.Network;
using ZB.MOM.WW.CBDDC.Network;
using ZB.MOM.WW.CBDDC.Network.Security;
using ZB.MOM.WW.CBDDC.Persistence.Surreal;
using ZB.MOM.WW.CBDDC.Sample.Console;

namespace ZB.MOM.WW.CDBBC.E2E.Benchmark.Tests;

internal sealed class BenchmarkPeerNode : IAsyncDisposable
{
    private readonly ICBDDCNode _node;
    private readonly ServiceProvider _serviceProvider;
    private readonly string _workDir;
    private bool _started;

    private BenchmarkPeerNode(
        ServiceProvider serviceProvider,
        ICBDDCNode node,
        SampleDbContext context,
        string workDir)
    {
        _serviceProvider = serviceProvider;
        _node = node;
        Context = context;
        _workDir = workDir;
    }

    public SampleDbContext Context { get; }

    public static BenchmarkPeerNode Create(
        string nodeId,
        int tcpPort,
        string authToken,
        IReadOnlyList<KnownPeerConfiguration> knownPeers)
    {
        string workDir = Path.Combine(Path.GetTempPath(), $"cbddc-benchmark-{nodeId}-{Guid.NewGuid():N}");
        Directory.CreateDirectory(workDir);

        string dbPath = Path.Combine(workDir, "node.rocksdb");
        string databaseName = nodeId.Replace("-", "_", StringComparison.Ordinal);

        var configurationProvider = new StaticPeerNodeConfigurationProvider(new PeerNodeConfiguration
        {
            NodeId = nodeId,
            TcpPort = tcpPort,
            AuthToken = authToken,
            KnownPeers = knownPeers.ToList(),
            RetryDelayMs = 25,
            RetryAttempts = 5
        });

        var services = new ServiceCollection();
        services.AddLogging(builder => builder.SetMinimumLevel(LogLevel.Warning));
        services.AddSingleton(configurationProvider);
        services.AddSingleton<IPeerNodeConfigurationProvider>(configurationProvider);
        services.AddSingleton<ICBDDCSurrealSchemaInitializer, SampleSurrealSchemaInitializer>();
        services.AddSingleton<SampleDbContext>();

        services.AddCBDDCCore()
            .AddCBDDCSurrealEmbedded<SampleDocumentStore>(_ => new CBDDCSurrealEmbeddedOptions
            {
                Endpoint = "rocksdb://local",
                DatabasePath = dbPath,
                Namespace = "cbddc_benchmark",
                Database = databaseName,
                Cdc = new CBDDCSurrealCdcOptions
                {
                    Enabled = true,
                    ConsumerId = $"{nodeId}-benchmark",
                    PollingInterval = TimeSpan.FromMilliseconds(50),
                    EnableLiveSelectAccelerator = true
                }
            })
            .AddCBDDCNetwork<StaticPeerNodeConfigurationProvider>(false);

        // Benchmark runs use explicit known peers; disable UDP discovery and handshake overhead.
        services.AddSingleton<IDiscoveryService, PassiveDiscoveryService>();
        services.AddSingleton<IPeerHandshakeService, NoOpHandshakeService>();

        ServiceProvider provider = services.BuildServiceProvider();
        ICBDDCNode node = provider.GetRequiredService<ICBDDCNode>();
        SampleDbContext context = provider.GetRequiredService<SampleDbContext>();

        return new BenchmarkPeerNode(provider, node, context, workDir);
    }

    public async Task StartAsync()
    {
        if (_started) return;
        await _node.Start();
        _started = true;
    }

    public async Task StopAsync()
    {
        if (!_started) return;

        try
        {
            await _node.Stop();
        }
        catch (ObjectDisposedException)
        {
        }
        catch (AggregateException ex) when (ex.InnerExceptions.All(e => e is ObjectDisposedException))
        {
        }

        _started = false;
    }

    public async Task UpsertUserAsync(User user)
    {
        User? existing = Context.Users.Find(u => u.Id == user.Id).FirstOrDefault();
        if (existing == null)
            await Context.Users.InsertAsync(user);
        else
            await Context.Users.UpdateAsync(user);

        await Context.SaveChangesAsync();
    }

    public bool ContainsUser(string userId)
    {
        return Context.Users.Find(u => u.Id == userId).Any();
    }

    public async ValueTask DisposeAsync()
    {
        try
        {
            await StopAsync();
        }
        finally
        {
            _serviceProvider.Dispose();
            TryDeleteDirectory(_workDir);
        }
    }

    private static void TryDeleteDirectory(string path)
    {
        if (!Directory.Exists(path)) return;

        for (var attempt = 0; attempt < 5; attempt++)
            try
            {
                Directory.Delete(path, true);
                return;
            }
            catch when (attempt < 4)
            {
                Thread.Sleep(50);
            }
    }

    private sealed class PassiveDiscoveryService : IDiscoveryService
    {
        public IEnumerable<PeerNode> GetActivePeers()
        {
            return Array.Empty<PeerNode>();
        }

        public Task Start()
        {
            return Task.CompletedTask;
        }

        public Task Stop()
        {
            return Task.CompletedTask;
        }
    }
}