How to Enforce a Global Max Connection Limit When Load Balancing Npgsql Connections

				
					using System;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using System.Collections.Generic;
using Npgsql;

class Program
{
    // Counters
    private static int _started;
    private static int _succeeded;
    private static int _failed;
    private static long _maxMonBackends;

    // Logging
    private static string? _logFile = null;
    private static readonly object _logLock = new();

    static void Log(string msg)
    {
        if (_logFile == null)
        {
            Console.WriteLine(msg);
            return;
        }
        lock (_logLock)
        {
            File.AppendAllText(_logFile, msg + Environment.NewLine);
        }
    }

    static async Task Main(string[] args)
    {
        // -----------------------
        // DEFAULT CONFIG
        // -----------------------
        string seedHost = "";
        int workers = -1;
        int maxPool = -1;

        bool loadBalance = false;
        int sleepSeconds = 10;
        int timeoutSeconds = 5;

        string? logfile = null;

        // -----------------------
        // PARSE NAMED PARAMETERS
        // -----------------------
        foreach (var arg in args)
        {
            if (arg == "--help")
            {
                PrintHelp();
                return;
            }
            if (arg.StartsWith("--seed="))
                seedHost = arg.Substring("--seed=".Length);

            else if (arg.StartsWith("--workers="))
                int.TryParse(arg.Substring("--workers=".Length), out workers);

            else if (arg.StartsWith("--maxpool="))
                int.TryParse(arg.Substring("--maxpool=".Length), out maxPool);

            else if (arg == "--lb")
                loadBalance = true;

            else if (arg.StartsWith("--logfile="))
            {
                logfile = arg.Substring("--logfile=".Length);
                try
                {
                    File.WriteAllText(logfile, "=== Log Start ===\n");
                    _logFile = logfile;
                }
                catch (Exception ex)
                {
                    Console.WriteLine($"Failed to open logfile '{logfile}': {ex.Message}");
                    _logFile = null;
                }
            }
            else if (arg.StartsWith("--sleep="))
                int.TryParse(arg.Substring("--sleep=".Length), out sleepSeconds);

            else if (arg.StartsWith("--timeout="))
                int.TryParse(arg.Substring("--timeout=".Length), out timeoutSeconds);

            else
            {
                Console.WriteLine($"Unknown parameter: {arg}");
                PrintHelp();
                return;
            }
        }

        // -----------------------
        // VALIDATION
        // -----------------------
        if (string.IsNullOrWhiteSpace(seedHost) ||
            workers <= 0 ||
            maxPool <= 0)
        {
            PrintHelp();
            return;
        }

        // Fixed DB settings
        int port = 5433;
        string dbName = "yugabyte";
        string dbUser = "yugabyte";
        string dbPass = "yugabyte";

        string workerAppName = "YbNpgsqlPoolDemo";
        string monitorAppName = "YbNpgsqlPoolDemo_MON";

        // -----------------------
        // DISCOVER CLUSTER HOSTS
        // -----------------------
        List<string> allHosts = await DiscoverClusterAsync(
            seedHost, port, dbName, dbUser, dbPass, monitorAppName);

        List<string> workloadHosts =
            loadBalance ? allHosts : new List<string> { seedHost };

        string hostList = string.Join(",", workloadHosts);
        string lbClause = loadBalance ? "Load Balance Hosts=true;" : "";

        // -----------------------
        // BUILD CONNECTION STRING
        // -----------------------
        string endpointBase =
            $"Host={hostList};Port={port};Database={dbName};Username={dbUser};Password={dbPass};{lbClause}";

        string pooledConnString =
            endpointBase +
            $"Application Name={workerAppName};Maximum Pool Size={maxPool};Timeout={timeoutSeconds};";

        Log("=== FULL LOG START ===");
        Log($"Seed host            = {seedHost}");
        Log($"Workers              = {workers}");
        Log($"MaxPool              = {maxPool}");
        Log($"Load balance         = {loadBalance}");
        Log($"SleepSeconds         = {sleepSeconds}");
        Log($"TimeoutSeconds       = {timeoutSeconds}");
        Log($"Discovered hosts     = {string.Join(",", allHosts)}");
        Log($"Workload host list   = {hostList}");
        Log($"Worker connstring    = {pooledConnString}");

        // -----------------------
        // PRINT CLEAN SUMMARY HEADER
        // -----------------------
        Console.WriteLine("=== Npgsql/YugabyteDB Connection Pool Demo ===");
        Console.WriteLine($"Seed host          : {seedHost}");
        Console.WriteLine($"Workers            : {workers}");
        Console.WriteLine($"Max pool size      : {maxPool}");
        Console.WriteLine($"Load balancing     : {loadBalance}");
        Console.WriteLine($"Discovered hosts   : {string.Join(",", allHosts)}");
        Console.WriteLine($"Workload hosts     : {hostList}");
        Console.WriteLine($"Sleep (seconds)    : {sleepSeconds}");
        Console.WriteLine($"Timeout (seconds)  : {timeoutSeconds}");
        if (logfile != null)
            Console.WriteLine($"Logfile            : {logfile}");
        Console.WriteLine();

        // -----------------------
        // START WORKERS & MONITOR
        // -----------------------
        var cts = new CancellationTokenSource();

        var monitorTask = MonitorClusterAsync(
            allHosts, port, dbName, dbUser, dbPass,
            workerAppName, monitorAppName, cts.Token);

        var workerTasks = new Task[workers];

        for (int i = 0; i < workers; i++)
        {
            int id = i;
            workerTasks[i] = Task.Run(() =>
                WorkerAsync(id, pooledConnString, sleepSeconds), cts.Token);
        }

        await Task.WhenAll(workerTasks);

        cts.Cancel();
        try { await monitorTask; } catch { }

        // -----------------------
        // FINAL PER-NODE COUNTS
        // -----------------------
        var perNode = await PerNodeCountsAsync(


            allHosts, port, dbName, dbUser, dbPass,
            workerAppName, monitorAppName);

        // -----------------------
        // FINAL SUMMARY
        // -----------------------
        Console.WriteLine("=== Summary (App-side) ===");
        Console.WriteLine($"Started      : {_started}");
        Console.WriteLine($"Succeeded    : {_succeeded}");
        Console.WriteLine($"Failed       : {_failed}");

        Console.WriteLine("\n=== Summary (DB-side, via pg_stat_activity per node) ===");
        Console.WriteLine($"Max concurrent backend count across cluster: {_maxMonBackends}");

        foreach (var kvp in perNode)
            Console.WriteLine($"Host {kvp.Key} = {kvp.Value} backend(s)");

        if (logfile != null)
            Console.WriteLine($"\nFull log written to: {logfile}");
    }

