AsyncResourceOwner.cs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

namespace Lilith.Core.Kokoro;

/// <summary>Bounds callers' waits while retaining ownership until native work actually ends.</summary>
internal sealed class AsyncResourceOwner(Action cleanup, TimeSpan timeout, TimeSpan shutdownTimeout)
{
    private readonly object _sync = new();
    private readonly CancellationTokenSource _lifetime = new();
    private readonly HashSet<Task> _work = [];
    private bool _closing;
    private bool _faulted;
    private Task? _dispose;

    public bool IsClosing { get { lock (_sync) return _closing; } }
    public void WhileOpen(Action action) { lock (_sync) { if (!_closing) action(); } }

    public Task RunAsync(Func<CancellationToken, Task> action, CancellationToken ct = default)
    {
        Task work;
        lock (_sync)
        {
            ObjectDisposedException.ThrowIf(_closing, this);
            if (_faulted) throw new InvalidOperationException("Speech engine timed out; restart the agent before retrying speech.");
            ct.ThrowIfCancellationRequested();
            work = Task.Run(async () =>
            {
                using var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _lifetime.Token);
                await action(linked.Token);
            });
            _work.Add(work);
        }
        _ = work.ContinueWith(completed =>
        {
            _ = completed.Exception; // Observe errors even if the caller stopped waiting.
            lock (_sync) _work.Remove(completed);
        }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
        return WaitAsync(work, ct);
    }

    private async Task WaitAsync(Task work, CancellationToken ct)
    {
        try { await work.WaitAsync(timeout, ct); }
        catch (TimeoutException)
        {
            lock (_sync)
            {
                _faulted = true;
                if (!_closing) _lifetime.Cancel();
            }
            throw new TimeoutException("Kokoro operation timed out. Native work still owns its resources until it finishes.");
        }
    }

    public Task DisposeAsync()
    {
        lock (_sync)
        {
            if (_dispose is not null) return _dispose;
            _closing = true;
            _lifetime.Cancel();
            Task[] pending = [.. _work];
            _dispose = DisposeCoreAsync(pending);
            return _dispose;
        }
    }

    private async Task DisposeCoreAsync(Task[] pending)
    {
        var cleanupTask = Task.Run(async () =>
        {
            try { await Task.WhenAll(pending); }
            catch { /* Completion, including failure, releases ownership. */ }
            finally { try { cleanup(); } finally { _lifetime.Dispose(); } }
        });
        // If native work is stuck, cleanup remains attached to it and never frees live resources.
        try { await cleanupTask.WaitAsync(shutdownTimeout); }
        catch (TimeoutException)
        {
            _ = cleanupTask.ContinueWith(task => _ = task.Exception, TaskScheduler.Default);
        }
    }
}