Understanding Smart-Driver Pooling with Multiple Connection Pools (and Important Known Issues)

				
					// Program.cs
// YugabyteDB Npgsql Smart Driver - Multi Pool / Search Path Demo
//
// Demonstrates:
//  1) YugabyteDB Npgsql Smart Driver load balancing (Load Balance Hosts=true)
//  2) Multiple connection strings => (ideally) multiple pools
//  3) Forced application_name AFTER OpenAsync (Option B - decisive)
//  4) Forced search_path per pool + verification by reading same table name from different schemas
//
// Logging behavior:
//  - Header + pool config: console + file
//  - Per-worker/per-iteration details: file only
//  - Summary: console only

using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using YBNpgsql; // YugabyteDB Npgsql Smart Driver

internal enum PoolBehavior
{
    Block,
    Failfast
}

internal sealed class PoolConfig
{
    public string Name { get; init; } = "";
    public string Schema { get; init; } = "";
    public int MaxPoolSize { get; init; }
    public string AppName { get; init; } = "";
    public string ConnectionString { get; init; } = "";
}

internal sealed class PoolStats
{
    public int WorkersStarted;
    public int WorkersSucceeded;
    public int WorkersFailed;

    public long SchemaMatches;
    public long SchemaMismatches;

    private int _activeNow;
    public int MaxConcurrentActive;

    public void IncActive()
    {
        var now = Interlocked.Increment(ref _activeNow);
        UpdateMax(now);
    }

    public void DecActive()
    {
        Interlocked.Decrement(ref _activeNow);
    }

    public int ActiveNow => Volatile.Read(ref _activeNow);

    private void UpdateMax(int now)
    {
        while (true)
        {
            var max = Volatile.Read(ref MaxConcurrentActive);
            if (now <= max) return;
            if (Interlocked.CompareExchange(ref MaxConcurrentActive, now, max) == max) return;
        }
    }
}

internal sealed class DualLogger : IDisposable
{
    private readonly StreamWriter _file;
    private readonly object _lock = new();

    public DualLogger(string logFilePath)
    {
        _file = new StreamWriter(File.Open(logFilePath, FileMode.Create, FileAccess.Write, FileShare.Read))
        {
            AutoFlush = true,
            NewLine = "\n",
        };
    }

    // goes to console + file
    public void Info(string msg)
    {
        lock (_lock)
        {
            Console.WriteLine(msg);
            _file.WriteLine(msg);
        }
    }

    // goes to file only
    public void Verbose(string msg)
    {
        lock (_lock)
        {
            _file.WriteLine(msg);
        }
    }

    public void Dispose()
    {
        lock (_lock)
        {
            _file.Flush();
            _file.Dispose();
        }
    }
}