    // ------------------------------------------------------
    // WORKER THREAD
    // ------------------------------------------------------
    private static async Task WorkerAsync(int id, string connString, int sleepSeconds)
    {
        int started = Interlocked.Increment(ref _started);
        Log($"[WORKER {id:00}] STARTED (#{started})");

        try
        {
            await using var conn = new NpgsqlConnection(connString);
            await conn.OpenAsync();

            int succ = Interlocked.Increment(ref _succeeded);
            Log($"[WORKER {id:00}] GOT CONNECTION (succ={succ})");

            string sql = $"select pg_sleep({sleepSeconds});";

            await using var cmd = new NpgsqlCommand(sql, conn);
            await cmd.ExecuteNonQueryAsync();

            Log($"[WORKER {id:00}] FINISHED");
        }
        catch (Exception ex)
        {
            int f = Interlocked.Increment(ref _failed);
            Log($"[WORKER {id:00}] FAILED (#{f}): {ex.Message}");
        }
    }

    // ------------------------------------------------------
    // CLUSTER DISCOVERY (yb_servers)
    // ------------------------------------------------------
    private static async Task<List<string>> DiscoverClusterAsync(
        string seed, int port, string db, string user, string pass, string appName)
    {
        var hosts = new List<string>();

        string connStr =
            $"Host={seed};Port={port};Database={db};Username={user};Password={pass};" +
            $"Application Name={appName}_DISC;Pooling=false;";

        try
        {
            await using var conn = new NpgsqlConnection(connStr);
            await conn.OpenAsync();

            const string sql = "select host from yb_servers();";

            await using var cmd = new NpgsqlCommand(sql, conn);
            await using var reader = await cmd.ExecuteReaderAsync();

            while (await reader.ReadAsync())
            {
                if (!reader.IsDBNull(0))
                {
                    string h = reader.GetString(0).Trim();
                    if (!string.IsNullOrEmpty(h))
                        hosts.Add(h);
                }
            }

        }
        catch (Exception ex)
        {
            Log($"[DISCOVERY] Failed: {ex.Message}");
        }

        if (hosts.Count == 0)
        {
            Log($"[DISCOVERY] No results from yb_servers(); using seed host {seed}");
            hosts.Add(seed);
        }

        return hosts;
    }

