stgsharp/StgSharp.Common/Threading/SingleConsumerBell.cs

// -----------------------------------------------------------------------------
// file="SingleConsumerBell"
// Project: StgSharp
// Copyright (c) Nitload.
// SPDX-License-Identifier: MIT
// -----------------------------------------------------------------------------

using System;
using System.Threading;

namespace StgSharp.Threading
{
    /// <summary>
    ///   A single-consumer bell/signal primitive.
    /// </summary>
    /// <remarks>
    ///   <para>Producers can call <see cref="Set" /> multiple times to publish signals. The bell
    ///   keeps a pending count of not-yet-consumed signals. The consumer uses <see
    ///   cref="Wait(object, CancellationToken)" /> or <see cref="TryWait(object)" /> to consume
    ///   signals. </para> <para>When a producer increments the pending count from 0 to 1, the
    ///   internal event is set to wake the consumer. Consuming via <see cref="TryWait(object)" />
    ///   or <see cref="Wait(object, CancellationToken)" /> clears the entire pending count
    ///   (aggregate consumption). If new signals arrive while resetting, the event is set again to
    ///   avoid missed signals. </para> <para>Ownership: only the specified <c> owner </c> can call
    ///   the consumer APIs. </para>
    /// </remarks>
    public sealed class SingleConsumerBell : IDisposable
    {

        private volatile bool _disposed;
        private int _pending; // not volatile: used with Interlocked/Volatile APIs to avoid CS0420
        private readonly ManualResetEventSlim _evt;
        private readonly object _owner;

        /// <summary>
        ///   Initializes a new instance of <see cref="SingleConsumerBell" />.
        /// </summary>
        /// <param name="owner">
        ///   The unique consumer identity; subsequent calls to Wait/TryWait must pass the same
        ///   reference.
        /// </param>
        /// <param name="spinCount">
        ///   Spin count passed to the underlying <see cref="ManualResetEventSlim" />.
        /// </param>
        /// <exception cref="ArgumentNullException">
        ///   Thrown if <paramref name="owner" /> is null.
        /// </exception>
        public SingleConsumerBell(
               object owner,
               int spinCount = 100
        )
        {
            _owner = owner ?? throw new ArgumentNullException(nameof(owner));
            _evt = new ManualResetEventSlim(false, spinCount);
        }

        /// <summary>
        ///   Disposes the bell, sets the event to wake any waiting consumer, and releases
        ///   resources.
        /// </summary>
        public void Dispose()
        {
            _disposed = true;
            _evt.Set();
            _evt.Dispose();
        }

        /// <summary>
        ///   Publishes a signal by incrementing the not-yet-consumed count and returns the current
        ///   pending count.
        /// </summary>
        /// <remarks>
        ///   If the pending count transitions from 0 to 1, the event is set to wake the consumer.
        /// </remarks>
        /// <returns>
        ///   The current pending count including this call.
        /// </returns>
        public int Set()
        {
            if (_disposed) {
                return Volatile.Read(ref _pending);
            }

            // Increment pending; if we transitioned 0 -> 1, set the event to wake the consumer
            int newCount = Interlocked.Increment(ref _pending);
            if (newCount == 1) {
                _evt.Set();
            }
            return newCount;
        }

        /// <summary>
        ///   Attempts to consume any pending signals without blocking.
        /// </summary>
        /// <param name="consumer">
        ///   The consumer object; must be the same reference provided at construction.
        /// </param>
        /// <returns>
        ///   <c> true </c> if a signal was available (pending or event set); otherwise
        ///   <c>false</c>.
        /// </returns>
        /// <exception cref="InvalidOperationException">
        ///   Thrown if <paramref name="consumer" /> is not the owner.
        /// </exception>
        public bool TryWait(
                    object consumer
        )
        {
            if (!ReferenceEquals(_owner, consumer)) {
                throw new InvalidOperationException("This bell is owned by another consumer.");
            }

            if (!_evt.IsSet && Volatile.Read(ref _pending) == 0) {
                return false;
            }

            // Aggregate consumption: clear pending and reset event
            _ = Interlocked.Exchange(ref _pending, 0);
            _evt.Reset();

            // If concurrent Set happened during reset, re-set the event to avoid lost signals
            if (Volatile.Read(ref _pending) != 0)
            {
                _evt.Set();
            }
            return true;
        }

        /// <summary>
        ///   Blocks until a signal is available, then consumes all pending signals.
        /// </summary>
        /// <param name="consumer">
        ///   The consumer object; must be the same reference provided at construction.
        /// </param>
        /// <param name="cancellationToken">
        ///   A token to cancel the wait.
        /// </param>
        /// <exception cref="InvalidOperationException">
        ///   Thrown if <paramref name="consumer" /> is not the owner.
        /// </exception>
        public void Wait(
                    object consumer,
                    CancellationToken cancellationToken = default
        )
        {
            if (!ReferenceEquals(_owner, consumer)) {
                throw new InvalidOperationException("This bell is owned by another consumer.");
            }

            _evt.Wait(cancellationToken);

            // Consume this round of signals: clear pending and handle reset/concurrency
            _ = Interlocked.Exchange(ref _pending, 0);

            _evt.Reset();

            if (Volatile.Read(ref _pending) != 0) {
                _evt.Set();
            }
        }

    }
}