Track database handles over metrics

PluralKit.Bot/Bot.cs

@@ -79,7 +79,8 @@ namespace PluralKit.Bot
             .AddTransient(_ => _config.GetSection("PluralKit").Get<CoreConfig>() ?? new CoreConfig())
             .AddTransient(_ => _config.GetSection("PluralKit").GetSection("Bot").Get<BotConfig>() ?? new BotConfig())
 
-            .AddTransient(svc => new DbConnectionFactory(svc.GetRequiredService<CoreConfig>().Database))
+            .AddSingleton<DbConnectionCountHolder>()
+            .AddTransient<DbConnectionFactory>()
 
             .AddSingleton<IDiscordClient, DiscordShardedClient>(_ => new DiscordShardedClient(new DiscordSocketConfig
             {

PluralKit.Bot/Services/PeriodicStatCollector.cs

@@ -24,9 +24,11 @@ namespace PluralKit.Bot
 
         private WebhookCacheService _webhookCache;
 
+        private DbConnectionCountHolder _countHolder;
+
         private ILogger _logger;
 
-        public PeriodicStatCollector(IDiscordClient client, IMetrics metrics, SystemStore systems, MemberStore members, SwitchStore switches, MessageStore messages, ILogger logger, WebhookCacheService webhookCache)
+        public PeriodicStatCollector(IDiscordClient client, IMetrics metrics, SystemStore systems, MemberStore members, SwitchStore switches, MessageStore messages, ILogger logger, WebhookCacheService webhookCache, DbConnectionCountHolder countHolder)
         {
             _client = (DiscordShardedClient) client;
             _metrics = metrics;
@@ -35,6 +37,7 @@ namespace PluralKit.Bot
             _switches = switches;
             _messages = messages;
             _webhookCache = webhookCache;
+            _countHolder = countHolder;
             _logger = logger.ForContext<PeriodicStatCollector>();
         }
 
@@ -75,6 +78,9 @@ namespace PluralKit.Bot
             _metrics.Measure.Gauge.SetValue(CoreMetrics.ProcessHandles, process.HandleCount);
             _metrics.Measure.Gauge.SetValue(CoreMetrics.CpuUsage, await EstimateCpuUsage());
             
+            // Database info
+            _metrics.Measure.Gauge.SetValue(CoreMetrics.DatabaseConnections, _countHolder.ConnectionCount);
+            
             // Other shiz
             _metrics.Measure.Gauge.SetValue(BotMetrics.WebhookCacheSize, _webhookCache.CacheSize);
 

PluralKit.Core/CoreMetrics.cs

@@ -17,5 +17,8 @@ namespace PluralKit.Core
         public static GaugeOptions ProcessThreads => new GaugeOptions { Name = "Process Thread Count", MeasurementUnit = Unit.Threads, Context = "Process" };
         public static GaugeOptions ProcessHandles => new GaugeOptions { Name = "Process Handle Count", MeasurementUnit = Unit.Items, Context = "Process" };
         public static GaugeOptions CpuUsage => new GaugeOptions { Name = "CPU Usage", MeasurementUnit = Unit.Percent, Context = "Process" };
+        
+        public static MeterOptions DatabaseRequests => new MeterOptions() { Name = "Database Requests", MeasurementUnit = Unit.Requests, Context = "Database" };
+        public static GaugeOptions DatabaseConnections => new GaugeOptions() { Name = "Database Connections", MeasurementUnit = Unit.Connections, Context = "Database" };
     }
 }

PluralKit.Core/Utils.cs

@@ -5,7 +5,9 @@ using System.IO;
 using System.Linq;
 using System.Security.Cryptography;
 using System.Text.RegularExpressions;
+using System.Threading;
 using System.Threading.Tasks;
+using App.Metrics;
 using Dapper;
 using Microsoft.Extensions.Configuration;
 using Newtonsoft.Json;
@@ -13,6 +15,7 @@ using NodaTime;
 using NodaTime.Serialization.JsonNet;
 using NodaTime.Text;
 using Npgsql;
+using PluralKit.Core;
 using Serilog;
 using Serilog.Formatting.Compact;
 using Serilog.Sinks.SystemConsole.Themes;
@@ -349,18 +352,109 @@ namespace PluralKit
     
     public class DbConnectionFactory
     {
-        private string _connectionString;
+        private CoreConfig _config;
+        private IMetrics _metrics;
+        private DbConnectionCountHolder _countHolder;
 
-        public DbConnectionFactory(string connectionString)
+        public DbConnectionFactory(CoreConfig config, DbConnectionCountHolder countHolder, IMetrics metrics)
         {
-            _connectionString = connectionString;
+            _config = config;
+            _countHolder = countHolder;
+            _metrics = metrics;
         }
 
         public async Task<IDbConnection> Obtain()
         {
-            var conn = new NpgsqlConnection(_connectionString);
+            // Mark the request (for a handle, I guess) in the metrics
+            _metrics.Measure.Meter.Mark(CoreMetrics.DatabaseRequests);
+            
+            // Actually create and try to open the connection
+            var conn = new NpgsqlConnection(_config.Database);
             await conn.OpenAsync();
-            return conn;
+            
+            // Increment the count
+            _countHolder.Increment();
+            // Return a wrapped connection which will decrement the counter on dispose
+            return new DbConnectionTrackingConnection(conn, _countHolder);
         }
     }
+
+    public class DbConnectionCountHolder
+    {
+        private int _connectionCount;
+        public int ConnectionCount => _connectionCount;
+
+        public void Increment()
+        {
+            Interlocked.Increment(ref _connectionCount);
+        }
+
+        public void Decrement()
+        {
+            Interlocked.Decrement(ref _connectionCount);
+        }
+    } 
+    
+    public class DbConnectionTrackingConnection: IDbConnection
+    {
+        // Simple delegation of everything.
+        private IDbConnection _impl;
+
+        private DbConnectionCountHolder _countHolder;
+
+        public DbConnectionTrackingConnection(IDbConnection impl, DbConnectionCountHolder countHolder)
+        {
+            _impl = impl;
+            _countHolder = countHolder;
+        }
+
+        public void Dispose()
+        {
+            _impl.Dispose();
+            
+            _countHolder.Decrement();
+        }
+
+        public IDbTransaction BeginTransaction()
+        {
+            return _impl.BeginTransaction();
+        }
+
+        public IDbTransaction BeginTransaction(IsolationLevel il)
+        {
+            return _impl.BeginTransaction(il);
+        }
+
+        public void ChangeDatabase(string databaseName)
+        {
+            _impl.ChangeDatabase(databaseName);
+        }
+
+        public void Close()
+        {
+            _impl.Close();
+        }
+
+        public IDbCommand CreateCommand()
+        {
+            return _impl.CreateCommand();
+        }
+
+        public void Open()
+        {
+            _impl.Open();
+        }
+
+        public string ConnectionString
+        {
+            get => _impl.ConnectionString;
+            set => _impl.ConnectionString = value;
+        }
+
+        public int ConnectionTimeout => _impl.ConnectionTimeout;
+
+        public string Database => _impl.Database;
+
+        public ConnectionState State => _impl.State;
+    }
 }