Browse Source

Add: support exception handling from hook

Akeit0 7 months ago
parent
commit
01c919f4b2
2 changed files with 70 additions and 1 deletions
  1. 22 1
      src/Lua/Runtime/LuaVirtualMachine.cs
  2. 48 0
      tests/Lua.Tests/CancellationTest.cs

+ 22 - 1
src/Lua/Runtime/LuaVirtualMachine.cs

@@ -242,10 +242,12 @@ public static partial class LuaVirtualMachine
         public async ValueTask<int> ExecuteClosureAsyncImpl()
         public async ValueTask<int> ExecuteClosureAsyncImpl()
         {
         {
             var returnFrameBase = CurrentReturnFrameBase;
             var returnFrameBase = CurrentReturnFrameBase;
+            var toCatchFlag = false;
             try
             try
             {
             {
                 while (MoveNext(this))
                 while (MoveNext(this))
                 {
                 {
+                    toCatchFlag = true;
                     await Task;
                     await Task;
                     Task = default;
                     Task = default;
                     if (PostOperation is not (PostOperationType.TailCall or PostOperationType.DontPop))
                     if (PostOperation is not (PostOperationType.TailCall or PostOperationType.DontPop))
@@ -258,11 +260,30 @@ public static partial class LuaVirtualMachine
                         break;
                         break;
                     }
                     }
 
 
+                    toCatchFlag = false;
+
                     ThrowIfCancellationRequested();
                     ThrowIfCancellationRequested();
                 }
                 }
 
 
                 return Thread.Stack.Count - returnFrameBase;
                 return Thread.Stack.Count - returnFrameBase;
             }
             }
+            catch (Exception e)
+            {
+                if (toCatchFlag)
+                {
+                    State.CloseUpValues(Thread, FrameBase);
+                    if (e is not (LuaRuntimeException or LuaCanceledException))
+                    {
+                        Exception newException = e is OperationCanceledException ? new LuaCanceledException(Thread, CancellationToken, e) : new LuaRuntimeException(Thread, e);
+                        PopOnTopCallStackFrames();
+                        throw newException;
+                    }
+
+                    PopOnTopCallStackFrames();
+                }
+
+                throw;
+            }
             finally
             finally
             {
             {
                 pool.TryPush(this);
                 pool.TryPush(this);
@@ -823,7 +844,7 @@ public static partial class LuaVirtualMachine
             context.State.CloseUpValues(context.Thread, context.FrameBase);
             context.State.CloseUpValues(context.Thread, context.FrameBase);
             if (e is not (LuaRuntimeException or LuaCanceledException))
             if (e is not (LuaRuntimeException or LuaCanceledException))
             {
             {
-                var newException = new LuaRuntimeException(context.Thread, e);
+                Exception newException = e is OperationCanceledException ? new LuaCanceledException(context.Thread, context.CancellationToken, e) : new LuaRuntimeException(context.Thread, e);
                 context.PopOnTopCallStackFrames();
                 context.PopOnTopCallStackFrames();
                 throw newException;
                 throw newException;
             }
             }

+ 48 - 0
tests/Lua.Tests/CancellationTest.cs

@@ -1,4 +1,5 @@
 using Lua.Standard;
 using Lua.Standard;
+using System.Diagnostics;
 
 
 namespace Lua.Tests;
 namespace Lua.Tests;
 
 
@@ -177,4 +178,51 @@ public class CancellationTest
             }
             }
         }
         }
     }
     }
+    
+    [Test]
+    public async Task CancelByHookTest()
+    {
+        var source = """
+                     local ret = 0
+                     ::loop::
+                     ret = ret + 1
+                     goto loop
+                     return ret
+                     """;
+        var cancellationTokenSource = new CancellationTokenSource();
+        var sw = Stopwatch.StartNew();
+        state.MainThread.SetHook(new LuaFunction("timeout",async (context, cancellationToken) =>
+        {
+            if (sw.ElapsedMilliseconds > 100)
+            {
+                await Task.Delay(1,cancellationToken);
+                cancellationTokenSource.Cancel();
+                cancellationToken.ThrowIfCancellationRequested();
+            }
+            return  context.Return();
+        }),"",10000);
+        cancellationTokenSource.Token.Register(() =>
+        {
+            Console.WriteLine("Cancellation requested");
+        });
+        try
+        {
+            var r = await state.DoStringAsync(source, "@test.lua", cancellationTokenSource.Token);
+            Console.WriteLine(r[0]);
+            Assert.Fail("Expected TaskCanceledException was not thrown.");
+        }
+        catch (Exception e)
+        {
+            Assert.That(e, Is.TypeOf<LuaCanceledException>());
+            Console.WriteLine(e.StackTrace);
+            var luaCancelledException = (LuaCanceledException)e;
+            Assert.That(luaCancelledException.InnerException, Is.TypeOf<OperationCanceledException>());
+            var traceback = luaCancelledException.LuaTraceback;
+            if (traceback != null)
+            {
+                var luaStackTrace = traceback.ToString();
+                Console.WriteLine(luaStackTrace);
+            }
+        }
+    }
 }
 }