meeting-assistant/MeetingAssistant/Recording/CompositeMeetingAudioSource.cs

using System.Runtime.CompilerServices;
using System.Threading.Channels;

namespace MeetingAssistant.Recording;

public sealed class CompositeMeetingAudioSource : IMeetingAudioSource
{
    private static readonly TimeSpan SourceSkewTimeout = TimeSpan.FromMilliseconds(250);
    private readonly IMeetingAudioSource microphone;
    private readonly IMeetingAudioSource systemAudio;
    private readonly IAcousticEchoCancellerFactory echoCancellerFactory;
    private readonly ILogger<CompositeMeetingAudioSource> logger;

    public CompositeMeetingAudioSource(
        IMeetingAudioSource microphone,
        IMeetingAudioSource systemAudio,
        IAcousticEchoCancellerFactory echoCancellerFactory,
        ILogger<CompositeMeetingAudioSource> logger)
    {
        this.microphone = microphone;
        this.systemAudio = systemAudio;
        this.echoCancellerFactory = echoCancellerFactory;
        this.logger = logger;
    }

    public async IAsyncEnumerable<AudioChunk> CaptureAsync(
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        await foreach (var chunk in CaptureAsync(new MeetingAssistantOptions(), cancellationToken))
        {
            yield return chunk;
        }
    }

    public async IAsyncEnumerable<AudioChunk> CaptureAsync(
        MeetingAssistantOptions options,
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var echoCanceller = echoCancellerFactory.Create();
        echoCanceller.Reset();
        var microphoneGain = NormalizeGain(options.Recording.MicrophoneMixGain);
        var systemAudioGain = NormalizeGain(options.Recording.SystemAudioMixGain);
        var channel = Channel.CreateUnbounded<SourceChunk>();
        var microphoneTask = PumpAsync(AudioSourceKind.Microphone, microphone, options, channel.Writer, cancellationToken);
        var systemTask = PumpAsync(AudioSourceKind.System, systemAudio, options, channel.Writer, cancellationToken);
        _ = Task.WhenAll(microphoneTask, systemTask)
            .ContinueWith(
                task =>
                {
                    if (task.Exception is not null)
                    {
                        channel.Writer.TryComplete(task.Exception);
                        return;
                    }

                    channel.Writer.TryComplete();
                },
                CancellationToken.None,
                TaskContinuationOptions.ExecuteSynchronously,
                TaskScheduler.Default);

        var pendingMicrophone = new PendingAudioStream();
        var pendingSystem = new PendingAudioStream();

        await foreach (var sourceChunk in channel.Reader.ReadAllAsync(cancellationToken))
        {
            logger.LogDebug(
                "Captured {ByteCount} byte(s) from {SourceKind}",
                sourceChunk.Chunk.Pcm.Length,
                sourceChunk.Kind);
            if (sourceChunk.Kind == AudioSourceKind.Microphone)
            {
                pendingMicrophone.Enqueue(sourceChunk.Chunk);
            }
            else
            {
                pendingSystem.Enqueue(sourceChunk.Chunk);
            }

            while (TryReadAlignedPair(pendingMicrophone, pendingSystem, out var microphoneChunk, out var systemChunk))
            {
                yield return CleanAndMix(microphoneChunk, systemChunk, echoCanceller, microphoneGain, systemAudioGain);
            }

            if (ShouldFlushTimedOutRemainder(pendingMicrophone, pendingSystem))
            {
                while (TryReadRemainder(pendingMicrophone, pendingSystem, out var bufferedMicrophoneChunk, out var bufferedSystemChunk))
                {
                    yield return CleanAndMix(
                        bufferedMicrophoneChunk,
                        bufferedSystemChunk,
                        echoCanceller,
                        microphoneGain,
                        systemAudioGain);
                }
            }

            using var timeout = new CancellationTokenSource(SourceSkewTimeout);
            using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeout.Token);
            try
            {
                if (await channel.Reader.WaitToReadAsync(linked.Token))
                {
                    continue;
                }
            }
            catch (OperationCanceledException) when (timeout.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
            {
            }

            while (TryReadRemainder(pendingMicrophone, pendingSystem, out var timeoutMicrophoneChunk, out var timeoutSystemChunk))
            {
                yield return CleanAndMix(
                    timeoutMicrophoneChunk,
                    timeoutSystemChunk,
                    echoCanceller,
                    microphoneGain,
                    systemAudioGain);
            }
        }

        while (TryReadRemainder(pendingMicrophone, pendingSystem, out var microphoneChunk, out var systemChunk))
        {
            yield return CleanAndMix(microphoneChunk, systemChunk, echoCanceller, microphoneGain, systemAudioGain);
        }
    }

