Kowalski/CoTask/Core/CoTask.cs

using System;
using System.Collections;
using System.Collections.Generic;
namespace Kowalski
{
public enum CoTaskStatus
{
// The task is running but has not yet completed
Running,
// The task completed execution successfully
Succeeded,
// The task was canceled
Canceled,
}
public sealed partial class CoTask
{
public static readonly object Canceled = new object();
private CoTaskStatus m_Status;
private IEnumerator m_Routine;
private readonly long m_Id;
private readonly CoTaskAgent m_Agent;
private readonly Action<Action<CoTask>> m_CancelingAction;
private readonly Action<CoTask> m_CancelingMethod;
private Stack<IEnumerator> m_Jobs;
private int m_NextUpdateFrame;
private bool m_IsUpdating;
public long Id => m_Id;
public CoTaskAgent Agent => m_Agent;
public bool IsCompleted
{
get
{
return m_Status == CoTaskStatus.Succeeded ||
m_Status == CoTaskStatus.Canceled;
}
}
internal CoTaskStatus Status => m_Status;
internal CoTask(IEnumerator routine, long id, CoTaskAgent agent, Action<Action<CoTask>> cancelingAction)
{
m_Status = CoTaskStatus.Running;
m_Routine = routine;
m_Id = id;
m_Agent = agent;
m_CancelingAction = cancelingAction;
m_CancelingMethod = CancelTask;
m_Jobs = new Stack<IEnumerator>();
m_Jobs.Push(routine);
m_NextUpdateFrame = agent.TaskUpdater.FrameNumber;
m_IsUpdating = false;
}
internal bool IsWrapping(IEnumerator routine)
{
return ReferenceEquals(m_Routine, routine);
}
internal void StartTask()
{
InitRoutine(m_Jobs.Peek());
UpdateTask();
}
internal void UpdateTask()
{
if (m_IsUpdating)
{
throw new InvalidOperationException("Task is updating");
}
int currentFrame = m_Agent.TaskUpdater.FrameNumber;
// Frame number can overflow so we do not use < or >
if (currentFrame != m_NextUpdateFrame || IsCompleted)
{
return;
}
while (m_Jobs.Count > 0)
{
var job = m_Jobs.Peek();
bool hasValue;
try
{
m_IsUpdating = true;
hasValue = job.MoveNext();
// Process pending canceling
m_CancelingAction(m_CancelingMethod);
}
catch
{
SetCanceled();
throw;
}
finally
{
m_IsUpdating = false;
}
if (IsCompleted)
{
return;
}
if (hasValue)
{
// Assertion: job.Current should not throw exception
var val = job.Current;
if (val is null)
{
ScheduleForNextFrame();
return;
}
else if (ReferenceEquals(val, Canceled))
{
SetCanceled();
return;
}
else if (val is IEnumerator nested)
{
InitRoutine(nested);
m_Jobs.Push(nested);
}
else
{
throw new NotSupportedException($"Yield value is not supported: {val.GetType()}");
}
}
else
{
ReleaseRoutine(m_Jobs.Pop());
}
}
m_Status = CoTaskStatus.Succeeded;
}
internal bool TryCancel()
{
if (IsCompleted)
{
return true;
}
if (m_IsUpdating)
{
// Should be canceled after updating
return false;
}
ForceCancel();
return true;
}
private static void CancelTask(CoTask task)
{
task.SetCanceled();
}
private void SetCanceled()
{
if (!IsCompleted)
{
ForceCancel();
}
}
private void ForceCancel()
{
// Mark as canceled first to avoid canceling multiple times (inside ReleaseTask)
m_Status = CoTaskStatus.Canceled;
ReleaseTask();
}
private void InitRoutine(IEnumerator routine)
{
if (routine is ICoTaskEnumerator t)
{
t.InitCoTaskEnumerator(this);
}
}
private void ReleaseRoutine(IEnumerator routine)
{
if (routine is ICoTaskEnumerator t)
{
t.ReleaseCoTaskEnumerator();
}
}
private void ReleaseTask()
{
// Throwing exception in finally block is bad
// This is the outermost exception
Exception bad = null;
while (m_Jobs.TryPop(out var job))
{
try
{
(job as IDisposable)?.Dispose();
ReleaseRoutine(job);
}
catch (System.Exception e)
{
bad = e;
}
}
if (!(bad is null))
{
throw bad;
}
}
private void ScheduleForNextFrame()
{
m_NextUpdateFrame = m_Agent.TaskUpdater.FrameNumber + 1;
}
}
}