Rate Policy

Libraries/Esiur/Core/InvocationContext.cs

@@ -1,15 +1,47 @@
-using Esiur.Protocol;
+using Esiur.Data.Types;
+using Esiur.Protocol;
 using System;
+using System.Collections;
 using System.Collections.Generic;
-using System.Text;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Esiur.Core
 {
     public class InvocationContext
     {
-        private uint CallbackId;
+        readonly uint CallbackId;
+        readonly object _stateLock = new object();
+        readonly CancellationTokenSource _cancellation = new CancellationTokenSource();
+        readonly SemaphoreSlim _moveLock = new SemaphoreSlim(1, 1);
+        TaskCompletionSource<bool> _resumeSignal = CompletedSignal();
 
-        internal bool Ended;
+        IAsyncEnumeratorAdapter _asyncEnumerator;
+        IEnumerator _enumerator;
+        bool _halted;
+
+        internal StreamMode StreamMode { get; private set; }
+        internal bool Pausable { get; private set; }
+
+        internal volatile bool Ended;
+
+        /// <summary>
+        /// Gets a token that is cancelled when the caller terminates this execution.
+        /// </summary>
+        public CancellationToken CancellationToken => _cancellation.Token;
+
+        /// <summary>
+        /// Gets whether a pausable execution is currently halted.
+        /// </summary>
+        public bool IsHalted
+        {
+            get
+            {
+                lock (_stateLock)
+                    return _halted;
+            }
+        }
 
         public void Chunk(object value)
         {
@@ -19,8 +51,8 @@ namespace Esiur.Core
             Connection.SendChunk(CallbackId, value);
         }
 
-        public void Progress(uint value, uint max) {
-
+        public void Progress(uint value, uint max)
+        {
             if (Ended)
                 throw new Exception("Execution has ended.");
 
@@ -29,7 +61,6 @@ namespace Esiur.Core
 
         public void Warning(byte level, string message)
         {
-
             if (Ended)
                 throw new Exception("Execution has ended.");
 
@@ -38,12 +69,192 @@ namespace Esiur.Core
 
         public EpConnection Connection { get; internal set; }
 
-
         internal InvocationContext(EpConnection connection, uint callbackId)
         {
             Connection = connection;
             CallbackId = callbackId;
+        }
 
+        internal void InitializeStream(StreamMode streamMode, bool pausable)
+        {
+            StreamMode = streamMode;
+            Pausable = pausable;
+        }
+
+        internal bool SetAsyncEnumerable(object enumerable)
+        {
+            var asyncEnumerableInterface = enumerable.GetType()
+                .GetInterfaces()
+                .Concat(new[] { enumerable.GetType() })
+                .FirstOrDefault(x => x.IsGenericType &&
+                                     x.GetGenericTypeDefinition() == typeof(IAsyncEnumerable<>));
+
+            if (asyncEnumerableInterface == null)
+                return false;
+
+            var itemType = asyncEnumerableInterface.GetGenericArguments()[0];
+            var adapterType = typeof(AsyncEnumeratorAdapter<>).MakeGenericType(itemType);
+            _asyncEnumerator = (IAsyncEnumeratorAdapter)Activator.CreateInstance(
+                adapterType, enumerable, _cancellation.Token);
+            return true;
+        }
+
+        internal void SetEnumerable(IEnumerable enumerable)
+        {
+            _enumerator = enumerable.GetEnumerator();
+        }
+
+        internal async Task<(bool HasValue, object Value)> PullAsync()
+        {
+            if (StreamMode != StreamMode.Pull || _asyncEnumerator == null)
+                throw new InvalidOperationException("Execution is not an asynchronous pull stream.");
+
+            await _moveLock.WaitAsync(_cancellation.Token);
+            try
+            {
+                await WaitWhileHaltedAsync();
+                _cancellation.Token.ThrowIfCancellationRequested();
+
+                var hasValue = await _asyncEnumerator.MoveNextAsync();
+                return (hasValue, hasValue ? _asyncEnumerator.Current : null);
+            }
+            finally
+            {
+                _moveLock.Release();
+            }
+        }
+
+        internal async Task<(bool HasValue, object Value)> MoveNextAsync()
+        {
+            if (_enumerator == null)
+                throw new InvalidOperationException("Execution is not a synchronous stream.");
+
+            await WaitWhileHaltedAsync();
+            _cancellation.Token.ThrowIfCancellationRequested();
+
+            var hasValue = _enumerator.MoveNext();
+            return (hasValue, hasValue ? _enumerator.Current : null);
+        }
+
+        internal bool Halt()
+        {
+            lock (_stateLock)
+            {
+                if (Ended || !Pausable || _halted)
+                    return false;
+
+                _halted = true;
+                _resumeSignal = NewSignal();
+                return true;
+            }
+        }
+
+        internal bool Resume()
+        {
+            TaskCompletionSource<bool> resumeSignal;
+
+            lock (_stateLock)
+            {
+                if (Ended || !Pausable || !_halted)
+                    return false;
+
+                _halted = false;
+                resumeSignal = _resumeSignal;
+            }
+
+            resumeSignal.TrySetResult(true);
+            return true;
+        }
+
+        /// <summary>
+        /// Asynchronously waits until a halted execution is resumed.
+        /// </summary>
+        public async Task WaitWhileHaltedAsync()
+        {
+            Task resumeTask;
+
+            lock (_stateLock)
+                resumeTask = _resumeSignal.Task;
+
+            await resumeTask;
+            _cancellation.Token.ThrowIfCancellationRequested();
+        }
+
+        internal async Task TerminateAsync()
+        {
+            TaskCompletionSource<bool> resumeSignal;
+
+            lock (_stateLock)
+            {
+                if (Ended)
+                    return;
+
+                Ended = true;
+                _halted = false;
+                resumeSignal = _resumeSignal;
+            }
+
+            resumeSignal.TrySetResult(true);
+            _cancellation.Cancel();
+
+            if (_asyncEnumerator != null)
+                await _asyncEnumerator.DisposeAsync();
+
+            (_enumerator as IDisposable)?.Dispose();
+        }
+
+        internal async Task EndAsync()
+        {
+            TaskCompletionSource<bool> resumeSignal;
+
+            lock (_stateLock)
+            {
+                if (Ended)
+                    return;
+
+                Ended = true;
+                _halted = false;
+                resumeSignal = _resumeSignal;
+            }
+
+            resumeSignal.TrySetResult(true);
+
+            if (_asyncEnumerator != null)
+                await _asyncEnumerator.DisposeAsync();
+
+            (_enumerator as IDisposable)?.Dispose();
+        }
+
+        static TaskCompletionSource<bool> NewSignal()
+            => new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        static TaskCompletionSource<bool> CompletedSignal()
+        {
+            var signal = NewSignal();
+            signal.SetResult(true);
+            return signal;
         }
     }
+
+    interface IAsyncEnumeratorAdapter
+    {
+        object Current { get; }
+        ValueTask<bool> MoveNextAsync();
+        ValueTask DisposeAsync();
+    }
+
+    sealed class AsyncEnumeratorAdapter<T> : IAsyncEnumeratorAdapter
+    {
+        readonly IAsyncEnumerator<T> _enumerator;
+
+        public object Current => _enumerator.Current;
+
+        public AsyncEnumeratorAdapter(IAsyncEnumerable<T> enumerable, CancellationToken cancellationToken)
+        {
+            _enumerator = enumerable.GetAsyncEnumerator(cancellationToken);
+        }
+
+        public ValueTask<bool> MoveNextAsync() => _enumerator.MoveNextAsync();
+        public ValueTask DisposeAsync() => _enumerator.DisposeAsync();
+    }
 }