Limit postgres database concurrency (#17246)

Content.Server/Database/ServerDbManager.cs

@@ -283,11 +283,11 @@ namespace Content.Server.Database
             {
                 case "sqlite":
                     SetupSqlite(out var contextFunc, out var inMemory);
-                    _db = new ServerDbSqlite(contextFunc, inMemory);
+                    _db = new ServerDbSqlite(contextFunc, inMemory, _cfg);
                     break;
                 case "postgres":
                     var pgOptions = CreatePostgresOptions();
-                    _db = new ServerDbPostgres(pgOptions);
+                    _db = new ServerDbPostgres(pgOptions, _cfg);
                     break;
                 default:
                     throw new InvalidDataException($"Unknown database engine {engine}.");

Content.Server/Database/ServerDbPostgres.cs

@@ -4,7 +4,9 @@ using System.Linq;
 using System.Net;
 using System.Threading;
 using System.Threading.Tasks;
+using Content.Shared.CCVar;
 using Microsoft.EntityFrameworkCore;
+using Robust.Shared.Configuration;
 using Robust.Shared.Network;
 using Robust.Shared.Utility;
 
@@ -13,11 +15,15 @@ namespace Content.Server.Database
     public sealed class ServerDbPostgres : ServerDbBase
     {
         private readonly DbContextOptions<PostgresServerDbContext> _options;
+        private readonly SemaphoreSlim _prefsSemaphore;
         private readonly Task _dbReadyTask;
 
-        public ServerDbPostgres(DbContextOptions<PostgresServerDbContext> options)
+        public ServerDbPostgres(DbContextOptions<PostgresServerDbContext> options, IConfigurationManager cfg)
         {
+            var concurrency = cfg.GetCVar(CCVars.DatabasePgConcurrency);
+
             _options = options;
+            _prefsSemaphore = new SemaphoreSlim(concurrency, concurrency);
 
             _dbReadyTask = Task.Run(async () =>
             {
@@ -485,8 +491,9 @@ namespace Content.Server.Database
         private async Task<DbGuardImpl> GetDbImpl()
         {
             await _dbReadyTask;
+            await _prefsSemaphore.WaitAsync();
 
-            return new DbGuardImpl(new PostgresServerDbContext(_options));
+            return new DbGuardImpl(this, new PostgresServerDbContext(_options));
         }
 
         protected override async Task<DbGuard> GetDb()
@@ -496,17 +503,21 @@ namespace Content.Server.Database
 
         private sealed class DbGuardImpl : DbGuard
         {
-            public DbGuardImpl(PostgresServerDbContext dbC)
+            private readonly ServerDbPostgres _db;
+
+            public DbGuardImpl(ServerDbPostgres db, PostgresServerDbContext dbC)
             {
+                _db = db;
                 PgDbContext = dbC;
             }
 
             public PostgresServerDbContext PgDbContext { get; }
             public override ServerDbContext DbContext => PgDbContext;
 
-            public override ValueTask DisposeAsync()
+            public override async ValueTask DisposeAsync()
             {
-                return DbContext.DisposeAsync();
+                await DbContext.DisposeAsync();
+                _db._prefsSemaphore.Release();
             }
         }
     }

Content.Server/Database/ServerDbSqlite.cs

@@ -20,22 +20,21 @@ namespace Content.Server.Database
     {
         private readonly Func<DbContextOptions<SqliteServerDbContext>> _options;
 
-        // This doesn't allow concurrent access so that's what the semaphore is for.
-        // That said, this is bloody SQLite, I don't even think EFCore bothers to truly async it.
         private readonly SemaphoreSlim _prefsSemaphore;
 
         private readonly Task _dbReadyTask;
 
         private int _msDelay;
 
-        public ServerDbSqlite(Func<DbContextOptions<SqliteServerDbContext>> options, bool inMemory)
+        public ServerDbSqlite(
+            Func<DbContextOptions<SqliteServerDbContext>> options,
+            bool inMemory,
+            IConfigurationManager cfg)
         {
             _options = options;
 
             var prefsCtx = new SqliteServerDbContext(options());
 
-            var cfg = IoCManager.Resolve<IConfigurationManager>();
-
             // When inMemory we re-use the same connection, so we can't have any concurrency.
             var concurrency = inMemory ? 1 : cfg.GetCVar(CCVars.DatabaseSqliteConcurrency);
             _prefsSemaphore = new SemaphoreSlim(concurrency, concurrency);

Content.Shared/CCVar/CCVars.cs

@@ -479,11 +479,16 @@ namespace Content.Shared.CCVar
         public static readonly CVarDef<string> DatabasePgPassword =
             CVarDef.Create("database.pg_password", "", CVar.SERVERONLY | CVar.CONFIDENTIAL);
 
+        /// <summary>
+        /// Max amount of concurrent Postgres database operations.
+        /// </summary>
+        public static readonly CVarDef<int> DatabasePgConcurrency =
+            CVarDef.Create("database.pg_concurrency", 8, CVar.SERVERONLY);
+
         // Basically only exists for integration tests to avoid race conditions.
         public static readonly CVarDef<bool> DatabaseSynchronous =
             CVarDef.Create("database.sync", false, CVar.SERVERONLY);
 
-
         /*
          * Outline
          */

Content.Tests/Server/Preferences/ServerDbSqliteTests.cs

@@ -10,6 +10,7 @@ using Content.Shared.Preferences;
 using Microsoft.Data.Sqlite;
 using Microsoft.EntityFrameworkCore;
 using NUnit.Framework;
+using Robust.Shared.Configuration;
 using Robust.Shared.Enums;
 using Robust.Shared.IoC;
 using Robust.Shared.Maths;
@@ -74,7 +75,7 @@ namespace Content.Tests.Server.Preferences
             var conn = new SqliteConnection("Data Source=:memory:");
             conn.Open();
             builder.UseSqlite(conn);
-            return new ServerDbSqlite(() => builder.Options, true);
+            return new ServerDbSqlite(() => builder.Options, true, IoCManager.Resolve<IConfigurationManager>());
         }
 
         [Test]