    // ------------------------------------------------------
    // CLUSTER-WIDE MONITOR
    // ------------------------------------------------------
    private static async Task MonitorClusterAsync(
        List<string> hosts, int port, string db, string user, string pass,
        string workerApp, string monitorApp, CancellationToken token)
    {
        const int intervalMs = 1000;

        while (!token.IsCancellationRequested)
        {
            long total = 0;

            foreach (var h in hosts)
            {
                string connStr =
                    $"Host={h};Port={port};Database={db};Username={user};Password={pass};" +
                    $"Application Name={monitorApp};Pooling=false;";

                try
                {
                    await using var conn = new NpgsqlConnection(connStr);
                    await conn.OpenAsync(token);

                    const string sql =
                        "select count(*) from pg_stat_activity where application_name=@app;";

                    await using var cmd = new NpgsqlCommand(sql, conn);
                    cmd.Parameters.AddWithValue("app", workerApp);

                    long count = (long)(await cmd.ExecuteScalarAsync(token))!;
                    total += count;
                }
                catch
                {
                    // Ignore per-node monitor errors in this loop
                }
            }

            if (total > _maxMonBackends)
                _maxMonBackends = total;

            Log($"[MON] cluster-active-backends = {total}");

            try { await Task.Delay(intervalMs, token); }
            catch { break; }
        }

        Log("[MON] stopped");
    }


    // ------------------------------------------------------
    // FINAL PER-NODE COUNTS
    // ------------------------------------------------------
    private static async Task<Dictionary<string, long>> PerNodeCountsAsync(
        List<string> hosts, int port, string db, string user, string pass,
        string workerApp, string monitorApp)
    {
        var result = new Dictionary<string, long>();

        foreach (var h in hosts)
        {
            string connStr =
                $"Host={h};Port={port};Database={db};Username={user};Password={pass};" +
                $"Application Name={monitorApp}_SUMMARY;Pooling=false;";

            try
            {
                await using var conn = new NpgsqlConnection(connStr);
                await conn.OpenAsync();

                const string sql =
                    "select count(*) from pg_stat_activity where application_name=@app;";

                await using var cmd = new NpgsqlCommand(sql, conn);
                cmd.Parameters.AddWithValue("app", workerApp);

                long count = (long)(await cmd.ExecuteScalarAsync())!;
                Log($"[SUMMARY] Host {h} has {count} backend(s)");

                result[h] = count;
            }
            catch (Exception ex)
            {
                Log($"[SUMMARY] Failed on {h}: {ex.Message}");
                result[h] = -1;
            }
        }

        return result;
    }

    // ------------------------------------------------------
    // HELP TEXT
    // ------------------------------------------------------
    private static void PrintHelp()
    {
        Console.WriteLine();
        Console.WriteLine("Usage:");
        Console.WriteLine("  dotnet run -- \\");
        Console.WriteLine("     --seed=<host>        REQUIRED  YSQL seed host for discovery");
        Console.WriteLine("     --workers=<int>      REQUIRED  Number of workers");
        Console.WriteLine("     --maxpool=<int>      REQUIRED  Max Npgsql pool size");
        Console.WriteLine();
        Console.WriteLine("Optional flags:");
        Console.WriteLine("     --lb                 Enable load balance across YB nodes");
        Console.WriteLine("     --logfile=<file>     Save detailed logs to file");
        Console.WriteLine("     --sleep=<seconds>    How long workers hold the connection (default 10)");
        Console.WriteLine("     --timeout=<seconds>  Npgsql connection timeout (default 5)");
        Console.WriteLine();
        Console.WriteLine("Example:");
        Console.WriteLine("  dotnet run -- --seed=127.0.0.1 --workers=20 --maxpool=5 --lb --logfile=run.log");
        Console.WriteLine();
    }
}