ReferenceCountingContainer.cs

Go to the documentation of this file.

using System;
using System.Threading.Tasks;
 
namespace Tgstation.Server.Host.Utils
{
    sealed class ReferenceCountingContainer<TWrapped, TReference>
        where TWrapped : class
        where TReference : ReferenceCounter<TWrapped>, new()
    {
        public TWrapped Instance { get; }
 
        public Task OnZeroReferences
        {
            get
            {
                lock (referenceCountLock)
                {
                    if (referenceCount == 0)
                        return Task.CompletedTask;
                    return onZeroReferencesTcs!.Task;
                }
            }
        }
 
        readonly object referenceCountLock;
 
        TaskCompletionSource? onZeroReferencesTcs;
 
        ulong referenceCount;
 
        public ReferenceCountingContainer(TWrapped instance)
        {
            Instance = instance ?? throw new ArgumentNullException(nameof(instance));
 
            referenceCountLock = new object();
        }
 
        public TReference AddReference()
        {
            lock (referenceCountLock)
            {
                if (referenceCount++ == 0)
                    onZeroReferencesTcs = new TaskCompletionSource();
 
                try
                {
                    var reference = new TReference();
                    reference.Initialize(Instance, () =>
                    {
                        lock (referenceCountLock)
                            if (--referenceCount == 0)
                                onZeroReferencesTcs!.SetResult();
                    });
                    return reference;
                }
                catch
                {
                    --referenceCount;
                    throw;
                }
            }
        }
    }
}