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; + } + } + } + } +}