meeting-assistant/MeetingAssistant/Transcription/FunAsrDockerBackendLifecycle.cs

using Microsoft.Extensions.Options;

namespace MeetingAssistant.Transcription;

public interface IFunAsrBackendLifecycle
{
    Task EnsureStartedAsync(CancellationToken cancellationToken);
}

public sealed class FunAsrDockerBackendLifecycle : IFunAsrBackendLifecycle, IAsyncDisposable
{
    private readonly ICommandRunner commandRunner;
    private readonly IFunAsrBackendReadinessProbe readinessProbe;
    private readonly MeetingAssistantOptions options;
    private readonly ILogger<FunAsrDockerBackendLifecycle> logger;
    private readonly SemaphoreSlim gate = new(1, 1);
    private bool containerStarted;
    private bool started;
    private bool disposed;

    public FunAsrDockerBackendLifecycle(
        ICommandRunner commandRunner,
        IFunAsrBackendReadinessProbe readinessProbe,
        IOptions<MeetingAssistantOptions> options,
        ILogger<FunAsrDockerBackendLifecycle> logger)
    {
        this.commandRunner = commandRunner;
        this.readinessProbe = readinessProbe;
        this.options = options.Value;
        this.logger = logger;
    }

    public async Task EnsureStartedAsync(CancellationToken cancellationToken)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        var backend = options.FunAsr.Backend;
        if (!backend.Enabled)
        {
            return;
        }

        await gate.WaitAsync(cancellationToken);
        try
        {
            if (started)
            {
                return;
            }

            PrepareModelsFolder(backend);
            await RunRequiredAsync(backend.DockerCommand, ["pull", backend.Image], cancellationToken);
            await commandRunner.RunAsync(
                backend.DockerCommand,
                ["rm", "-f", backend.ContainerName],
                cancellationToken);
            await RunRequiredAsync(backend.DockerCommand, BuildDockerRunArguments(), cancellationToken);
            containerStarted = true;
            await readinessProbe.WaitUntilReadyAsync(
                new Uri(options.FunAsr.Endpoint),
                backend.StartupTimeout,
                cancellationToken);

            started = true;
            logger.LogInformation(
                "Started managed FunASR backend container {ContainerName}; endpoint {Endpoint} is ready",
                backend.ContainerName,
                options.FunAsr.Endpoint);
        }
        finally
        {
            gate.Release();
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (disposed)
        {
            return;
        }

        var backend = options.FunAsr.Backend;
        if (!backend.Enabled || !backend.StopOnDispose)
        {
            disposed = true;
            gate.Dispose();
            return;
        }

        await gate.WaitAsync();
        try
        {
            if (containerStarted)
            {
                await commandRunner.RunAsync(
                    backend.DockerCommand,
                    ["stop", backend.ContainerName],
                    CancellationToken.None);
                containerStarted = false;
                started = false;
            }
        }
        finally
        {
            disposed = true;
            gate.Release();
            gate.Dispose();
        }
    }

    private async Task RunRequiredAsync(
        string fileName,
        IReadOnlyList<string> arguments,
        CancellationToken cancellationToken)
    {
        var backend = options.FunAsr.Backend;
        using var timeoutSource = backend.CommandTimeout > TimeSpan.Zero
            ? new CancellationTokenSource(backend.CommandTimeout)
            : null;
        using var linkedSource = timeoutSource is null
            ? null
            : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutSource.Token);
        var effectiveToken = linkedSource?.Token ?? cancellationToken;
        CommandResult result;
        try
        {
            result = await commandRunner.RunAsync(fileName, arguments, effectiveToken);
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested && timeoutSource?.IsCancellationRequested == true)
        {
            throw new TimeoutException(
                $"Command '{fileName} {string.Join(' ', arguments)}' timed out after {backend.CommandTimeout}.");
        }

        if (result.ExitCode != 0)
        {
            throw new InvalidOperationException(
                $"Command '{fileName} {string.Join(' ', arguments)}' failed with exit code {result.ExitCode}: {result.StandardError}");
        }
    }

    private string[] BuildDockerRunArguments()
    {
        var backend = options.FunAsr.Backend;
        List<string> arguments =
        [
            "run",
            "--rm",
            "-d",
        ];

        if (backend.Privileged)
        {
            arguments.Add("--privileged=true");
        }

        arguments.AddRange(
        [
            "--name",
            backend.ContainerName,
            "-p",
            $"{GetEndpointPort()}:{backend.ContainerPort}",
            "-v",
            $"{VaultPath.Resolve(backend.ModelsFolder)}:/workspace/models",
            backend.Image,
            "bash",
            "-lc",
            backend.ServerCommand
        ]);

        return [.. arguments];
    }

    private static void PrepareModelsFolder(FunAsrBackendOptions backend)
    {
        var modelsFolder = VaultPath.Resolve(backend.ModelsFolder);
        Directory.CreateDirectory(modelsFolder);

        var hotwordsPath = Path.Combine(modelsFolder, "hotwords.txt");
        if (!File.Exists(hotwordsPath))
        {
            File.WriteAllText(hotwordsPath, "");
        }
    }

    private int GetEndpointPort()
    {
        return new Uri(options.FunAsr.Endpoint).Port;
    }
}