Restore count hooks after aborted callbacks

README.md

@@ -313,6 +313,18 @@ This code can resume the execution of the Lua script after waiting with await, a
 
 ![img](docs/img1.png)
 
+## Debug Hooks
+
+`state.SetHook(hook, mask, count)` installs a hook for that Lua thread. The mask combines `c` (call), `r` (return), and `l` (line); a positive `count` enables instruction events even with an empty mask. The callback receives the event name and a line number (`nil` for non-line events). Return values, including `false`, are ignored. `state.SetHook(null, "")` removes the hook.
+
+A hook can call `SetHook` to replace or remove itself or change its interval. Setting a hook resets the instruction counter, so subsequent count events use the new interval. Hooks are suppressed while the callback executes, including Lua code it invokes.
+
+Count events measure Lua instructions, as in [Lua 5.2](https://www.lua.org/manual/5.2/manual.html#lua_sethook). They do not fire during a C# host function's own execution or while awaiting it. If that host function calls back into Lua, the nested Lua instructions can trigger hooks. A timeout checked by a hook therefore cannot interrupt a blocking C# function. Long-running host functions should check the supplied `CancellationToken` or pass it to cancellable operations. A host-caught exception or cancellation from a hook does not prevent hooks from running on later executions using the same state.
+
+Hook-based memory sampling is approximate: one Lua instruction or host operation can allocate substantially before the next event. `GC.GetTotalMemory(false)` reports the managed heap, not memory attributable to one Lua state, so it cannot enforce a strict per-state memory limit.
+
+A `LuaRuntimeException` thrown by a hook is a Lua error that `pcall` can catch; it is not an uncatchable execution limit.
+
 ## Coroutines
 
 Lua coroutines are represented by the `LuaState` type.

README_JA.md

@@ -310,6 +310,18 @@ print "goodbye!"
 
 ![img](docs/img1.png)
 
+## デバッグフック
+
+`state.SetHook(hook, mask, count)`で、そのLuaスレッドにフックを設定できます。`mask`には`c`（呼び出し）、`r`（戻り）、`l`（行）を組み合わせます。正の`count`を指定すると、空のmaskでも命令数に応じたイベントが有効になります。コールバックにはイベント名と行番号（行イベント以外では`nil`）が渡されます。`false`を含む戻り値は無視されます。`state.SetHook(null, "")`でフックを解除できます。
+
+フック内で`SetHook`を呼び、自身を置き換えたり解除したり、間隔を変更したりできます。設定時に命令カウンターがリセットされ、その後のcountイベントは新しい間隔を使います。コールバックの実行中は、そこから呼び出したLuaコードも含め、フックは再帰的に発火しません。
+
+countイベントが数えるのは[Lua 5.2](https://www.lua.org/manual/5.2/manual.html#lua_sethook)と同様にLuaの命令です。C#ホスト関数自体の実行中や、そのawait中には発火しません。ホスト関数からLuaを呼び出した場合は、そのLua命令によって発火できます。そのため、フックで確認するタイムアウトでは、実行中のC#関数を強制中断できません。長時間動くホスト関数では、渡された`CancellationToken`を確認するか、キャンセル可能な処理に渡してください。フックからの例外やキャンセルをホスト側で捕捉した後も、同じstateの次の実行ではフックが動作します。
+
+フックによるメモリ使用量の確認は近似的なものです。次のイベントまでに、1つのLua命令やホスト処理が大量のメモリを確保する場合があります。また、`GC.GetTotalMemory(false)`は特定のLua stateの使用量ではなくマネージドヒープの値なので、stateごとの厳密なメモリ上限には使えません。
+
+フックが投げた`LuaRuntimeException`はLuaのエラーとして`pcall`で捕捉できます。そのため、捕捉できない強制的な実行上限にはなりません。
+
 ## コルーチン
 
 Luaのコルーチンは`LuaState`型で表現されます。

src/Lua/LuaState.cs

@@ -416,6 +416,17 @@ public LuaTable GetCurrentEnvironment()
         return Environment;
     }
 
+    /// <summary>Sets the debug hook for this Lua thread, or removes it when <paramref name="hook"/> is null.</summary>
+    /// <param name="hook">The callback receiving an event name and a line number (nil for non-line events). Its return values are ignored.</param>
+    /// <param name="mask">A combination of 'c' (call), 'r' (return), and 'l' (line).</param>
+    /// <param name="count">A positive Lua instruction interval, independent of the mask. Zero disables count events.</param>
+    /// <remarks>
+    /// The callback can call SetHook to replace or remove itself or change the interval.
+    /// Setting the hook resets the count interval; callbacks do not recursively trigger hooks.
+    /// Count events track Lua instructions, not C# execution or time spent awaiting host functions.
+    /// A host function that calls back into Lua can trigger count events for that Lua execution.
+    /// Long-running host functions must cooperate with the execution's CancellationToken to stop promptly.
+    /// </remarks>
     public void SetHook(LuaFunction? hook, string mask, int count = 0)
     {
         if (hook is null)

src/Lua/Runtime/LuaVirtualMachine.Debug.cs

@@ -40,16 +40,24 @@ static async ValueTask<int> Impl(VirtualMachineExecutionContext context)
                 stack.Push("count");
                 stack.Push(LuaValue.Nil);
                 context.State.IsInHook = true;
-                var frame = context.State.CreateCallStackFrame(hook, 2, top, pc);
-                context.State.PushCallStackFrame(frame);
-                LuaFunctionExecutionContext funcContext = new()
+                try
                 {
-                    State = context.State,
-                    ArgumentCount = stack.Count - frame.Base,
-                    ReturnFrameBase = frame.ReturnBase,
-                };
-                await hook.Func(funcContext, context.CancellationToken);
-                context.State.IsInHook = false;
+                    var frame = context.State.CreateCallStackFrame(hook, 2, top, pc);
+                    context.State.PushCallStackFrame(frame);
+                    LuaFunctionExecutionContext funcContext = new()
+                    {
+                        State = context.State,
+                        ArgumentCount = stack.Count - frame.Base,
+                        ReturnFrameBase = frame.ReturnBase,
+                    };
+                    await hook.Func(funcContext, context.CancellationToken);
+                }
+                finally
+                {
+                    // An aborted hook must not suppress hooks when this state is reused.
+                    // Frame creation can also fail before the callback is invoked.
+                    context.State.IsInHook = false;
+                }
 
                 countHookIsDone = true;
             }