    private static async Task PumpAsync(
        AudioSourceKind kind,
        IMeetingAudioSource source,
        MeetingAssistantOptions options,
        ChannelWriter<SourceChunk> writer,
        CancellationToken cancellationToken)
    {
        try
        {
            await foreach (var chunk in source.CaptureAsync(options, cancellationToken))
            {
                await writer.WriteAsync(new SourceChunk(kind, chunk), cancellationToken);
            }
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            throw;
        }
        catch (Exception exception)
        {
            writer.TryComplete(new InvalidOperationException($"{kind} audio capture failed.", exception));
            throw;
        }
    }

    private AudioChunk CleanAndMix(
        AudioChunk microphoneChunk,
        AudioChunk systemChunk,
        IAcousticEchoCanceller echoCanceller,
        double microphoneGain,
        double systemAudioGain)
    {
        if (microphoneChunk.SampleRate != systemChunk.SampleRate || microphoneChunk.Channels != systemChunk.Channels)
        {
            logger.LogWarning(
                "Cannot mix audio chunks with different formats. Using microphone format {MicrophoneSampleRate}/{MicrophoneChannels} and dropping system chunk format {SystemSampleRate}/{SystemChannels}.",
                microphoneChunk.SampleRate,
                microphoneChunk.Channels,
                systemChunk.SampleRate,
                systemChunk.Channels);
            return microphoneChunk;
        }

        var cleanedMicrophone = echoCanceller.CleanMicrophone(microphoneChunk, systemChunk);
        var mixed = new byte[Math.Max(cleanedMicrophone.Pcm.Length, systemChunk.Pcm.Length)];
        var sampleCount = mixed.Length / sizeof(short);

        for (var sampleIndex = 0; sampleIndex < sampleCount; sampleIndex++)
        {
            var byteIndex = sampleIndex * sizeof(short);
            var leftSample = ReadSample(cleanedMicrophone.Pcm, byteIndex);
            var rightSample = ReadSample(systemChunk.Pcm, byteIndex);
            var sample = Math.Clamp(
                (int)Math.Round(leftSample * microphoneGain + rightSample * systemAudioGain),
                short.MinValue,
                short.MaxValue);
            BitConverter.GetBytes((short)sample).CopyTo(mixed, byteIndex);
        }

        return new AudioChunk(
            mixed,
            cleanedMicrophone.SampleRate,
            cleanedMicrophone.Channels);
    }

    private static double NormalizeGain(double gain)
    {
        return double.IsFinite(gain) && gain >= 0 ? gain : 0;
    }

    private static int ReadSample(byte[] pcm, int byteIndex)
    {
        return byteIndex + 1 < pcm.Length ? BitConverter.ToInt16(pcm, byteIndex) : 0;
    }

    private static bool TryReadAlignedPair(
        PendingAudioStream microphone,
        PendingAudioStream systemAudio,
        out AudioChunk microphoneChunk,
        out AudioChunk systemChunk)
    {
        microphoneChunk = null!;
        systemChunk = null!;
        var byteCount = Math.Min(microphone.ByteCount, systemAudio.ByteCount);
        byteCount -= byteCount % sizeof(short);
        if (byteCount <= 0 || microphone.Format is null || systemAudio.Format is null)
        {
            return false;
        }

        microphoneChunk = microphone.Read(byteCount);
        systemChunk = systemAudio.Read(byteCount);
        return true;
    }