public static class Program
{
    public static async Task<int> Main(string[] args)
    {
        // Defaults
        string host = GetArg(args, "--host", "127.0.0.1");
        int port = GetArgInt(args, "--port", 5433);
        string database = GetArg(args, "--db", "yugabyte");
        string username = GetArg(args, "--user", "yugabyte");
        string password = GetArg(args, "--pass", "yugabyte");

        int workersPerPool = GetArgInt(args, "--workers-per-pool", 40);

        int max1 = GetArgInt(args, "--maxpool-tenant1", 2);
        int max2 = GetArgInt(args, "--maxpool-tenant2", 4);
        int max3 = GetArgInt(args, "--maxpool-tenant3", 8);

        int sleepSeconds = GetArgInt(args, "--sleep", 10);

        // Npgsql "Timeout" controls connect + pool wait time. 0 = wait indefinitely for pool.
        int timeoutSeconds = GetArgInt(args, "--timeout", 0);

        string behaviorRaw = GetArg(args, "--pool-behavior", "block").Trim().ToLowerInvariant();
        PoolBehavior behavior = behaviorRaw switch
        {
            "block" => PoolBehavior.Block,
            "failfast" => PoolBehavior.Failfast,
            _ => PoolBehavior.Block
        };

        int iterations = GetArgInt(args, "--iterations", 1);
        int iterSleepMs = GetArgInt(args, "--iter-sleep-ms", 0);

        string logFile = GetArg(args, "--log-file",
            $"yb_smart_driver_search_path_demo_{DateTime.Now:yyyyMMdd_HHmmss}.log");

        using var logger = new DualLogger(logFile);

        // Header
        logger.Info("=== YugabyteDB Npgsql Smart Driver - Multi Pool / Search Path Demo ===");
        logger.Info($"Args              : {string.Join(" ", args.Select(a => a.Contains(' ') ? $"\"{a}\"" : a))}");
        logger.Info($"Host              : {host}");
        logger.Info($"Port              : {port}");
        logger.Info($"Database          : {database}");
        logger.Info($"User              : {username}");
        logger.Info($"Workers per pool  : {workersPerPool}");
        logger.Info($"Sleep (seconds)   : {sleepSeconds}");
        logger.Info($"Timeout (seconds) : {timeoutSeconds}  (0 = wait indefinitely for a pooled connection)");
        logger.Info($"Pool behavior     : {behavior.ToString().ToLowerInvariant()}  (block | failfast)");
        logger.Info($"Iterations        : {iterations}");
        logger.Info($"Iter sleep (ms)   : {iterSleepMs}");
        logger.Info($"Log file          : {logFile}");
        logger.Info("");

        var pools = BuildPools(host, port, database, username, password, timeoutSeconds, max1, max2, max3);

        foreach (var p in pools)
            logger.Info($"Pool {p.Name}: schema={p.Schema}, MaxPoolSize={p.MaxPoolSize}, Application Name={p.AppName}");

        logger.Info("");
        logger.Info("Running workers across ALL pools in parallel...");
        logger.Info("");

        var statsByPool = pools.ToDictionary(p => p.Name, _ => new PoolStats());

        var tasks = new List<Task>();
        foreach (var pool in pools)
        {
            var stats = statsByPool[pool.Name];
            for (int w = 0; w < workersPerPool; w++)
            {
                int workerId = w;
                Interlocked.Increment(ref stats.WorkersStarted);
                tasks.Add(RunWorkerAsync(pool, stats, workerId, iterations, iterSleepMs, sleepSeconds, logger));
            }
        }

        await Task.WhenAll(tasks);

        PrintSummaryToConsoleOnly(pools, statsByPool);

        Console.WriteLine();
        Console.WriteLine("Demo complete.");
        Console.WriteLine($"See the log file for per-worker details:\n  {logFile}");
        return 0;
    }

    private static List<PoolConfig> BuildPools(
        string host,
        int port,
        string database,
        string username,
        string password,
        int timeoutSeconds,
        int max1,
        int max2,
        int max3)
    {
        // Use a single seed host; smart driver should discover topology if supported in this environment.
        string common =
            $"Host={host};" +
            $"Port={port};" +
            $"Database={database};" +
            $"Username={username};" +
            $"Password={EscapeConnStringValue(password)};" +
            $"Load Balance Hosts=true;" +
            $"Timeout={timeoutSeconds};" +
            $"No Reset On Close=true;"; // helps surface session-state reuse issues

        return new List<PoolConfig>
        {
            new PoolConfig
            {
                Name = "Conn1",
                Schema = "tenant1",
                MaxPoolSize = max1,
                AppName = "yb-smart-pool-tenant1",
                ConnectionString = common +
                                   $"Application Name=yb-smart-pool-tenant1;" +
                                   $"Maximum Pool Size={max1};"
            },
            new PoolConfig
            {
                Name = "Conn2",
                Schema = "tenant2",
                MaxPoolSize = max2,
                AppName = "yb-smart-pool-tenant2",
                ConnectionString = common +
                                   $"Application Name=yb-smart-pool-tenant2;" +
                                   $"Maximum Pool Size={max2};"
            },
            new PoolConfig
            {
                Name = "Conn3",
                Schema = "tenant3",
                MaxPoolSize = max3,
                AppName = "yb-smart-pool-tenant3",
                ConnectionString = common +
                                   $"Application Name=yb-smart-pool-tenant3;" +
                                   $"Maximum Pool Size={max3};"
            }
        };
    }

    private static async Task RunWorkerAsync(
        PoolConfig pool,
        PoolStats stats,
        int workerId,
        int iterations,
        int iterSleepMs,
        int sleepSeconds,
        DualLogger logger)
    {
        string ts() => DateTime.Now.ToString("HH:mm:ss.fff", CultureInfo.InvariantCulture);

        bool opened = false;

        try
        {
            await using var conn = new NpgsqlConnection(pool.ConnectionString);

            await conn.OpenAsync();

            opened = true;
            stats.IncActive();

            // OPTION B: force application_name explicitly after OpenAsync
            await using (var setApp = new NpgsqlCommand($"SET application_name TO '{SqlLiteral(pool.AppName)}';", conn))
                await setApp.ExecuteNonQueryAsync();

            // Force schema via search_path
            await using (var setPath = new NpgsqlCommand($"SET search_path TO {SqlIdent(pool.Schema)}, public;", conn))
                await setPath.ExecuteNonQueryAsync();

            // Verify server-side view (file only)
            string appNameNow = await ScalarStringAsync(conn, "SHOW application_name;");
            string schemaNow = await ScalarStringAsync(conn, "SELECT current_schema();");
            logger.Verbose($"{ts()} [POOL {pool.Name}] Worker {workerId:D2} OPEN ok => app='{appNameNow}', schema='{schemaNow}', activeNow={stats.ActiveNow}");

            long localMatches = 0;
            long localMismatches = 0;

            for (int i = 0; i < iterations; i++)
            {
                await using var cmd = new NpgsqlCommand("SELECT current_schema(), value FROM tenant_data WHERE id = 1;", conn);

                await using var reader = await cmd.ExecuteReaderAsync();
                if (await reader.ReadAsync())
                {
                    string gotSchema = reader.GetString(0);
                    string value = reader.IsDBNull(1) ? "<null>" : reader.GetString(1);

                    bool match = string.Equals(gotSchema, pool.Schema, StringComparison.OrdinalIgnoreCase);
                    if (match) localMatches++;
                    else localMismatches++;

                    logger.Verbose($"{ts()} [POOL {pool.Name}] Worker {workerId:D2} iter={i:D3} => schema={gotSchema}, expected={pool.Schema}, match={match}, value='{value}'");
                }
                else
                {
                    localMismatches++;
                    logger.Verbose($"{ts()} [POOL {pool.Name}] Worker {workerId:D2} iter={i:D3} => ERROR: no rows returned");
                }

                if (iterSleepMs > 0)
                    await Task.Delay(iterSleepMs);
            }

            if (sleepSeconds > 0)
            {
                logger.Verbose($"{ts()} [POOL {pool.Name}] Worker {workerId:D2} sleeping for {sleepSeconds}s with an open connection (activeNow={stats.ActiveNow}).");
                await Task.Delay(TimeSpan.FromSeconds(sleepSeconds));
            }

            Interlocked.Add(ref stats.SchemaMatches, localMatches);
            Interlocked.Add(ref stats.SchemaMismatches, localMismatches);

            Interlocked.Increment(ref stats.WorkersSucceeded);
        }
        catch (Exception ex)
        {
            Interlocked.Increment(ref stats.WorkersFailed);
            logger.Verbose($"[POOL {pool.Name}] Worker {workerId:D2} ERROR: {ex.Message}");
        }
        finally
        {
            // Accurate: decrement only if OpenAsync succeeded and we incremented.
            if (opened)
                stats.DecActive();
        }
    }

    private static async Task<string> ScalarStringAsync(NpgsqlConnection conn, string sql)
    {
        await using var cmd = new NpgsqlCommand(sql, conn);
        var obj = await cmd.ExecuteScalarAsync();
        return obj?.ToString() ?? "<null>";
    }

    private static void PrintSummaryToConsoleOnly(List<PoolConfig> pools, Dictionary<string, PoolStats> statsByPool)
    {
        Console.WriteLine();
        Console.WriteLine("=== Summary Per Pool (App-side) ===");

        foreach (var pool in pools)
        {
            var s = statsByPool[pool.Name];

            int maxActive = s.MaxConcurrentActive;
            bool exceeded = maxActive > pool.MaxPoolSize;

            long matches = Interlocked.Read(ref s.SchemaMatches);
            long mismatches = Interlocked.Read(ref s.SchemaMismatches);

            Console.WriteLine($"Pool {pool.Name} (schema={pool.Schema})");
            Console.WriteLine($"  MaxPoolSize configured          : {pool.MaxPoolSize}");
            Console.WriteLine($"  Workers started                 : {s.WorkersStarted}");
            Console.WriteLine($"  Workers succeeded               : {s.WorkersSucceeded}");
            Console.WriteLine($"  Workers failed (timeouts/etc.)  : {s.WorkersFailed}");
            Console.WriteLine($"  Max concurrent active workers   : {maxActive}");
            Console.WriteLine($"  Did MaxActive exceed MaxPool?   : {(exceeded ? "YES (unexpected!)" : "NO (respected)")}");
            Console.WriteLine($"  Schema matches (search_path OK) : {matches}");
            Console.WriteLine($"  Schema mismatches               : {mismatches}");
            Console.WriteLine($"  Did all successful reads use the expected schema? : {(mismatches == 0 && matches > 0 ? "YES" : "NO (check log file)")}");
            Console.WriteLine();
        }
    }

    // Args helpers
    private static string GetArg(string[] args, string name, string defaultValue)
    {
        for (int i = 0; i < args.Length; i++)
        {
            if (string.Equals(args[i], name, StringComparison.OrdinalIgnoreCase))
                return (i + 1 < args.Length) ? args[i + 1] : defaultValue;
        }
        return defaultValue;
    }

    private static int GetArgInt(string[] args, string name, int defaultValue)
    {
        var s = GetArg(args, name, defaultValue.ToString(CultureInfo.InvariantCulture));
        return int.TryParse(s, NumberStyles.Integer, CultureInfo.InvariantCulture, out var v) ? v : defaultValue;
    }

    // SQL helpers
    private static string SqlIdent(string ident) => "\"" + ident.Replace("\"", "\"\"") + "\"";
    private static string SqlLiteral(string s) => s.Replace("'", "''");

    private static string EscapeConnStringValue(string s)
    {
        bool needsQuotes =
            s.Contains(';') ||
            s.StartsWith(' ') ||
            s.EndsWith(' ') ||
            s.Contains('\'') ||
            s.Contains('"');

        if (!needsQuotes) return s;
        return $"'{s.Replace("'", "''")}'";
    }
}