// A prototype implementation of the subsystems mechanism for provably safe exception handling
// Bart Jacobs <bart.jacobs@cs.kuleuven.be> and Frank Piessens <frank.piessens@cs.kuleuven.be>, 2008
// http://www.cs.kuleuven.be/~bartj/subsystems
// Last update: 2008-03-20

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.CompilerServices;
using System.Threading;

namespace Subsystems
{
    class ThreadListNode
    {
        public WeakReference/*<PerThreadInfo>*/ perThreadInfo;
        public volatile ThreadListNode next;
    }

    class AbortObject
    {
        public Thread thread;
        public volatile bool signalSent;

        public void Signal()
        {
            thread.Abort(); // I'm assuming this call is allowed in a Constrained Execution Region.
            signalSent = true;
        }

        public void Wait()
        {
            while (!signalSent)
                Thread.Sleep(0); // I'm assuming this call is allowed in a Constrained Execution Region.
            Thread.ResetAbort(); // I'm assuming this call is allowed in a Constrained Execution Region.
            signalSent = false;
        }
    }

    class Compensation
    {
        public IDisposable action;
        public Subsystem subsystem;
        public Compensation next;
    }

    class CompensationStack
    {
        public Compensation entries;
    }

    public class Subsystem
    {
        static volatile int failCount;
        static ThreadListNode threads = new ThreadListNode(); // sentinel
        static Subsystem failureSignal = new Subsystem(null);

        class PerThreadInfo
        {
            public Thread thread;
            public volatile Subsystem subsystem;
            public AbortObject abortObject;
            public CompensationStack compensationStack;
        }

        [ThreadStatic]
        static PerThreadInfo info;

        static PerThreadInfo Info
        {
            get
            {
                PerThreadInfo info = Subsystem.info;
                if (info == null)
                {
                    info = new PerThreadInfo();
                    info.thread = Thread.CurrentThread;
                    info.abortObject = new AbortObject();
                    info.abortObject.thread = info.thread;
                    Subsystem.info = info;
                    lock (threads)
                    {
                        // Clean up entries of dead threads
                        ThreadListNode previous = threads;
                        ThreadListNode node = threads.next;
                        while (node != null)
                        {
                            ThreadListNode next = node.next;
                            if (node.perThreadInfo.Target == null)
                                previous.next = next;
                            previous = node;
                            node = next;
                        }
                        // Add the new entry to the front
                        ThreadListNode newNode = new ThreadListNode();
                        newNode.perThreadInfo = new WeakReference(info);
                        newNode.next = threads.next;
                        threads.next = newNode;
                    }
                }
                return info;
            }
        }

        public static Subsystem Current
        {
            get
            {
                Subsystem s = Info.subsystem;
                if (s == failureSignal)
                    Thread.CurrentThread.Abort();
                return s;
            }
        }

        readonly Subsystem parent;
        volatile Subsystem failRoot; // Ancestor where the exception occurred that caused this subsystem to fail.
        Exception cause; // Only used if this is a fail root.
        volatile int lastFailCountSeen;

        public Subsystem(Subsystem parent)
        {
            this.parent = parent;
            CheckNotFailed();
        }

        public Subsystem() : this(Current) { }

        public Subsystem Parent
        {
            get
            {
                return parent;
            }
        }

        public Subsystem FailRoot
        {
            get
            {
                // Optimized for the case where failures occur rarely.
                if (failRoot == null)
                {
                    int failCountSeen = failCount;
                    if (failCountSeen > lastFailCountSeen)
                    {
                        Subsystem ancestor = this.parent; // We assume we never see an old value (i.e. null for a non-root subsystem) here.
                        while (ancestor != null)
                        {
                            Subsystem r = ancestor.failRoot;
                            if (r != null)
                            {
                                failRoot = r;
                                return r;
                            }
                            ancestor = ancestor.parent;
                        }
                        lastFailCountSeen = failCountSeen;
                    }
                }
                return failRoot;
            }
        }

        public bool HasFailed
        {
            get
            {
                return FailRoot != null;
            }
        }