tests/Lua.Tests/HookTests.cs

@@ -0,0 +1,307 @@
+namespace Lua.Tests;
+
+public class HookTests
+{
+    const string Loop = "local x=0 for i=1,100 do x=x+i end return x";
+
+    [TestCase(false, "runtime")]
+    [TestCase(true, "runtime")]
+    [TestCase(false, "host")]
+    [TestCase(true, "host")]
+    [TestCase(false, "cancel")]
+    [TestCase(true, "cancel")]
+    public async Task CountHook_AfterException_StateCanRunHooksAgain(
+        bool asynchronous,
+        string failure
+    )
+    {
+        using var state = LuaState.Create();
+        using var cancellation = new CancellationTokenSource();
+        var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        state.SetHook(
+            new("abort", (c, ct) => asynchronous ? ThrowAsync(c, ct) : Throw(c, ct)),
+            "",
+            4
+        );
+
+        var execution = state.DoStringAsync(Loop, "hook.lua", cancellation.Token).AsTask();
+        if (asynchronous)
+        {
+            Assert.That(execution.IsCompleted, Is.False);
+            gate.SetResult();
+        }
+
+        if (failure == "cancel")
+        {
+            Assert.ThrowsAsync<LuaCanceledException>(async () => await execution);
+        }
+        else
+        {
+            Assert.ThrowsAsync<LuaRuntimeException>(async () => await execution);
+        }
+
+        Assert.That(state.CallStackFrameCount, Is.Zero);
+        var stackTop = state.Stack.Count;
+
+        state.Environment["echo"] = new LuaFunction((c, ct) => new(c.Return(c.GetArgument(0))));
+        for (var attempt = 0; attempt < 2; attempt++)
+        {
+            if (attempt == 1)
+            {
+                state.SetHook(null, "");
+            }
+
+            var events = new HashSet<string>();
+            state.SetHook(
+                new(
+                    "reuse",
+                    (c, ct) =>
+                    {
+                        events.Add(c.GetArgument<string>(0));
+                        return new(c.Return());
+                    }
+                ),
+                "lcr",
+                4
+            );
+
+            var result = await state.DoStringAsync(
+                "local x=0 for i=1,100 do x=x+i end echo(x) return x"
+            );
+            Assert.That(result[0], Is.EqualTo((LuaValue)5050));
+            Assert.That(events, Is.SupersetOf(new[] { "count", "line", "call", "return" }));
+            Assert.That(state.CallStackFrameCount, Is.Zero);
+            Assert.That(state.Stack.Count, Is.EqualTo(stackTop));
+        }
+
+        ValueTask<int> Throw(LuaFunctionExecutionContext context, CancellationToken ct)
+        {
+            if (failure == "cancel")
+            {
+                cancellation.Cancel();
+                ct.ThrowIfCancellationRequested();
+            }
+
+            throw failure switch
+            {
+                "host" => new InvalidOperationException("hook failed"),
+                _ => new LuaRuntimeException(context.State, "hook failed"),
+            };
+        }
+
+        async ValueTask<int> ThrowAsync(LuaFunctionExecutionContext context, CancellationToken ct)
+        {
+            await gate.Task;
+            return await Throw(context, ct);
+        }
+    }
+
+    [Test]
+    public async Task CountHook_WhenFrameLimitIsExceeded_StateCanRunHooksAgain()
+    {
+        using var state = LuaState.Create();
+        state.MaxCallDepth = 1;
+        state.SetHook(new((c, ct) => new(c.Return())), "", 4);
+        Assert.ThrowsAsync<LuaRuntimeException>(async () => await state.DoStringAsync(Loop));
+
+        state.MaxCallDepth = 10;
+        var fires = 0;
+        state.SetHook(
+            new(
+                (c, ct) =>
+                {
+                    fires++;
+                    return new(c.Return());
+                }
+            ),
+            "",
+            4
+        );
+        var result = await state.DoStringAsync(Loop);
+        Assert.That(result[0], Is.EqualTo((LuaValue)5050));
+        Assert.That(fires, Is.GreaterThan(0));
+    }
+
+    [TestCase(false)]
+    [TestCase(true)]
+    public async Task CountHook_CanChangeIntervalInsideCallback(bool asynchronous)
+    {
+        using var state = LuaState.Create();
+        var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var fires = 0;
+        LuaFunction hook = null!;
+        hook = new(
+            "adaptive",
+            async (c, ct) =>
+            {
+                fires++;
+                if (fires == 1)
+                {
+                    if (asynchronous)
+                    {
+                        await gate.Task;
+                    }
+
+                    state.SetHook(hook, "", 64);
+                }
+
+                return c.Return();
+            }
+        );
+        state.SetHook(hook, "", 4);
+        var execution = state.DoStringAsync(Loop).AsTask();
+        if (asynchronous)
+        {
+            Assert.That(execution.IsCompleted, Is.False);
+            gate.SetResult();
+        }
+
+        var result = await execution;
+        Assert.That(result[0], Is.EqualTo((LuaValue)5050));
+        Assert.That(fires, Is.EqualTo(4));
+    }
+
+    [Test]
+    public async Task CountHook_CanDisableItselfInsideCallback()
+    {
+        using var state = LuaState.Create();
+        var fires = 0;
+        state.SetHook(
+            new(
+                (c, ct) =>
+                {
+                    fires++;
+                    state.SetHook(null, "");
+                    return new(c.Return());
+                }
+            ),
+            "",
+            4
+        );
+        var result = await state.DoStringAsync(Loop);
+        Assert.That(result[0], Is.EqualTo((LuaValue)5050));
+        Assert.That(fires, Is.EqualTo(1));
+    }
+
+    [TestCase(false)]
+    [TestCase(true)]
+    public async Task CountHook_CanReplaceItselfInsideCallback(bool asynchronous)
+    {
+        using var state = LuaState.Create();
+        var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var originalFires = 0;
+        var replacementFires = 0;
+        var replacement = new LuaFunction(
+            (c, ct) =>
+            {
+                replacementFires++;
+                return new(c.Return());
+            }
+        );
+        state.SetHook(
+            new(
+                async (c, ct) =>
+                {
+                    originalFires++;
+                    if (asynchronous)
+                    {
+                        await gate.Task;
+                    }
+
+                    state.SetHook(replacement, "", 64);
+                    return c.Return();
+                }
+            ),
+            "",
+            4
+        );
+        var execution = state.DoStringAsync(Loop).AsTask();
+        if (asynchronous)
+        {
+            Assert.That(execution.IsCompleted, Is.False);
+            gate.SetResult();
+        }
+
+        var result = await execution;
+        Assert.That(result[0], Is.EqualTo((LuaValue)5050));
+        Assert.That(originalFires, Is.EqualTo(1));
+        Assert.That(replacementFires, Is.EqualTo(3));
+    }
+
+    [Test]
+    public async Task CountHook_DoesNotRecursivelyFireForLuaCalledByHook()
+    {
+        using var state = LuaState.Create();
+        var fires = 0;
+        state.SetHook(
+            new(
+                async (c, ct) =>
+                {
+                    fires++;
+                    if (fires == 1)
+                    {
+                        var result = await state.DoStringAsync(Loop, "hook-body.lua", ct);
+                        Assert.That(result[0], Is.EqualTo((LuaValue)5050));
+                        Assert.That(fires, Is.EqualTo(1));
+                    }
+
+                    return c.Return();
+                }
+            ),
+            "",
+            4
+        );
+        var result = await state.DoStringAsync(Loop);
+        Assert.That(result[0], Is.EqualTo((LuaValue)5050));
+        Assert.That(fires, Is.EqualTo(51));
+    }
+
+    [TestCase(false)]
+    [TestCase(true)]
+    public async Task CountHook_DoesNotFireInsideHostFunction(bool asynchronous)
+    {
+        using var state = LuaState.Create();
+        var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var fires = 0;
+        var firesAtEntry = -1;
+        var firesAtExit = -1;
+        state.SetHook(
+            new(
+                (c, ct) =>
+                {
+                    fires++;
+                    return new(c.Return());
+                }
+            ),
+            "",
+            1
+        );
+        state.Environment["host"] = new LuaFunction(
+            async (c, ct) =>
+            {
+                firesAtEntry = fires;
+                if (asynchronous)
+                {
+                    await gate.Task;
+                }
+                firesAtExit = fires;
+                return c.Return();
+            }
+        );
+
+        var execution = state
+            .DoStringAsync("host() local x=0 for i=1,100 do x=x+i end return x")
+            .AsTask();
+        Assert.That(firesAtEntry, Is.GreaterThanOrEqualTo(0));
+        if (asynchronous)
+        {
+            Assert.That(execution.IsCompleted, Is.False);
+            Assert.That(fires, Is.EqualTo(firesAtEntry));
+            gate.SetResult();
+        }
+        var result = await execution;
+        Assert.That(result[0], Is.EqualTo((LuaValue)5050));
+        Assert.That(firesAtExit, Is.EqualTo(firesAtEntry));
+        Assert.That(fires, Is.GreaterThan(firesAtExit));
+    }
+}