Split PoolManager into separate classes. (#19370)

Content.IntegrationTests/Pair/TestMapData.cs

@@ -0,0 +1,19 @@
+using Robust.Shared.GameObjects;
+using Robust.Shared.Map;
+using Robust.Shared.Map.Components;
+
+namespace Content.IntegrationTests.Pair;
+
+/// <summary>
+/// Simple data class that stored information about a map being used by a test.
+/// </summary>
+public sealed class TestMapData
+{
+    public EntityUid MapUid { get; set; }
+    public EntityUid GridUid { get; set; }
+    public MapId MapId { get; set; }
+    public MapGridComponent MapGrid { get; set; } = default!;
+    public EntityCoordinates GridCoords { get; set; }
+    public MapCoordinates MapCoords { get; set; }
+    public TileRef Tile { get; set; }
+}

Content.IntegrationTests/Pair/TestPair.Helpers.cs

@@ -0,0 +1,39 @@
+#nullable enable
+using System.Linq;
+using Robust.Shared.Map;
+
+namespace Content.IntegrationTests.Pair;
+
+// Contains misc helper functions to make writing tests easier.
+public sealed partial class TestPair
+{
+    /// <summary>
+    /// Creates a map, a grid, and a tile, and gives back references to them.
+    /// </summary>
+    public async Task<TestMapData> CreateTestMap()
+    {
+        await Server.WaitIdleAsync();
+        var tileDefinitionManager = Server.ResolveDependency<ITileDefinitionManager>();
+
+        var mapData = new TestMapData();
+        await Server.WaitPost(() =>
+        {
+            mapData.MapId = Server.MapMan.CreateMap();
+            mapData.MapUid = Server.MapMan.GetMapEntityId(mapData.MapId);
+            mapData.MapGrid = Server.MapMan.CreateGrid(mapData.MapId);
+            mapData.GridUid = mapData.MapGrid.Owner; // Fixing this requires an engine PR.
+            mapData.GridCoords = new EntityCoordinates(mapData.GridUid, 0, 0);
+            var plating = tileDefinitionManager["Plating"];
+            var platingTile = new Tile(plating.TileId);
+            mapData.MapGrid.SetTile(mapData.GridCoords, platingTile);
+            mapData.MapCoords = new MapCoordinates(0, 0, mapData.MapId);
+            mapData.Tile = mapData.MapGrid.GetAllTiles().First();
+        });
+        
+        if (Settings.Connected)
+            await RunTicksSync(10);
+
+        TestMap = mapData;
+        return mapData;
+    }
+}

Content.IntegrationTests/Pair/TestPair.Prototypes.cs

@@ -0,0 +1,64 @@
+#nullable enable
+using System.Collections.Generic;
+using Robust.Shared.Prototypes;
+using Robust.Shared.Utility;
+using Robust.UnitTesting;
+
+namespace Content.IntegrationTests.Pair;
+
+// This partial class contains helper methods to deal with yaml prototypes.
+public sealed partial class TestPair
+{
+    private Dictionary<Type, HashSet<string>> _loadedPrototypes = new();
+    private HashSet<string> _loadedEntityPrototypes = new();
+
+    public async Task LoadPrototypes(List<string> prototypes)
+    {
+        await LoadPrototypes(Server, prototypes);
+        await LoadPrototypes(Client, prototypes);
+    }
+
+    private async Task LoadPrototypes(RobustIntegrationTest.IntegrationInstance instance, List<string> prototypes)
+    {
+        var changed = new Dictionary<Type, HashSet<string>>();
+        foreach (var file in prototypes)
+        {
+            instance.ProtoMan.LoadString(file, changed: changed);
+        }
+
+        await instance.WaitPost(() => instance.ProtoMan.ReloadPrototypes(changed));
+
+        foreach (var (kind, ids) in changed)
+        {
+            _loadedPrototypes.GetOrNew(kind).UnionWith(ids);
+        }
+
+        if (_loadedPrototypes.TryGetValue(typeof(EntityPrototype), out var entIds))
+            _loadedEntityPrototypes.UnionWith(entIds);
+    }
+
+    public bool IsTestPrototype(EntityPrototype proto)
+    {
+        return _loadedEntityPrototypes.Contains(proto.ID);
+    }
+
+    public bool IsTestEntityPrototype(string id)
+    {
+        return _loadedEntityPrototypes.Contains(id);
+    }
+
+    public bool IsTestPrototype<TPrototype>(string id) where TPrototype : IPrototype
+    {
+        return IsTestPrototype(typeof(TPrototype), id);
+    }
+
+    public bool IsTestPrototype<TPrototype>(TPrototype proto) where TPrototype : IPrototype
+    {
+        return IsTestPrototype(typeof(TPrototype), proto.ID);
+    }
+
+    public bool IsTestPrototype(Type kind, string id)
+    {
+        return _loadedPrototypes.TryGetValue(kind, out var ids) && ids.Contains(id);
+    }
+}

Content.IntegrationTests/Pair/TestPair.Recycle.cs

@@ -0,0 +1,218 @@
+#nullable enable
+using System.IO;
+using System.Linq;
+using Content.Server.GameTicking;
+using Content.Server.Mind.Components;
+using Content.Shared.CCVar;
+using Content.Shared.GameTicking;
+using Robust.Client;
+using Robust.Server.Player;
+using Robust.Shared.Exceptions;
+using Robust.Shared.GameObjects;
+using Robust.Shared.Network;
+
+namespace Content.IntegrationTests.Pair;
+
+// This partial class contains logic related to recycling & disposing test pairs.
+public sealed partial class TestPair : IAsyncDisposable
+{
+    public PairState State { get; private set; } = PairState.Ready;
+
+    private async Task OnDirtyDispose()
+    {
+        var usageTime = Watch.Elapsed;
+        Watch.Restart();
+        await _testOut.WriteLineAsync($"{nameof(DisposeAsync)}: Test gave back pair {Id} in {usageTime.TotalMilliseconds} ms");
+        Kill();
+        var disposeTime = Watch.Elapsed;
+        await _testOut.WriteLineAsync($"{nameof(DisposeAsync)}: Disposed pair {Id} in {disposeTime.TotalMilliseconds} ms");
+        // Test pairs should only dirty dispose if they are failing. If they are not failing, this probably happened
+        // because someone forgot to clean-return the pair.
+        Assert.Warn("Test was dirty-disposed.");
+    }
+
+    private async Task OnCleanDispose()
+    {
+        if (TestMap != null)
+        {
+            await Server.WaitPost(() => Server.EntMan.DeleteEntity(TestMap.MapUid));
+            TestMap = null;
+        }
+   
+        var usageTime = Watch.Elapsed;
+        Watch.Restart();
+        await _testOut.WriteLineAsync($"{nameof(CleanReturnAsync)}: Test borrowed pair {Id} for {usageTime.TotalMilliseconds} ms");
+        // Let any last minute failures the test cause happen.
+        await ReallyBeIdle();
+        if (!Settings.Destructive)
+        {
+            if (Client.IsAlive == false)
+            {
+                throw new Exception($"{nameof(CleanReturnAsync)}: Test killed the client in pair {Id}:", Client.UnhandledException);
+            }
+
+            if (Server.IsAlive == false)
+            {
+                throw new Exception($"{nameof(CleanReturnAsync)}: Test killed the server in pair {Id}:", Server.UnhandledException);
+            }
+        }
+
+        if (Settings.MustNotBeReused)
+        {
+            Kill();
+            await ReallyBeIdle();
+            await _testOut.WriteLineAsync($"{nameof(CleanReturnAsync)}: Clean disposed in {Watch.Elapsed.TotalMilliseconds} ms");
+            return;
+        }
+
+        var sRuntimeLog = Server.ResolveDependency<IRuntimeLog>();
+        if (sRuntimeLog.ExceptionCount > 0)
+            throw new Exception($"{nameof(CleanReturnAsync)}: Server logged exceptions");
+        var cRuntimeLog = Client.ResolveDependency<IRuntimeLog>();
+        if (cRuntimeLog.ExceptionCount > 0)
+            throw new Exception($"{nameof(CleanReturnAsync)}: Client logged exceptions");
+
+        var returnTime = Watch.Elapsed;
+        await _testOut.WriteLineAsync($"{nameof(CleanReturnAsync)}: PoolManager took {returnTime.TotalMilliseconds} ms to put pair {Id} back into the pool");
+    }
+
+    public async ValueTask CleanReturnAsync()
+    {
+        if (State != PairState.InUse)
+            throw new Exception($"{nameof(CleanReturnAsync)}: Unexpected state. Pair: {Id}. State: {State}.");
+        
+        await _testOut.WriteLineAsync($"{nameof(CleanReturnAsync)}: Return of pair {Id} started");
+        State = PairState.CleanDisposed;
+        await OnCleanDispose();
+        State = PairState.Ready;
+        PoolManager.NoCheckReturn(this);
+        ClearContext();
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        switch (State)
+        {
+            case PairState.Dead:
+            case PairState.Ready:
+                break;
+            case PairState.InUse:
+                await _testOut.WriteLineAsync($"{nameof(DisposeAsync)}: Dirty return of pair {Id} started");
+                await OnDirtyDispose();
+                PoolManager.NoCheckReturn(this);
+                ClearContext();
+                break;
+            default:
+                throw new Exception($"{nameof(DisposeAsync)}: Unexpected state. Pair: {Id}. State: {State}.");
+        }
+    }
+    
+    public async Task CleanPooledPair(PoolSettings settings, TextWriter testOut)
+    {
+        Settings = default!;
+        Watch.Restart();
+        await testOut.WriteLineAsync($"Recycling...");
+
+        var gameTicker = Server.System<GameTicker>();
+        var cNetMgr = Client.ResolveDependency<IClientNetManager>();
+
+        await RunTicksSync(1);
+
+        // Disconnect the client if they are connected.
+        if (cNetMgr.IsConnected)
+        {
+            await testOut.WriteLineAsync($"Recycling: {Watch.Elapsed.TotalMilliseconds} ms: Disconnecting client.");
+            await Client.WaitPost(() => cNetMgr.ClientDisconnect("Test pooling cleanup disconnect"));
+            await RunTicksSync(1);
+        }
+        Assert.That(cNetMgr.IsConnected, Is.False);
+
+        // Move to pre-round lobby. Required to toggle dummy ticker on and off
+        if (gameTicker.RunLevel != GameRunLevel.PreRoundLobby)
+        {
+            await testOut.WriteLineAsync($"Recycling: {Watch.Elapsed.TotalMilliseconds} ms: Restarting server.");
+            Assert.That(gameTicker.DummyTicker, Is.False);
+            Server.CfgMan.SetCVar(CCVars.GameLobbyEnabled, true);
+            await Server.WaitPost(() => gameTicker.RestartRound());
+            await RunTicksSync(1);
+        }
+
+        //Apply Cvars
+        await testOut.WriteLineAsync($"Recycling: {Watch.Elapsed.TotalMilliseconds} ms: Setting CVar ");
+        await PoolManager.SetupCVars(Client, settings);
+        await PoolManager.SetupCVars(Server, settings);
+        await RunTicksSync(1);
+
+        // Restart server.
+        await testOut.WriteLineAsync($"Recycling: {Watch.Elapsed.TotalMilliseconds} ms: Restarting server again");
+        await Server.WaitPost(() => gameTicker.RestartRound());
+        await RunTicksSync(1);
+
+        // Connect client
+        if (settings.ShouldBeConnected)
+        {
+            await testOut.WriteLineAsync($"Recycling: {Watch.Elapsed.TotalMilliseconds} ms: Connecting client");
+            Client.SetConnectTarget(Server);
+            await Client.WaitPost(() => cNetMgr.ClientConnect(null!, 0, null!));
+        }
+
+        await testOut.WriteLineAsync($"Recycling: {Watch.Elapsed.TotalMilliseconds} ms: Idling");
+        await ReallyBeIdle();
+        await testOut.WriteLineAsync($"Recycling: {Watch.Elapsed.TotalMilliseconds} ms: Done recycling");
+    }
+    
+    public void ValidateSettings(PoolSettings settings)
+    {
+        var cfg = Server.CfgMan;
+        Assert.That(cfg.GetCVar(CCVars.AdminLogsEnabled), Is.EqualTo(settings.AdminLogsEnabled));
+        Assert.That(cfg.GetCVar(CCVars.GameLobbyEnabled), Is.EqualTo(settings.InLobby));
+        Assert.That(cfg.GetCVar(CCVars.GameDummyTicker), Is.EqualTo(settings.UseDummyTicker));
+
+        var entMan = Server.ResolveDependency<EntityManager>();
+        var ticker = entMan.System<GameTicker>();
+        Assert.That(ticker.DummyTicker, Is.EqualTo(settings.UseDummyTicker));
+
+        var expectPreRound = settings.InLobby | settings.DummyTicker;
+        var expectedLevel = expectPreRound ? GameRunLevel.PreRoundLobby : GameRunLevel.InRound;
+        Assert.That(ticker.RunLevel, Is.EqualTo(expectedLevel));
+
+        var baseClient = Client.ResolveDependency<IBaseClient>();
+        var netMan = Client.ResolveDependency<INetManager>();
+        Assert.That(netMan.IsConnected, Is.Not.EqualTo(!settings.ShouldBeConnected));
+
+        if (!settings.ShouldBeConnected)
+            return;
+
+        Assert.That(baseClient.RunLevel, Is.EqualTo(ClientRunLevel.InGame));
+        var cPlayer = Client.ResolveDependency<Robust.Client.Player.IPlayerManager>();
+        var sPlayer = Server.ResolveDependency<IPlayerManager>();
+        Assert.That(sPlayer.Sessions.Count(), Is.EqualTo(1));
+        var session = sPlayer.Sessions.Single();
+        Assert.That(cPlayer.LocalPlayer?.Session.UserId, Is.EqualTo(session.UserId));
+
+        if (ticker.DummyTicker)
+            return;
+
+        var status = ticker.PlayerGameStatuses[session.UserId];
+        var expected = settings.InLobby
+            ? PlayerGameStatus.NotReadyToPlay
+            : PlayerGameStatus.JoinedGame;
+
+        Assert.That(status, Is.EqualTo(expected));
+
+        if (settings.InLobby)
+        {
+            Assert.Null(session.AttachedEntity);
+            return;
+        }
+
+        Assert.NotNull(session.AttachedEntity);
+        Assert.That(entMan.EntityExists(session.AttachedEntity));
+        Assert.That(entMan.HasComponent<MindContainerComponent>(session.AttachedEntity));
+        var mindCont = entMan.GetComponent<MindContainerComponent>(session.AttachedEntity!.Value);
+        Assert.NotNull(mindCont.Mind);
+        Assert.Null(mindCont.Mind?.VisitingEntity);
+        Assert.That(mindCont.Mind!.OwnedEntity, Is.EqualTo(session.AttachedEntity!.Value));
+        Assert.That(mindCont.Mind.UserId, Is.EqualTo(session.UserId));
+    }
+}

Content.IntegrationTests/Pair/TestPair.Timing.cs

@@ -0,0 +1,62 @@
+#nullable enable
+using Robust.Shared.Timing;
+
+namespace Content.IntegrationTests.Pair;
+
+// This partial class contains methods for running the server/client pairs for some number of ticks
+public sealed partial class TestPair
+{
+    /// <summary>
+    /// Runs the server-client pair in sync
+    /// </summary>
+    /// <param name="ticks">How many ticks to run them for</param>
+    public async Task RunTicksSync(int ticks)
+    {
+        for (var i = 0; i < ticks; i++)
+        {
+            await Server.WaitRunTicks(1);
+            await Client.WaitRunTicks(1);
+        }
+    }
+
+    /// <summary>
+    /// Runs the server-client pair in sync, but also ensures they are both idle each tick.
+    /// </summary>
+    /// <param name="runTicks">How many ticks to run</param>
+    public async Task ReallyBeIdle(int runTicks = 25)
+    {
+        for (var i = 0; i < runTicks; i++)
+        {
+            await Client.WaitRunTicks(1);
+            await Server.WaitRunTicks(1);
+            for (var idleCycles = 0; idleCycles < 4; idleCycles++)
+            {
+                await Client.WaitIdleAsync();
+                await Server.WaitIdleAsync();
+            }
+        }
+    }
+
+    /// <summary>
+    /// Run the server/clients until the ticks are synchronized.
+    /// By default the client will be one tick ahead of the server.
+    /// </summary>
+    public async Task SyncTicks(int targetDelta = 1)
+    {
+        var sTick = (int)Server.Timing.CurTick.Value;
+        var cTick = (int)Client.Timing.CurTick.Value;
+        var delta = cTick - sTick;
+
+        if (delta == targetDelta)
+            return;
+        if (delta > targetDelta)
+            await Server.WaitRunTicks(delta - targetDelta);
+        else
+            await Client.WaitRunTicks(targetDelta - delta);
+
+        sTick = (int)Server.Timing.CurTick.Value;
+        cTick = (int)Client.Timing.CurTick.Value;
+        delta = cTick - sTick;
+        Assert.That(delta, Is.EqualTo(targetDelta));
+    }
+}

Content.IntegrationTests/Pair/TestPair.cs

@@ -0,0 +1,111 @@
+#nullable enable
+using System.Collections.Generic;
+using System.IO;
+using Content.Server.GameTicking;
+using Robust.Shared.GameObjects;
+using Robust.Shared.IoC;
+using Robust.Shared.Network;
+using Robust.Shared.Timing;
+using Robust.UnitTesting;
+
+namespace Content.IntegrationTests.Pair;
+
+/// <summary>
+/// This object wraps a pooled server+client pair.
+/// </summary>
+public sealed partial class TestPair
+{
+    // TODO remove this.
+    [Obsolete("Field access is redundant")]
+    public TestPair Pair => this;
+
+    public readonly int Id;
+    private bool _initialized;
+    private TextWriter _testOut = default!;
+    public readonly Stopwatch Watch = new();
+    public readonly List<string> TestHistory = new();
+    public PoolSettings Settings = default!;
+    public TestMapData? TestMap;
+    public RobustIntegrationTest.ServerIntegrationInstance Server { get; private set; } = default!;
+    public RobustIntegrationTest.ClientIntegrationInstance Client { get;  private set; } = default!;
+
+    public PoolTestLogHandler ServerLogHandler { get;  private set; } = default!;
+    public PoolTestLogHandler ClientLogHandler { get;  private set; } = default!;
+    
+    public TestPair(int id)
+    {
+        Id = id;
+    }
+
+    public async Task Initialize(PoolSettings settings, TextWriter testOut, List<string> testPrototypes)
+    {
+        if (_initialized)
+            throw new InvalidOperationException("Already initialized");
+
+        _initialized = true;
+        Settings = settings;
+        (Client, ClientLogHandler) = await PoolManager.GenerateClient(settings, testOut);
+        (Server, ServerLogHandler) = await PoolManager.GenerateServer(settings, testOut);
+        ActivateContext(testOut);
+
+        if (!settings.NoLoadTestPrototypes)
+            await LoadPrototypes(testPrototypes!);
+
+        if (!settings.UseDummyTicker)
+        {
+            var gameTicker = Server.ResolveDependency<IEntityManager>().System<GameTicker>();
+            await Server.WaitPost(() => gameTicker.RestartRound());
+        }
+
+        if (settings.ShouldBeConnected)
+        {
+            Client.SetConnectTarget(Server);
+            await Client.WaitPost(() =>
+            {
+                var netMgr = IoCManager.Resolve<IClientNetManager>();
+                if (!netMgr.IsConnected)
+                {
+                    netMgr.ClientConnect(null!, 0, null!);
+                }
+            });
+            await ReallyBeIdle(10);
+            await Client.WaitRunTicks(1);
+        }
+    }
+    
+    public void Kill()
+    {
+        State = PairState.Dead;
+        Server.Dispose();
+        Client.Dispose();
+    }
+
+    private void ClearContext()
+    {
+        _testOut = default!;
+        ServerLogHandler.ClearContext();
+        ClientLogHandler.ClearContext();
+    }
+
+    public void ActivateContext(TextWriter testOut)
+    {
+        _testOut = testOut;
+        ServerLogHandler.ActivateContext(testOut);
+        ClientLogHandler.ActivateContext(testOut);
+    }
+
+    public void Use()
+    {
+        if (State != PairState.Ready)
+            throw new InvalidOperationException($"Pair is not ready to use. State: {State}");
+        State = PairState.InUse;
+    }
+    
+    public enum PairState : byte
+    {
+        Ready = 0,
+        InUse = 1,
+        CleanDisposed = 2,
+        Dead = 3,
+    }
+}