        public Exception CauseOfFailure
        {
            get
            {
                Subsystem failRoot = FailRoot;
                if (failRoot == null)
                    throw new InvalidOperationException("Subsystem has not failed.");
                return failRoot.cause;
            }
        }

        public void Fail()
        {
            RuntimeHelpers.PrepareConstrainedRegions();
            try { }
            finally
            {
                FailCore();
            }
        }

        // Should be called only from a CER, to ensure that whenever a subsystem is marked as failed, all necessary aborts are performed.
        // This is necessary in order to be able to provide strong liveness guarantees.
        void FailCore() 
        {
            if (failRoot != null)
                return;
            failRoot = this;
#pragma warning disable 0420 // a reference to a volatile field will not be treated as volatile
            Interlocked.Increment(ref failCount);
#pragma warning restore 0420

            Thread currentThread = Thread.CurrentThread;
            ThreadListNode node = threads.next;
            while (node != null)
            {
                PerThreadInfo info = (PerThreadInfo)node.perThreadInfo.Target;
                if (info != null && info.thread != currentThread)
                {
                    Subsystem s = info.subsystem;
                    while (s != this && s != null)
                        s = s.parent;
                    if (s == this)
                    {
                        // Prevent the target thread from transitioning to another subsystem
                        // without catching the abort, and prevent other failing threads from sending
                        // duplicate abort signals.
#pragma warning disable 0420 // a reference to a volatile field will not be treated as volatile
                        Subsystem s2 = Interlocked.CompareExchange<Subsystem>(ref info.subsystem, failureSignal, s);
#pragma warning restore 0420
                        if (s2 == s)
                            info.abortObject.Signal();
                    }
                }
                node = node.next;
            }
        }

        public static void Leave(Action<Func<Exception, bool>> body, Func<Exception, bool> catchBlock)
        {
            PerThreadInfo info = Info;
            Subsystem currentSubsystem = info.subsystem;
#pragma warning disable 0420 // a reference to a volatile field will not be treated as volatile
            Subsystem s2 = Interlocked.CompareExchange<Subsystem>(ref info.subsystem, null, currentSubsystem);
#pragma warning restore 0420
            if (s2 == failureSignal)
                Thread.CurrentThread.Abort();
            if (currentSubsystem != null)
                currentSubsystem.CheckNotFailed();
            body((e) =>
            {
                info.subsystem = currentSubsystem; // Must happen before currentSubsystem.CheckNotFailed() to prevent abort signal loss.
                if (currentSubsystem != null)
                    currentSubsystem.CheckNotFailed(); // Fail fast.
                return catchBlock(e);
            });
            info.subsystem = currentSubsystem; // Must happen before currentSubsystem.CheckNotFailed() to prevent abort signal loss.
            if (currentSubsystem != null)
                currentSubsystem.CheckNotFailed(); // Fail fast.
        }

        void EnterCore(Action body, Func<Exception, bool> catchBlock)
        {
            PerThreadInfo info = Info;
            bool alreadyFailed = false; // Failed state on entry to the cleanup block.
            bool failed = false;
            try
            {
                RuntimeHelpers.ExecuteCodeWithGuaranteedCleanup(userData =>
                {
                    info.subsystem = this;
                    CheckNotFailed();
                    body();
                }, (userData, tryBlockFailed) =>
                {
                    failed = tryBlockFailed;
                    alreadyFailed = failRoot != null;
                    if (tryBlockFailed && !alreadyFailed)
                        FailCore();
                    Subsystem s = info.subsystem;
#pragma warning disable 0420 // a reference to a volatile field will not be treated as volatile
                    Subsystem s3 = Interlocked.CompareExchange<Subsystem>(ref info.subsystem, null, s); // Prevent abort signals after here.
#pragma warning restore 0420
                    if (s3 == failureSignal)
                        info.abortObject.Wait();
                    info.subsystem = null;
                },
                null);
            }
            catch (Exception e)
            {
                // We attempt here to record the exception.
                if (!alreadyFailed) // Does not prevent overwriting another cause in case of concurrent failures, but that's okay.
                {
                    cause = e;
                }
                if (catchBlock == null)
                    throw;
                else if (!failed) // An exception occurred before or after the ExecuteCodeWithGuaranteedCleanup call.
                    throw;
                else
                    if (catchBlock(e))
                        throw;
            }
        }

