diff --git a/src/Tgstation.Server.Host/Components/InstanceContainer.cs b/src/Tgstation.Server.Host/Components/InstanceContainer.cs
deleted file mode 100644
index 6a8c49500e..0000000000
--- a/src/Tgstation.Server.Host/Components/InstanceContainer.cs
+++ /dev/null
@@ -1,86 +0,0 @@
-using System;
-using System.Threading.Tasks;
-
-namespace Tgstation.Server.Host.Components
-{
- ///
- /// Wrapper for managing s.
- ///
- sealed class InstanceContainer
- {
- ///
- /// The .
- ///
- public IInstance Instance { get; }
-
- ///
- /// A that completes when there are no s active for the .
- ///
- public Task OnZeroReferences
- {
- get
- {
- lock (referenceCountLock)
- {
- if (referenceCount == 0)
- return Task.CompletedTask;
- return onZeroReferencesTcs.Task;
- }
- }
- }
-
- ///
- /// for .
- ///
- readonly object referenceCountLock;
-
- ///
- /// Backing for .
- ///
- TaskCompletionSource onZeroReferencesTcs;
-
- ///
- /// Count of active s.
- ///
- ulong referenceCount;
-
- ///
- /// Initializes a new instance of the class.
- ///
- /// The value of .
- public InstanceContainer(IInstance instance)
- {
- Instance = instance ?? throw new ArgumentNullException(nameof(instance));
-
- referenceCountLock = new object();
- }
-
- ///
- /// Create a new .
- ///
- /// A new .
- public IInstanceReference AddReference()
- {
- lock (referenceCountLock)
- {
- if (referenceCount++ == 0)
- onZeroReferencesTcs = new TaskCompletionSource();
-
- try
- {
- return new InstanceWrapper(Instance, () =>
- {
- lock (referenceCountLock)
- if (--referenceCount == 0)
- onZeroReferencesTcs.SetResult();
- });
- }
- catch
- {
- --referenceCount;
- throw;
- }
- }
- }
- }
-}
diff --git a/src/Tgstation.Server.Host/Components/InstanceManager.cs b/src/Tgstation.Server.Host/Components/InstanceManager.cs
index a9d5c57a04..1b7bed2dd4 100644
--- a/src/Tgstation.Server.Host/Components/InstanceManager.cs
+++ b/src/Tgstation.Server.Host/Components/InstanceManager.cs
@@ -93,14 +93,14 @@ namespace Tgstation.Server.Host.Components
readonly ILogger logger;
///
- /// Map of instance s to respective s. Also used as a .
+ /// Map of instance s to the respective for s. Also used as a .
///
- readonly IDictionary instances;
+ readonly Dictionary> instances;
///
/// Map of s to their respective s.
///
- readonly IDictionary bridgeHandlers;
+ readonly Dictionary bridgeHandlers;
///
/// used to guard calls to and .
@@ -182,7 +182,7 @@ namespace Tgstation.Server.Host.Components
swarmConfiguration = swarmConfigurationOptions?.Value ?? throw new ArgumentNullException(nameof(swarmConfigurationOptions));
this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
- instances = new Dictionary();
+ instances = new Dictionary>();
bridgeHandlers = new Dictionary();
readyTcs = new TaskCompletionSource();
instanceStateChangeSemaphore = new SemaphoreSlim(1);
@@ -308,7 +308,7 @@ namespace Tgstation.Server.Host.Components
using var lockContext = await SemaphoreSlimContext.Lock(instanceStateChangeSemaphore, cancellationToken);
logger.LogInformation("Offlining instance ID {instanceId}", metadata.Id);
- InstanceContainer container;
+ ReferenceCountingContainer container;
lock (instances)
{
if (!instances.TryGetValue(metadata.Id.Value, out container))
@@ -375,7 +375,11 @@ namespace Tgstation.Server.Host.Components
try
{
lock (instances)
- instances.Add(metadata.Id.Value, new InstanceContainer(instance));
+ instances.Add(
+ metadata.Id.Value,
+ new ReferenceCountingContainer(
+ instance,
+ (wrappedInstance, disposeAction) => new InstanceWrapper(wrappedInstance, disposeAction)));
}
catch (Exception ex)
{
diff --git a/src/Tgstation.Server.Host/Utils/ReferenceCountingContainer.cs b/src/Tgstation.Server.Host/Utils/ReferenceCountingContainer.cs
new file mode 100644
index 0000000000..32ab3e8031
--- /dev/null
+++ b/src/Tgstation.Server.Host/Utils/ReferenceCountingContainer.cs
@@ -0,0 +1,96 @@
+using System;
+using System.Threading.Tasks;
+
+namespace Tgstation.Server.Host.Utils
+{
+ ///
+ /// Wrapper for managing some .
+ ///
+ /// The type being wrapped.
+ /// The disposable reference type returned.
+ sealed class ReferenceCountingContainer
+ where TReference : IDisposable
+ {
+ ///
+ /// The .
+ ///
+ public TWrapped Instance { get; }
+
+ ///
+ /// A that completes when there are no s active for the .
+ ///
+ public Task OnZeroReferences
+ {
+ get
+ {
+ lock (referenceCountLock)
+ {
+ if (referenceCount == 0)
+ return Task.CompletedTask;
+ return onZeroReferencesTcs.Task;
+ }
+ }
+ }
+
+ ///
+ /// The factory for generating s to the .
+ ///
+ readonly Func referenceFactory;
+
+ ///
+ /// for .
+ ///
+ readonly object referenceCountLock;
+
+ ///
+ /// Backing for .
+ ///
+ TaskCompletionSource onZeroReferencesTcs;
+
+ ///
+ /// Count of active s.
+ ///
+ ulong referenceCount;
+
+ ///
+ /// Initializes a new instance of the class.
+ ///
+ /// The value of .
+ /// The value of .
+ public ReferenceCountingContainer(TWrapped instance, Func referenceFactory)
+ {
+ Instance = instance ?? throw new ArgumentNullException(nameof(instance));
+ this.referenceFactory = referenceFactory ?? throw new ArgumentNullException(nameof(referenceFactory));
+
+ referenceCountLock = new object();
+ }
+
+ ///
+ /// Create a new to the .
+ ///
+ /// A new .
+ public TReference AddReference()
+ {
+ lock (referenceCountLock)
+ {
+ if (referenceCount++ == 0)
+ onZeroReferencesTcs = new TaskCompletionSource();
+
+ try
+ {
+ return referenceFactory(Instance, () =>
+ {
+ lock (referenceCountLock)
+ if (--referenceCount == 0)
+ onZeroReferencesTcs.SetResult();
+ });
+ }
+ catch
+ {
+ --referenceCount;
+ throw;
+ }
+ }
+ }
+ }
+}