    private static bool TryReadRemainder(
        PendingAudioStream microphone,
        PendingAudioStream systemAudio,
        out AudioChunk microphoneChunk,
        out AudioChunk systemChunk)
    {
        microphoneChunk = null!;
        systemChunk = null!;
        var byteCount = Math.Max(microphone.ByteCount, systemAudio.ByteCount);
        byteCount -= byteCount % sizeof(short);
        var format = microphone.Format ?? systemAudio.Format;
        if (byteCount <= 0 || format is null)
        {
            return false;
        }

        microphoneChunk = microphone.ByteCount > 0
            ? microphone.Read(Math.Min(byteCount, microphone.ByteCount))
            : new AudioChunk(new byte[byteCount], format.SampleRate, format.Channels);
        systemChunk = systemAudio.ByteCount > 0
            ? systemAudio.Read(Math.Min(byteCount, systemAudio.ByteCount))
            : new AudioChunk(new byte[byteCount], format.SampleRate, format.Channels);

        if (microphoneChunk.Pcm.Length < byteCount)
        {
            microphoneChunk = microphoneChunk with { Pcm = PadSilence(microphoneChunk.Pcm, byteCount) };
        }

        if (systemChunk.Pcm.Length < byteCount)
        {
            systemChunk = systemChunk with { Pcm = PadSilence(systemChunk.Pcm, byteCount) };
        }

        return true;
    }

    private static bool ShouldFlushTimedOutRemainder(PendingAudioStream microphone, PendingAudioStream systemAudio)
    {
        return microphone.ByteCount > 0 && systemAudio.ByteCount == 0 && microphone.HasBufferedFor(SourceSkewTimeout) ||
            systemAudio.ByteCount > 0 && microphone.ByteCount == 0 && systemAudio.HasBufferedFor(SourceSkewTimeout);
    }

    private static byte[] PadSilence(byte[] pcm, int byteCount)
    {
        var padded = new byte[byteCount];
        Buffer.BlockCopy(pcm, 0, padded, 0, pcm.Length);
        return padded;
    }

    private sealed class PendingAudioStream
    {
        private readonly LinkedList<ArraySegment<byte>> chunks = new();
        private long? firstBufferedAt;

        public int ByteCount { get; private set; }

        public AudioFormat? Format { get; private set; }

        public void Enqueue(AudioChunk chunk)
        {
            if (chunk.Pcm.Length == 0)
            {
                return;
            }

            firstBufferedAt ??= Environment.TickCount64;
            Format ??= new AudioFormat(chunk.SampleRate, chunk.Channels);
            chunks.AddLast(new ArraySegment<byte>(chunk.Pcm));
            ByteCount += chunk.Pcm.Length;
        }

        public bool HasBufferedFor(TimeSpan duration)
        {
            return firstBufferedAt is not null &&
                TimeSpan.FromMilliseconds(Environment.TickCount64 - firstBufferedAt.Value) >= duration;
        }

        public AudioChunk Read(int byteCount)
        {
            if (Format is null)
            {
                throw new InvalidOperationException("Cannot read audio before a stream format has been captured.");
            }

            var target = new byte[byteCount];
            var written = 0;
            while (written < byteCount && chunks.First is not null)
            {
                var chunk = chunks.First.Value;
                var copyCount = Math.Min(byteCount - written, chunk.Count);
                Buffer.BlockCopy(chunk.Array!, chunk.Offset, target, written, copyCount);
                written += copyCount;
                ByteCount -= copyCount;

                if (copyCount == chunk.Count)
                {
                    chunks.RemoveFirst();
                }
                else
                {
                    chunks.First.Value = new ArraySegment<byte>(
                        chunk.Array!,
                        chunk.Offset + copyCount,
                        chunk.Count - copyCount);
                }
            }

            if (ByteCount == 0)
            {
                firstBufferedAt = null;
            }

            return new AudioChunk(target, Format.SampleRate, Format.Channels);
        }
    }

    private sealed record AudioFormat(int SampleRate, int Channels);

    private sealed record SourceChunk(AudioSourceKind Kind, AudioChunk Chunk);

    private enum AudioSourceKind
    {
        Microphone,
        System
    }
}