Custom Coroutines in Unity3D.

I had a whim today. In my Python green threading framework (Fibra), it is possible to run a section of code inside the green thread function in another OS level thread. This lets you do fun things like break out of the main thread when waiting for IO, or doing something that might block the rest of the green threads.

In Unity3D-speak, a green thread is called a Coroutine. After doing some rather ugly things with threads today, I figured things could be much cleaner if I could simply execute certain parts of a Coroutine outside the Unity3D main loop. It is possible, and I've done it. Here is the code! :-)

This is the test MonoBehaviour script. Notice that it inherits from Fibra, and uses StartGreenThread to execute methods that look like real Unity3D Coroutines. Inside the Coroutine, a FiBranch instance is yielded, which is the signal for the scheduler to run the next section of the Coroutine in the thread pool.

using UnityEngine;
using System.Collections;
using System.Threading;

public class FibraTest : Fibra {
    
    public IEnumerator Busy() {
        Debug.Log("Before Continue.");
        yield return new FiContinue();    
        Debug.Log("After Continue.");
        yield return new FiContinue();
        Debug.Log("Finished.");
        yield return new FiSleep(3);
        Debug.Log("Ok really finished. Well with this CPU anyway:" + Thread.CurrentThread.ManagedThreadId.ToString());
        yield return new FiBranch();
        Debug.Log("I'm in another thread:" + Thread.CurrentThread.ManagedThreadId.ToString());
        Thread.Sleep(1000);
        yield return new FiContinue();    
        Debug.Log("Done:" + Thread.CurrentThread.ManagedThreadId.ToString());
        
    }
    
    // Use this for initialization
    void Start () {
        StartGreenThread("Busy");
    }
    
    // Update is called once per frame
    public new void Update () {
        base.Update();
    }
}

This is the Fibra class, which implements the green threading magic.

using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using System.Threading;


public abstract class FInstruction
{
    public abstract bool Execute (Fibra scheduler, IEnumerator task);
}

public class FiContinue : FInstruction
{
    public override bool Execute (Fibra scheduler, IEnumerator task)
    {
        return true;
    }
}

public class FiBranch : FInstruction
{
    public override bool Execute (Fibra scheduler, IEnumerator task)
    {
        ThreadPool.QueueUserWorkItem(delegate(object state) {
            bool removeTask = false;
            try {
                if (task.MoveNext ()) {
                    var fi = task.Current as FInstruction;
                    removeTask = !fi.Execute (scheduler, task);
                } else
                    removeTask = true;
            } catch(Exception e) {
            // Do something in here!    
            }
            if(!removeTask) {
                scheduler.StartGreenThread(task);    
            }
        });
        return false;
    }
}

public class FiSleep : FInstruction
{
    float seconds;
    public FiSleep (float seconds)
    {
        this.seconds = seconds;
    }

    IEnumerator Sleep (Fibra scheduler, IEnumerator task)
    {
        yield return new WaitForSeconds (seconds);
        scheduler.StartGreenThread (task);
    }

    public override bool Execute (Fibra scheduler, IEnumerator task)
    {
        scheduler.StartCoroutine (Sleep (scheduler, task));
        return false;
    }
}

public class Fibra : MonoBehaviour
{
    static List<IEnumerator> tasks = new List<IEnumerator> ();
    static Dictionary<string, IEnumerator> names = new Dictionary<string, IEnumerator> ();


    public void StartGreenThread (string taskName, params object[] args)
    {
        var method = this.GetType ().GetMethod (taskName);
        if (args.Length == 0)
            args = null;
        var task = method.Invoke (this, args) as IEnumerator;
        names.Add (taskName, task);
        StartGreenThread (task);
    }

    public void StopGreenThread (string taskName)
    {
        var task = names[taskName];
        Fibra.tasks.Remove (task);
        names.Remove (taskName);
    }

    public void StartGreenThread (IEnumerator task)
    {
        lock(Fibra.tasks) {
            Fibra.tasks.Add (task);
        }
    }

    public void Update ()
    {
        for (int i = 0; i < Fibra.tasks.Count; i++) {
            bool removeTask = false;
            var task = tasks[i];
            if (task.MoveNext ()) {
                var fi = task.Current as FInstruction;
                removeTask = !fi.Execute (this, task);
            } else
                removeTask = true;
            if (removeTask)
                tasks.RemoveAt (i--);
        }
    }

}

The above class uses a double dispatch mechanism, which makes it a piece of cake to add your own custom schedule instructions. Simply inherit from FInstruction, and write your own Execute method. Return true if you want the green thread to be rescheduled, or false if you don't.