        public void Enter(Action body)
        {
            EnterCore(body, null);
        }

        public static void TryCatch(Action tryBlock, Action<Exception> catchBlock)
        {
            TryCatch(tryBlock, (e) =>
            {
                catchBlock(e);
                return false;
            });
        }

        public static void TryCatch(Action tryBlock, Func<Exception, bool> catchBlock)
        {
            Leave((catchBlock2) => { new Subsystem().EnterCore(tryBlock, catchBlock2); }, catchBlock);
        }

        public void CheckNotFailed()
        {
            Subsystem failRoot = FailRoot;
            if (failRoot != null)
            {
                Exception cause = failRoot.cause;
                throw new Exception("The program attempted to enter a failed subsystem.\r\nCause of failure: " + (cause == null ? "<could not be recorded>" : cause.Message), cause);
            }
        }

        public static void Fork(Action body)
        {
            Subsystem s = Subsystem.Current;
            if (s != null)
                s.CheckNotFailed();
            new Thread(() =>
            {
                if (s == null)
                    body();
                else
                {
                    try
                    {
                        s.Enter(body);
                    }
                    catch (Exception)
                    {
                    }
                }
            }).Start();
        }

        public static void ResourceScope(Action body)
        {
            PerThreadInfo info = Info;
            CompensationStack currentCompensationStack = info.compensationStack;
            CompensationStack stack = new CompensationStack();
            Exception e = null;
            Compensation currentEntry = null;
            Subsystem currentSubsystem = info.subsystem;
            Leave((catchBlock) =>
            {
                RuntimeHelpers.ExecuteCodeWithGuaranteedCleanup((userData) =>
                {
                    info.compensationStack = stack;
                    try
                    {
                        if (currentSubsystem != null)
                            currentSubsystem.EnterCore(body, null);
                        else
                            body();
                    }
                    catch (Exception ex)
                    {
                        e = ex;
                    }
                    while (stack.entries != null)
                    {
                        currentEntry = stack.entries;
                        stack.entries = currentEntry.next; // Do this now so that things work correctly if the Dispose call pushes new entries onto the stack.
                        currentEntry.subsystem.EnterCore(() =>
                        {
                            currentEntry.action.Dispose();
                        }, (ex) =>
                        {
                            e = ex;
                            return false;
                        });
                    }
                    currentEntry = null;
                    if (e != null)
                        throw new Exception("Failure in resource scope.", e); // TODO: Try to use "throw;" when possible to preserve the exception type.
                }, (userData, tryBlockFailed) =>
                {
                    if (tryBlockFailed)
                    {
                        // This code is executed in the highly unlikely case that a CLR-internal exception occurred above outside of any try block.
                        if (currentEntry != null)
                            currentEntry.subsystem.Fail();
                        while (stack.entries != null)
                        {
                            stack.entries.subsystem.FailCore();
                            stack.entries = stack.entries.next;
                        }
                    }
                    info.compensationStack = currentCompensationStack;
                },
                null);
            }, (ex) => true);
        }

        public static void PushCompensationAction(IDisposable action)
        {
            PerThreadInfo info = Info;
            if (info.compensationStack == null)
                throw new Exception("Not in a resource scope.");
            if (info.subsystem == null)
                throw new Exception("Not in a subsystem.");
            Compensation compensation = new Compensation();
            compensation.next = info.compensationStack.entries;
            compensation.action = action;
            compensation.subsystem = info.subsystem;
            info.compensationStack.entries = compensation;
        }

        public T Allocate<T>(Func<T> body) where T : class, IDisposable
        {
            T action = null;
            Enter(() =>
            {
                action = body();
                PushCompensationAction(action);
            });
            return action;
        }
    }
}