using System;
|
using System.Collections.Generic;
|
using System.Threading;
|
|
namespace Cysharp.Threading.Tasks
|
{
|
public static class Channel
|
{
|
public static Channel<T> CreateSingleConsumerUnbounded<T>()
|
{
|
return new SingleConsumerUnboundedChannel<T>();
|
}
|
}
|
|
public abstract class Channel<TWrite, TRead>
|
{
|
public ChannelReader<TRead> Reader { get; protected set; }
|
public ChannelWriter<TWrite> Writer { get; protected set; }
|
|
public static implicit operator ChannelReader<TRead>(Channel<TWrite, TRead> channel) => channel.Reader;
|
public static implicit operator ChannelWriter<TWrite>(Channel<TWrite, TRead> channel) => channel.Writer;
|
}
|
|
public abstract class Channel<T> : Channel<T, T>
|
{
|
}
|
|
public abstract class ChannelReader<T>
|
{
|
public abstract bool TryRead(out T item);
|
public abstract UniTask<bool> WaitToReadAsync(CancellationToken cancellationToken = default(CancellationToken));
|
|
public abstract UniTask Completion { get; }
|
|
public virtual UniTask<T> ReadAsync(CancellationToken cancellationToken = default(CancellationToken))
|
{
|
if (this.TryRead(out var item))
|
{
|
return UniTask.FromResult(item);
|
}
|
|
return ReadAsyncCore(cancellationToken);
|
}
|
|
async UniTask<T> ReadAsyncCore(CancellationToken cancellationToken = default(CancellationToken))
|
{
|
if (await WaitToReadAsync(cancellationToken))
|
{
|
if (TryRead(out var item))
|
{
|
return item;
|
}
|
}
|
|
throw new ChannelClosedException();
|
}
|
|
public abstract IUniTaskAsyncEnumerable<T> ReadAllAsync(CancellationToken cancellationToken = default(CancellationToken));
|
}
|
|
public abstract class ChannelWriter<T>
|
{
|
public abstract bool TryWrite(T item);
|
public abstract bool TryComplete(Exception error = null);
|
|
public void Complete(Exception error = null)
|
{
|
if (!TryComplete(error))
|
{
|
throw new ChannelClosedException();
|
}
|
}
|
}
|
|
public partial class ChannelClosedException : InvalidOperationException
|
{
|
public ChannelClosedException() :
|
base("Channel is already closed.")
|
{ }
|
|
public ChannelClosedException(string message) : base(message) { }
|
|
public ChannelClosedException(Exception innerException) :
|
base("Channel is already closed", innerException)
|
{ }
|
|
public ChannelClosedException(string message, Exception innerException) : base(message, innerException) { }
|
}
|
|
internal class SingleConsumerUnboundedChannel<T> : Channel<T>
|
{
|
readonly Queue<T> items;
|
readonly SingleConsumerUnboundedChannelReader readerSource;
|
UniTaskCompletionSource completedTaskSource;
|
UniTask completedTask;
|
|
Exception completionError;
|
bool closed;
|
|
public SingleConsumerUnboundedChannel()
|
{
|
items = new Queue<T>();
|
Writer = new SingleConsumerUnboundedChannelWriter(this);
|
readerSource = new SingleConsumerUnboundedChannelReader(this);
|
Reader = readerSource;
|
}
|
|
sealed class SingleConsumerUnboundedChannelWriter : ChannelWriter<T>
|
{
|
readonly SingleConsumerUnboundedChannel<T> parent;
|
|
public SingleConsumerUnboundedChannelWriter(SingleConsumerUnboundedChannel<T> parent)
|
{
|
this.parent = parent;
|
}
|
|
public override bool TryWrite(T item)
|
{
|
bool waiting;
|
lock (parent.items)
|
{
|
if (parent.closed) return false;
|
|
parent.items.Enqueue(item);
|
waiting = parent.readerSource.isWaiting;
|
}
|
|
if (waiting)
|
{
|
parent.readerSource.SingalContinuation();
|
}
|
|
return true;
|
}
|
|
public override bool TryComplete(Exception error = null)
|
{
|
bool waiting;
|
lock (parent.items)
|
{
|
if (parent.closed) return false;
|
parent.closed = true;
|
waiting = parent.readerSource.isWaiting;
|
|
if (parent.items.Count == 0)
|
{
|
if (error == null)
|
{
|
if (parent.completedTaskSource != null)
|
{
|
parent.completedTaskSource.TrySetResult();
|
}
|
else
|
{
|
parent.completedTask = UniTask.CompletedTask;
|
}
|
}
|
else
|
{
|
if (parent.completedTaskSource != null)
|
{
|
parent.completedTaskSource.TrySetException(error);
|
}
|
else
|
{
|
parent.completedTask = UniTask.FromException(error);
|
}
|
}
|
|
if (waiting)
|
{
|
parent.readerSource.SingalCompleted(error);
|
}
|
}
|
|
parent.completionError = error;
|
}
|
|
return true;
|
}
|
}
|
|
sealed class SingleConsumerUnboundedChannelReader : ChannelReader<T>, IUniTaskSource<bool>
|
{
|
readonly Action<object> CancellationCallbackDelegate = CancellationCallback;
|
readonly SingleConsumerUnboundedChannel<T> parent;
|
|
CancellationToken cancellationToken;
|
CancellationTokenRegistration cancellationTokenRegistration;
|
UniTaskCompletionSourceCore<bool> core;
|
internal bool isWaiting;
|
|
public SingleConsumerUnboundedChannelReader(SingleConsumerUnboundedChannel<T> parent)
|
{
|
this.parent = parent;
|
|
TaskTracker.TrackActiveTask(this, 4);
|
}
|
|
public override UniTask Completion
|
{
|
get
|
{
|
if (parent.completedTaskSource != null) return parent.completedTaskSource.Task;
|
|
if (parent.closed)
|
{
|
return parent.completedTask;
|
}
|
|
parent.completedTaskSource = new UniTaskCompletionSource();
|
return parent.completedTaskSource.Task;
|
}
|
}
|
|
public override bool TryRead(out T item)
|
{
|
lock (parent.items)
|
{
|
if (parent.items.Count != 0)
|
{
|
item = parent.items.Dequeue();
|
|
// complete when all value was consumed.
|
if (parent.closed && parent.items.Count == 0)
|
{
|
if (parent.completionError != null)
|
{
|
if (parent.completedTaskSource != null)
|
{
|
parent.completedTaskSource.TrySetException(parent.completionError);
|
}
|
else
|
{
|
parent.completedTask = UniTask.FromException(parent.completionError);
|
}
|
}
|
else
|
{
|
if (parent.completedTaskSource != null)
|
{
|
parent.completedTaskSource.TrySetResult();
|
}
|
else
|
{
|
parent.completedTask = UniTask.CompletedTask;
|
}
|
}
|
}
|
}
|
else
|
{
|
item = default;
|
return false;
|
}
|
}
|
|
return true;
|
}
|
|
public override UniTask<bool> WaitToReadAsync(CancellationToken cancellationToken)
|
{
|
if (cancellationToken.IsCancellationRequested)
|
{
|
return UniTask.FromCanceled<bool>(cancellationToken);
|
}
|
|
lock (parent.items)
|
{
|
if (parent.items.Count != 0)
|
{
|
return CompletedTasks.True;
|
}
|
|
if (parent.closed)
|
{
|
if (parent.completionError == null)
|
{
|
return CompletedTasks.False;
|
}
|
else
|
{
|
return UniTask.FromException<bool>(parent.completionError);
|
}
|
}
|
|
cancellationTokenRegistration.Dispose();
|
|
core.Reset();
|
isWaiting = true;
|
|
this.cancellationToken = cancellationToken;
|
if (this.cancellationToken.CanBeCanceled)
|
{
|
cancellationTokenRegistration = this.cancellationToken.RegisterWithoutCaptureExecutionContext(CancellationCallbackDelegate, this);
|
}
|
|
return new UniTask<bool>(this, core.Version);
|
}
|
}
|
|
public void SingalContinuation()
|
{
|
core.TrySetResult(true);
|
}
|
|
public void SingalCancellation(CancellationToken cancellationToken)
|
{
|
TaskTracker.RemoveTracking(this);
|
core.TrySetCanceled(cancellationToken);
|
}
|
|
public void SingalCompleted(Exception error)
|
{
|
if (error != null)
|
{
|
TaskTracker.RemoveTracking(this);
|
core.TrySetException(error);
|
}
|
else
|
{
|
TaskTracker.RemoveTracking(this);
|
core.TrySetResult(false);
|
}
|
}
|
|
public override IUniTaskAsyncEnumerable<T> ReadAllAsync(CancellationToken cancellationToken = default)
|
{
|
return new ReadAllAsyncEnumerable(this, cancellationToken);
|
}
|
|
bool IUniTaskSource<bool>.GetResult(short token)
|
{
|
return core.GetResult(token);
|
}
|
|
void IUniTaskSource.GetResult(short token)
|
{
|
core.GetResult(token);
|
}
|
|
UniTaskStatus IUniTaskSource.GetStatus(short token)
|
{
|
return core.GetStatus(token);
|
}
|
|
void IUniTaskSource.OnCompleted(Action<object> continuation, object state, short token)
|
{
|
core.OnCompleted(continuation, state, token);
|
}
|
|
UniTaskStatus IUniTaskSource.UnsafeGetStatus()
|
{
|
return core.UnsafeGetStatus();
|
}
|
|
static void CancellationCallback(object state)
|
{
|
var self = (SingleConsumerUnboundedChannelReader)state;
|
self.SingalCancellation(self.cancellationToken);
|
}
|
|
sealed class ReadAllAsyncEnumerable : IUniTaskAsyncEnumerable<T>, IUniTaskAsyncEnumerator<T>
|
{
|
readonly Action<object> CancellationCallback1Delegate = CancellationCallback1;
|
readonly Action<object> CancellationCallback2Delegate = CancellationCallback2;
|
|
readonly SingleConsumerUnboundedChannelReader parent;
|
CancellationToken cancellationToken1;
|
CancellationToken cancellationToken2;
|
CancellationTokenRegistration cancellationTokenRegistration1;
|
CancellationTokenRegistration cancellationTokenRegistration2;
|
|
T current;
|
bool cacheValue;
|
bool running;
|
|
public ReadAllAsyncEnumerable(SingleConsumerUnboundedChannelReader parent, CancellationToken cancellationToken)
|
{
|
this.parent = parent;
|
this.cancellationToken1 = cancellationToken;
|
}
|
|
public IUniTaskAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
|
{
|
if (running)
|
{
|
throw new InvalidOperationException("Enumerator is already running, does not allow call GetAsyncEnumerator twice.");
|
}
|
|
if (this.cancellationToken1 != cancellationToken)
|
{
|
this.cancellationToken2 = cancellationToken;
|
}
|
|
if (this.cancellationToken1.CanBeCanceled)
|
{
|
this.cancellationTokenRegistration1 = this.cancellationToken1.RegisterWithoutCaptureExecutionContext(CancellationCallback1Delegate, this);
|
}
|
|
if (this.cancellationToken2.CanBeCanceled)
|
{
|
this.cancellationTokenRegistration2 = this.cancellationToken2.RegisterWithoutCaptureExecutionContext(CancellationCallback2Delegate, this);
|
}
|
|
running = true;
|
return this;
|
}
|
|
public T Current
|
{
|
get
|
{
|
if (cacheValue)
|
{
|
return current;
|
}
|
parent.TryRead(out current);
|
return current;
|
}
|
}
|
|
public UniTask<bool> MoveNextAsync()
|
{
|
cacheValue = false;
|
return parent.WaitToReadAsync(CancellationToken.None); // ok to use None, registered in ctor.
|
}
|
|
public UniTask DisposeAsync()
|
{
|
cancellationTokenRegistration1.Dispose();
|
cancellationTokenRegistration2.Dispose();
|
return default;
|
}
|
|
static void CancellationCallback1(object state)
|
{
|
var self = (ReadAllAsyncEnumerable)state;
|
self.parent.SingalCancellation(self.cancellationToken1);
|
}
|
|
static void CancellationCallback2(object state)
|
{
|
var self = (ReadAllAsyncEnumerable)state;
|
self.parent.SingalCancellation(self.cancellationToken2);
|
}
|
}
|
}
|
}
|
}
|
