Improved end-of-buffer detection for the producer side

Author: sakno

src/DotNext.Threading/Collections/Concurrent/RingBuffer.cs

@@ -1,5 +1,4 @@
 using System.Diagnostics;
-using System.Diagnostics.Contracts;
 using System.Numerics;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
@@ -13,21 +12,22 @@ namespace DotNext.Collections.Concurrent;
 internal struct RingBuffer<T>
 {
     private readonly Slot[] slots;
+    private readonly nuint indexMask;
+    private readonly int indexBits;
     private bool frozenForEnqueues;
-    private RingBuffer state;
+    private State state;
 
-    public RingBuffer(int maximumRetained)
+    public RingBuffer(int desiredSize)
     {
-        Debug.Assert(maximumRetained > 0);
+        Debug.Assert(desiredSize > 0);
 
-        var length = nuint.CreateChecked(BitOperations.RoundUpToPowerOf2((ulong)(uint)maximumRetained));
+        var length = nuint.CreateChecked(BitOperations.RoundUpToPowerOf2((ulong)(uint)desiredSize));
         slots = new Slot[length];
-        state = new(length);
+        indexMask = length - 1U;
+        indexBits = int.CreateChecked(nuint.Log2(length));
         frozenForEnqueues = false;
     }
-
-    public readonly bool IsFrozen => Volatile.Read(in frozenForEnqueues);
-
+    
     public bool Freeze()
     {
         var frozen = Interlocked.FalseToTrue(ref frozenForEnqueues);
@@ -39,32 +39,36 @@ public bool Freeze()
 
     private void WaitForPendingEnqueues()
     {
-        Debug.Assert(frozenForEnqueues);
+        Debug.Assert(IsFrozen);
 
         // the slots prior to the frozen position can be still in progress by the enqueuer, so wait for it
         for (var position = FreezeProducer() - 1U;
-             this[state.GetIndex(position)].WaitForPendingEnqueue(state.GetGeneration(position));
+             this[GetIndex(position)].WaitForPendingEnqueue(GetGeneration(position));
              position--) ;
     }
 
     private nuint FreezeProducer()
     {
-        Debug.Assert(frozenForEnqueues);
+        Debug.Assert(IsFrozen);
 
         const nuint shift = 2U;
-        var current = state.Positions.Producer;
-        for (nuint tmp, offset = (uint)slots.Length * shift;; current = tmp)
+        var current = state.Producer;
+        for (nuint tmp, offset = Array.GetLength(slots) * shift;; current = tmp)
         {
             // Advances the producer too far forward, so this position cannot be reached naturally even
             // if the buffer is full. 2 generations forward is enough, because the items in the buffer
             // can be in the current and the next generation.
-            tmp = Interlocked.CompareExchange(ref state.Positions.Producer, current + offset, current);
+            tmp = Interlocked.CompareExchange(ref state.Producer, current + offset, current);
             if (tmp == current)
                 break;
         }
 
         return current;
     }
+    
+    private static nuint StateBit => (nuint)nint.MinValue;
+
+    public readonly bool IsFrozen => Volatile.Read(in frozenForEnqueues);
 
     public readonly int Length => slots.Length;
 
@@ -79,40 +83,40 @@ private readonly ref Slot this[nuint index]
     }
 
     public ref Slot TryDequeue(out nuint sequence)
-        => ref DoOperation<RingBuffer.DequeueOperation>(ref state.Positions.Consumer, out sequence);
+        => ref DoOperation<DequeueOperation>(out sequence);
 
     public readonly bool IsEmpty
     {
         get
         {
-            var position = state.Positions.Consumer;
-            var generation = state.GetGeneration(position);
+            var position = Consumer;
+            var generation = GetGeneration(position);
 
-            return this[state.GetIndex(position)].Sequence != (generation | RingBuffer.StateBit);
+            return this[GetIndex(position)].Sequence != (generation | StateBit);
         }
     }
 
     public ref Slot TryEnqueue(out nuint sequence)
-        => ref DoOperation<RingBuffer.EnqueueOperation>(ref state.Positions.Producer, out sequence);
-
-    private readonly ref Slot DoOperation<TOperation>(scoped ref nuint position, out nuint newSeq)
-        where TOperation : struct, RingBuffer.IOperation<TOperation>, allows ref struct
+        => ref DoOperation<EnqueueOperation>(out sequence);
+    
+    private ref Slot DoOperation<TOperation>(out nuint newSeq)
+        where TOperation : struct, IOperation<TOperation>, allows ref struct
     {
         var spinner = new SpinWait();
-        for (nuint positionCopy = Volatile.Read(in position), tmp; TOperation.Retry(in frozenForEnqueues); positionCopy = tmp)
+        for (nuint positionCopy = TOperation.GetPosition(ref state), tmp; TOperation.Retry(in frozenForEnqueues); positionCopy = tmp)
         {
-            ref var slot = ref this[state.GetIndex(positionCopy)];
-            var context = TOperation.Create(in state, positionCopy);
+            ref var slot = ref this[GetIndex(positionCopy)];
+            var context = TOperation.Create(GetGeneration(positionCopy));
 
             if (!context.IsValidSequence(slot.Sequence))
             {
-                if (!context.CanRetry(in state))
+                if (!TOperation.CanRetry(in this, positionCopy))
                     break;
 
                 tmp = positionCopy + 1U;
                 spinner.SpinOnce(sleep1Threshold: -1);
             }
-            else if ((tmp = Interlocked.CompareExchange(ref position, positionCopy + 1U, positionCopy)) == positionCopy)
+            else if ((tmp = Interlocked.CompareExchange(ref TOperation.GetPosition(ref state), positionCopy + 1U, positionCopy)) == positionCopy)
             {
                 newSeq = context.NextSequence;
                 return ref slot;
@@ -122,110 +126,107 @@ private readonly ref Slot DoOperation<TOperation>(scoped ref nuint position, out
         Unsafe.SkipInit(out newSeq);
         return ref Unsafe.NullRef<Slot>();
     }
-
-    [StructLayout(LayoutKind.Auto)]
-    public struct Slot
-    {
-        public T? Item;
-        public volatile nuint Sequence; // higher bit is reserved for the value presence
-
-        public readonly bool WaitForPendingEnqueue(nuint frozenGen)
-        {
-            // The slot can be in three states:
-            // 1. Enqueued, so Sequence == (frozenGen | StateBit) => skip it and check the previous slot
-            // 2. Dequeued, so Sequence == (frozenGen + 1) & ~StateBit => leave the method
-            // 3. In-flight, so Sequence == frozenGen => wait for Enqueued or Dequeued
-            nuint sequence;
-            for (var spinner = new SpinWait();
-                 (sequence = Sequence) == frozenGen;
-                 spinner.SpinOnce()) ;
-
-            return sequence == (frozenGen | RingBuffer.StateBit);
-        }
-    }
-}
-
-[StructLayout(LayoutKind.Sequential)]
-internal struct RingBuffer(nuint length)
-{
-    private readonly nuint indexMask = length - 1U;
-    private readonly int indexBits = int.CreateChecked(nuint.Log2(length));
-    public State Positions;
-
-    [Pure]
-    public readonly nuint GetGeneration(nuint position) => position >>> indexBits;
-
-    [Pure]
-    public readonly nuint GetIndex(nuint position) => position & indexMask;
 
-    public static nuint StateBit => (nuint)nint.MinValue;
+    private readonly nuint GetGeneration(nuint position) => position >>> indexBits;
 
-    public interface IOperation<out TSelf>
-        where TSelf : struct, IOperation<TSelf>, allows ref struct
+    private readonly nuint GetIndex(nuint position) => position & indexMask;
+    
+    private interface IOperation
     {
+        static abstract ref nuint GetPosition(ref State state);
+        
         static abstract bool Retry(ref readonly bool frozenForEnqueues);
 
-        bool CanRetry(scoped ref readonly RingBuffer state);
+        static abstract bool CanRetry(ref readonly RingBuffer<T> state, nuint position);
 
         nuint NextSequence { get; }
 
         bool IsValidSequence(nuint sequence);
+    }
 
-        public static abstract TSelf Create(scoped ref readonly RingBuffer state, nuint position);
+    private interface IOperation<out TSelf> : IOperation
+        where TSelf : struct, IOperation<TSelf>, allows ref struct
+    {
+        public static abstract TSelf Create(nuint generation);
     }
 
     [StructLayout(LayoutKind.Auto)]
-    public readonly ref struct EnqueueOperation : IOperation<EnqueueOperation>
+    private readonly ref struct EnqueueOperation : IOperation<EnqueueOperation>
     {
-        private readonly nuint generation, position;
-
-        private EnqueueOperation(scoped ref readonly RingBuffer state, nuint producerPosition)
-            => generation = state.GetGeneration(position = producerPosition);
+        private readonly nuint generation;
 
-        bool IOperation<EnqueueOperation>.CanRetry(ref readonly RingBuffer state)
-        {
-            var consumerPos = Volatile.Read(in state.Positions.Consumer);
+        private EnqueueOperation(nuint generation) => this.generation = generation;
 
-            // the consumer must be in the same generation as the producer, or producer position must be less than the consumer position
-            return state.GetGeneration(consumerPos) == generation || state.GetIndex(position) < state.GetIndex(consumerPos);
-        }
+        static bool IOperation.CanRetry(ref readonly RingBuffer<T> state, nuint producerPosition)
+            => producerPosition != state.ConsumerNextGen;
 
-        static bool IOperation<EnqueueOperation>.Retry(ref readonly bool frozenForEnqueues) => !Volatile.Read(in frozenForEnqueues);
+        static bool IOperation.Retry(ref readonly bool frozenForEnqueues) => !Volatile.Read(in frozenForEnqueues);
 
         // the slot becomes available for consumption in the current generation
-        nuint IOperation<EnqueueOperation>.NextSequence => generation | StateBit;
+        nuint IOperation.NextSequence => generation | StateBit;
 
-        bool IOperation<EnqueueOperation>.IsValidSequence(nuint sequence) => sequence == generation;
+        bool IOperation.IsValidSequence(nuint sequence) => sequence == generation;
 
-        static EnqueueOperation IOperation<EnqueueOperation>.Create(scoped ref readonly RingBuffer state, nuint producerPosition)
-            => new(in state, producerPosition);
+        static EnqueueOperation IOperation<EnqueueOperation>.Create(nuint generation)
+            => new(generation);
+
+        static ref nuint IOperation.GetPosition(ref State state) => ref state.Producer;
 
         public override string ToString() => generation.ToString("X");
     }
 
     [StructLayout(LayoutKind.Auto)]
-    public readonly ref struct DequeueOperation : IOperation<DequeueOperation>
+    private readonly ref struct DequeueOperation : IOperation<DequeueOperation>
     {
-        private readonly nuint generation, position;
+        private readonly nuint generation;
 
-        private DequeueOperation(scoped ref readonly RingBuffer state, nuint consumerPosition)
-            => generation = state.GetGeneration(position = consumerPosition);
+        private DequeueOperation(nuint generation) => this.generation = generation;
 
-        static bool IOperation<DequeueOperation>.Retry(ref readonly bool frozenForEnqueues) => true;
+        static bool IOperation.Retry(ref readonly bool frozenForEnqueues) => true;
 
-        bool IOperation<DequeueOperation>.CanRetry(ref readonly RingBuffer state)
-            => position != Volatile.Read(in state.Positions.Producer);
+        static bool IOperation.CanRetry(ref readonly RingBuffer<T> state, nuint consumerPosition)
+            => consumerPosition != state.Producer;
 
         // the slot can be occupied in the next generation only
-        nuint IOperation<DequeueOperation>.NextSequence => (generation + 1U) & ~StateBit;
+        nuint IOperation.NextSequence => (generation + 1U) & ~StateBit;
 
-        bool IOperation<DequeueOperation>.IsValidSequence(nuint sequence) => sequence == (generation | StateBit);
+        bool IOperation.IsValidSequence(nuint sequence) => sequence == (generation | StateBit);
 
-        static DequeueOperation IOperation<DequeueOperation>.Create(scoped ref readonly RingBuffer state, nuint consumerPosition)
-            => new(in state, consumerPosition);
+        static DequeueOperation IOperation<DequeueOperation>.Create(nuint generation)
+            => new(generation);
+        
+        static ref nuint IOperation.GetPosition(ref State state) => ref state.Consumer;
 
         public override string ToString() => generation.ToString("X");
     }
+    
+    private readonly nuint Producer => Volatile.Read(in state.Producer);
+
+    private readonly nuint Consumer => Volatile.Read(in state.Consumer);
+
+    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    private readonly nuint ConsumerNextGen => Volatile.Read(in state.Consumer) + (uint)slots.Length;
+    
+    [StructLayout(LayoutKind.Auto)]
+    public struct Slot
+    {
+        public T? Item;
+        public volatile nuint Sequence; // higher bit is reserved for the value presence
+
+        public readonly bool WaitForPendingEnqueue(nuint frozenGen)
+        {
+            // The slot can be in three states:
+            // 1. Enqueued, so Sequence == (frozenGen | StateBit) => skip it and check the previous slot
+            // 2. Dequeued, so Sequence == (frozenGen + 1) & ~StateBit => leave the method
+            // 3. In-flight, so Sequence == frozenGen => wait for Enqueued or Dequeued
+            nuint sequence;
+            for (var spinner = new SpinWait();
+                 (sequence = Sequence) == frozenGen;
+                 spinner.SpinOnce()) ;
+
+            return sequence == (frozenGen | StateBit);
+        }
+    }
 }
 
 // producer/consumer positions are used by different threads, so it's better to avoid memory cache sharing