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.Lmdb;
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;

    /// <summary>
    ///     Gets the active EF/Core context for generated users.
    /// </summary>
    public SampleDbContext Context { get; }

    /// <summary>
    ///     Creates and initializes a benchmark peer node.
    /// </summary>
    /// <param name="serviceProvider">Service provider containing node dependencies.</param>
    /// <param name="node">Benchmark node abstraction.</param>
    /// <param name="context">Live data context used by the benchmark.</param>
    /// <param name="workDir">Temporary working directory for this node.</param>
    private BenchmarkPeerNode(
        ServiceProvider serviceProvider,
        ICBDDCNode node,
        SampleDbContext context,
        string workDir)
    {
        _serviceProvider = serviceProvider;
        _node = node;
        Context = context;
        _workDir = workDir;
    }

    /// <summary>
    ///     Creates and starts a benchmark peer node from configuration.
    /// </summary>
    /// <param name="nodeId">Unique peer identifier.</param>
    /// <param name="tcpPort">Local TCP port for the node.</param>
    /// <param name="authToken">Authentication token shared across peers.</param>
    /// <param name="knownPeers">Known peers to connect to at startup.</param>
    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>();

        bool useLmdb = GetBoolEnv("CBDDC_BENCH_USE_LMDB", defaultValue: true);
        bool dualWrite = GetBoolEnv("CBDDC_BENCH_DUAL_WRITE", defaultValue: true);
        bool preferLmdbReads = GetBoolEnv("CBDDC_BENCH_PREFER_LMDB_READS", defaultValue: true);
        bool enableShadowValidation = GetBoolEnv("CBDDC_BENCH_SHADOW_READ_VALIDATE", defaultValue: false);
        int reconcileIntervalMs = GetIntEnv("CBDDC_BENCH_RECONCILE_INTERVAL_MS", defaultValue: 0);

        var registration = 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
                }
            });

        if (useLmdb)
            registration.AddCBDDCLmdbOplog(
                _ => new LmdbOplogOptions
                {
                    EnvironmentPath = Path.Combine(workDir, "oplog-lmdb"),
                    MapSizeBytes = 256L * 1024 * 1024,
                    MaxDatabases = 16,
                    PruneBatchSize = 512
                },
                flags =>
                {
                    flags.UseLmdbOplog = true;
                    flags.DualWriteOplog = dualWrite;
                    flags.PreferLmdbReads = preferLmdbReads;
                    flags.EnableReadShadowValidation = enableShadowValidation;
                    flags.ReconciliationInterval = TimeSpan.FromMilliseconds(Math.Max(0, reconcileIntervalMs));
                });

        registration.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);
    }

    /// <summary>
    ///     Starts the node asynchronously.
    /// </summary>
    public async Task StartAsync()
    {
        if (_started) return;
        await _node.Start();
        _started = true;
    }

    /// <summary>
    ///     Stops the node asynchronously.
    /// </summary>
    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;
    }

    /// <summary>
    ///     Inserts or updates a user record.
    /// </summary>
    /// <param name="user">User payload.</param>
    public async Task UpsertUserAsync(User user)
    {
        await Context.Users.UpdateAsync(user);
        await Context.SaveChangesAsync();
    }

    /// <summary>
    ///     Returns whether a user identifier exists.
    /// </summary>
    /// <param name="userId">Target user identifier.</param>
    public bool ContainsUser(string userId)
    {
        return Context.Users.Find(u => u.Id == userId).Any();
    }

    /// <summary>
    ///     Counts user identifiers matching the provided prefix.
    /// </summary>
    /// <param name="prefix">User identifier prefix.</param>
    public int CountUsersWithPrefix(string prefix)
    {
        return Context.Users.FindAll().Count(u => u.Id.StartsWith(prefix, StringComparison.Ordinal));
    }

    /// <summary>
    ///     Disposes the node and any unmanaged resources.
    /// </summary>
    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 static bool GetBoolEnv(string key, bool defaultValue)
    {
        string? raw = Environment.GetEnvironmentVariable(key);
        if (string.IsNullOrWhiteSpace(raw)) return defaultValue;
        if (bool.TryParse(raw, out bool parsed)) return parsed;
        return defaultValue;
    }

    private static int GetIntEnv(string key, int defaultValue)
    {
        string? raw = Environment.GetEnvironmentVariable(key);
        if (string.IsNullOrWhiteSpace(raw)) return defaultValue;
        if (int.TryParse(raw, out int parsed)) return parsed;
        return defaultValue;
    }

    private sealed class PassiveDiscoveryService : IDiscoveryService
    {
        /// <summary>
        ///     Gets the current list of active peers.
        /// </summary>
        public IEnumerable<PeerNode> GetActivePeers()
        {
            return Array.Empty<PeerNode>();
        }

        /// <summary>
        ///     Starts discovery.
        /// </summary>
        public Task Start()
        {
            return Task.CompletedTask;
        }

        /// <summary>
        ///     Stops discovery.
        /// </summary>
        public Task Stop()
        {
            return Task.CompletedTask;
        }
    }
}