Permissions, RateControl and Auditing

Libraries/Esiur/Core/AsyncAwaiter.cs

@@ -1,52 +1,67 @@
-using System;
-using System.Collections.Generic;
+using System;
 using System.Runtime.CompilerServices;
-using System.Text;
-using System.Threading.Tasks;
+using System.Threading;
 
 namespace Esiur.Core;
 
 public class AsyncAwaiter : INotifyCompletion
 {
-    Action callback = null;
-
-    AsyncException exception = null;
-
-    object result;
+    private static readonly Action CompletedSentinel = () => { };
 
+    private Action continuation;
+    private AsyncException exception;
+    private object result;
+    private readonly AsyncReply reply;
 
     public AsyncAwaiter(AsyncReply reply)
     {
-        reply.Then(x =>
-        {
-            this.IsCompleted = true;
-            this.result = x;
-            this.callback?.Invoke();
-        }).Error(x =>
-        {
-            exception = x;
-            this.IsCompleted = true;
-            this.callback?.Invoke();
-        });
+        this.reply = reply;
+        reply.Then(Complete).Error(Fail);
     }
 
     public object GetResult()
     {
         if (exception != null)
-            throw exception;
+            throw reply.GetExceptionForAwait();
+
         return result;
     }
 
-    public bool IsCompleted { get; private set; }
+    public bool IsCompleted
+        => ReferenceEquals(Volatile.Read(ref continuation), CompletedSentinel);
 
     public void OnCompleted(Action continuation)
     {
-        if (IsCompleted)
-            continuation?.Invoke();
-        else
-            // Continue....
-            callback = continuation;
+        if (continuation == null)
+            throw new ArgumentNullException(nameof(continuation));
+
+        var previous = Interlocked.CompareExchange(ref this.continuation, continuation, null);
+        if (ReferenceEquals(previous, CompletedSentinel))
+        {
+            continuation();
+        }
+        else if (previous != null)
+        {
+            throw new InvalidOperationException("The awaiter already has a continuation.");
+        }
     }
 
+    private void Complete(object value)
+    {
+        result = value;
+        InvokeContinuation();
+    }
 
+    private void Fail(AsyncException value)
+    {
+        exception = value;
+        InvokeContinuation();
+    }
+
+    private void InvokeContinuation()
+    {
+        var registeredContinuation = Interlocked.Exchange(ref continuation, CompletedSentinel);
+        if (registeredContinuation != null && !ReferenceEquals(registeredContinuation, CompletedSentinel))
+            registeredContinuation();
+    }
 }