Stop server from overwhelming in-case of any epoch table slowness.

libs/storage/Tsavorite/cs/src/core/ClientSession/BasicContext.cs

@@ -16,6 +16,8 @@
         where TStoreFunctions : IStoreFunctions<TKey, TValue>
         where TAllocator : IAllocator<TKey, TValue, TStoreFunctions>
     {
+        private const int kMaxSpins = 10_000;
+
         readonly ClientSession<TKey, TValue, TInput, TOutput, TContext, TFunctions, TStoreFunctions, TAllocator> clientSession;
         internal readonly SessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TFunctions, BasicSessionLocker<TKey, TValue, TStoreFunctions, TAllocator>, TStoreFunctions, TAllocator> sessionFunctions;
 
@@ -33,7 +35,7 @@
         /// <inheritdoc/>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void UnsafeResumeThread()
-            => clientSession.UnsafeResumeThread(sessionFunctions);
+            => clientSession.UnsafeResumeThread(sessionFunctions, kMaxSpins);
 
         /// <inheritdoc/>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -479,7 +481,7 @@
         /// <returns>Status</returns>
         internal Status ContainsKeyInMemory(ref TKey key, out long logicalAddress, long fromAddress = -1)
         {
-            UnsafeResumeThread();
+            UnsafeResumeThread(kMaxSpins);
             try
             {
                 return store.InternalContainsKeyInMemory<TInput, TOutput, TContext, SessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TFunctions, BasicSessionLocker<TKey, TValue, TStoreFunctions, TAllocator>, TStoreFunctions, TAllocator>>(

libs/storage/Tsavorite/cs/src/core/ClientSession/ClientSession.cs

@@ -153,7 +153,7 @@
         internal void Refresh<TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions)
             where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
             UnsafeResumeThread(sessionFunctions);
             try
             {
                 store.InternalRefresh<TInput, TOutput, TContext, TSessionFunctionsWrapper>(sessionFunctions);
@@ -168,7 +168,7 @@
         internal void ResetModified<TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions, ref TKey key)
             where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
             UnsafeResumeThread(sessionFunctions);
             try
             {
                 UnsafeResetModified(sessionFunctions, ref key);
@@ -231,7 +231,7 @@
         internal bool CompletePending<TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions, bool getOutputs, bool wait, bool spinWaitForCommit)
             where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
             UnsafeResumeThread(sessionFunctions);
             try
             {
                 return UnsafeCompletePending(sessionFunctions, getOutputs, wait, spinWaitForCommit);
@@ -333,7 +333,7 @@
         internal bool IsModified<TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions, ref TKey key)
             where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
             UnsafeResumeThread(sessionFunctions);
             try
             {
                 return UnsafeIsModified(sessionFunctions, ref key);
@@ -506,12 +506,12 @@
         /// Resume session on current thread. IMPORTANT: Call SuspendThread before any async op.
         /// </summary>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        internal void UnsafeResumeThread<TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions)
+        internal void UnsafeResumeThread<TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions, int maxSpins)
             where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
             // We do not track any "acquired" state here; if someone mixes calls between safe and unsafe contexts, they will 
             // get the "trying to acquire already-acquired epoch" error.
-            store.epoch.Resume();
+            store.epoch.Resume(maxSpins);
             store.InternalRefresh<TInput, TOutput, TContext, TSessionFunctionsWrapper>(sessionFunctions);
         }
 

libs/storage/Tsavorite/cs/src/core/ClientSession/LockableContext.cs

@@ -24,6 +24,7 @@ namespace Tsavorite.core
         public bool IsNull => clientSession is null;
 
         const int KeyLockMaxRetryAttempts = 1000;
+        const int MaxSpinCount = KeyLockMaxRetryAttempts;
 
         internal LockableContext(ClientSession<TKey, TValue, TInput, TOutput, TContext, TFunctions, TStoreFunctions, TAllocator> clientSession)
         {
@@ -235,7 +236,7 @@ public void Lock<TLockableKey>(ReadOnlySpan<TLockableKey> keys)
             var lockAcquired = false;
             while (!lockAcquired)
             {
-                clientSession.UnsafeResumeThread(sessionFunctions);
+                clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
                 try
                 {
                     lockAcquired = DoManualLock(sessionFunctions, clientSession, keys);
@@ -269,7 +270,7 @@ public bool TryLock<TLockableKey>(ReadOnlySpan<TLockableKey> keys, TimeSpan time
             clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected(), "Trying to protect an already-protected epoch for LockableUnsafeContext.Lock()");
 
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return DoManualTryLock(sessionFunctions, clientSession, keys, timeout, cancellationToken);
@@ -302,7 +303,7 @@ public bool TryPromoteLock<TLockableKey>(TLockableKey key, TimeSpan timeout, Can
             clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected(), "Trying to protect an already-protected epoch for LockableUnsafeContext.Lock()");
 
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return DoManualTryPromoteLock(sessionFunctions, clientSession, key, timeout, cancellationToken);
@@ -320,7 +321,7 @@ public void Unlock<TLockableKey>(ReadOnlySpan<TLockableKey> keys)
             clientSession.CheckIsAcquiredLockable(sessionFunctions);
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected(), "Trying to protect an already-protected epoch for LockableUnsafeContext.Unlock()");
 
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, -1);
             try
             {
                 DoManualUnlock(clientSession, keys);
@@ -353,7 +354,7 @@ public void Unlock<TLockableKey>(ReadOnlySpan<TLockableKey> keys)
         public bool CompletePending(bool wait = false, bool spinWaitForCommit = false)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.UnsafeCompletePending(sessionFunctions, false, wait, spinWaitForCommit);
@@ -368,7 +369,7 @@ public bool CompletePending(bool wait = false, bool spinWaitForCommit = false)
         public bool CompletePendingWithOutputs(out CompletedOutputIterator<TKey, TValue, TInput, TOutput, TContext> completedOutputs, bool wait = false, bool spinWaitForCommit = false)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.UnsafeCompletePendingWithOutputs(sessionFunctions, out completedOutputs, wait, spinWaitForCommit);
@@ -392,7 +393,7 @@ public ValueTask<CompletedOutputIterator<TKey, TValue, TInput, TOutput, TContext
         public Status Read(ref TKey key, ref TInput input, ref TOutput output, TContext userContext = default)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.store.ContextRead(ref key, ref input, ref output, userContext, sessionFunctions);
@@ -481,7 +482,7 @@ public Status Read(TKey key, out TOutput output, ref ReadOptions readOptions, TC
         public Status Read(ref TKey key, ref TInput input, ref TOutput output, ref ReadOptions readOptions, out RecordMetadata recordMetadata, TContext userContext = default)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.store.ContextRead(ref key, ref input, ref output, ref readOptions, out recordMetadata, userContext, sessionFunctions);
@@ -497,7 +498,7 @@ public Status Read(ref TKey key, ref TInput input, ref TOutput output, ref ReadO
         public Status ReadAtAddress(long address, ref TInput input, ref TOutput output, ref ReadOptions readOptions, out RecordMetadata recordMetadata, TContext userContext = default)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.store.ContextReadAtAddress(address, ref input, ref output, ref readOptions, out recordMetadata, userContext, sessionFunctions);
@@ -513,7 +514,7 @@ public Status ReadAtAddress(long address, ref TInput input, ref TOutput output,
         public Status ReadAtAddress(long address, ref TKey key, ref TInput input, ref TOutput output, ref ReadOptions readOptions, out RecordMetadata recordMetadata, TContext userContext = default)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.store.ContextReadAtAddress(address, ref key, ref input, ref output, ref readOptions, out recordMetadata, userContext, sessionFunctions);
@@ -557,7 +558,7 @@ public Status Upsert(ref TKey key, ref TInput input, ref TValue desiredValue, re
         private Status Upsert(ref TKey key, long keyHash, ref TInput input, ref TValue desiredValue, ref TOutput output, TContext userContext = default)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.store.ContextUpsert(ref key, keyHash, ref input, ref desiredValue, ref output, userContext, sessionFunctions);
@@ -583,7 +584,7 @@ public Status Upsert(ref TKey key, ref TInput input, ref TValue desiredValue, re
         private Status Upsert(ref TKey key, long keyHash, ref TInput input, ref TValue desiredValue, ref TOutput output, out RecordMetadata recordMetadata, TContext userContext = default)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.store.ContextUpsert(ref key, keyHash, ref input, ref desiredValue, ref output, out recordMetadata, userContext, sessionFunctions);
@@ -639,7 +640,7 @@ public Status RMW(ref TKey key, ref TInput input, ref TOutput output, ref RMWOpt
         private Status RMW(ref TKey key, long keyHash, ref TInput input, ref TOutput output, out RecordMetadata recordMetadata, TContext userContext = default)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.store.ContextRMW(ref key, keyHash, ref input, ref output, out recordMetadata, userContext, sessionFunctions);
@@ -713,7 +714,7 @@ public Status Delete(ref TKey key, ref DeleteOptions deleteOptions, TContext use
         private Status Delete(ref TKey key, long keyHash, TContext userContext = default)
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 return clientSession.store.ContextDelete<TInput, TOutput, TContext, SessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TFunctions, LockableSessionLocker<TKey, TValue, TStoreFunctions, TAllocator>, TStoreFunctions, TAllocator>>(
@@ -749,7 +750,7 @@ internal bool IsModified(TKey key)
         public void Refresh()
         {
             Debug.Assert(!clientSession.store.epoch.ThisInstanceProtected());
-            clientSession.UnsafeResumeThread(sessionFunctions);
+            clientSession.UnsafeResumeThread(sessionFunctions, MaxSpinCount);
             try
             {
                 clientSession.store.InternalRefresh<TInput, TOutput, TContext, SessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TFunctions, LockableSessionLocker<TKey, TValue, TStoreFunctions, TAllocator>, TStoreFunctions, TAllocator>>(sessionFunctions);

libs/storage/Tsavorite/cs/src/core/Epochs/LightEpoch.cs

@@ -195,10 +195,13 @@ public void Suspend()
         /// <summary>
         /// Thread resumes its epoch entry
         /// </summary>
+        /// <param name="maxSpins">
+        /// Maximum number of spins to attempt when acquiring the epoch entry. Default -1 means spin forver.
+        /// </param>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public void Resume()
+        public void Resume(int maxSpins = -1)
         {
-            Acquire();
+            Acquire(maxSpins);
             ProtectAndDrain();
         }
 
@@ -366,10 +369,10 @@ void Drain(long nextEpoch)
         /// Thread acquires its epoch entry
         /// </summary>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        void Acquire()
+        void Acquire(int maxSpins)
         {
             if (Metadata.threadEntryIndex == kInvalidIndex)
-                Metadata.threadEntryIndex = ReserveEntryForThread();
+                Metadata.threadEntryIndex = ReserveEntryForThread(maxSpins);
 
             Debug.Assert((*(tableAligned + Metadata.threadEntryIndex)).localCurrentEpoch == 0,
                 "Trying to acquire protected epoch. Make sure you do not re-enter Tsavorite from callbacks or IDevice implementations. If using tasks, use TaskCreationOptions.RunContinuationsAsynchronously.");
@@ -407,9 +410,10 @@ void Release()
         /// thread will ever have ID 0.
         /// </summary>
         /// <returns>Reserved entry</returns>
-        static int ReserveEntry()
+        static int ReserveEntry(int maxSpins)
         {
-            while (true)
+            // uint means that overflows go to 0 and -1 maxSpin means spin forever; same as while (true)
+            for (uint i = 0; i < maxSpins; i++)
             {
                 // Try to acquire entry
                 if (0 == (threadIndexAligned + Metadata.startOffset1)->threadId)
@@ -433,14 +437,16 @@ static int ReserveEntry()
                     Thread.Yield();
                 }
             }
+
+            throw new TimeoutException($"Could not reserve epoch entry for thread after {maxSpins} spins");
         }
 
         /// <summary>
         /// Allocate a new entry in epoch table. This is called 
         /// once for a thread.
         /// </summary>
         /// <returns>Reserved entry</returns>
-        static int ReserveEntryForThread()
+        static int ReserveEntryForThread(int maxSpins)
         {
             if (Metadata.threadId == 0) // run once per thread for performance
             {
@@ -449,7 +455,7 @@ static int ReserveEntryForThread()
                 Metadata.startOffset1 = (ushort)(1 + (code % kTableSize));
                 Metadata.startOffset2 = (ushort)(1 + ((code >> 16) % kTableSize));
             }
-            return ReserveEntry();
+            return ReserveEntry(maxSpins);
         }
 
         /// <inheritdoc/>