Improve meeting audio mixing pipeline

MeetingAssistant/Recording/CompositeMeetingAudioSource.cs

@@ -5,26 +5,44 @@ 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, channel.Writer, cancellationToken);
-        var systemTask = PumpAsync(AudioSourceKind.System, systemAudio, channel.Writer, cancellationToken);
+        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 =>
@@ -41,8 +59,8 @@ public sealed class CompositeMeetingAudioSource : IMeetingAudioSource
                 TaskContinuationOptions.ExecuteSynchronously,
                 TaskScheduler.Default);
 
-        AudioChunk? pendingMicrophone = null;
-        AudioChunk? pendingSystem = null;
+        var pendingMicrophone = new PendingAudioStream();
+        var pendingSystem = new PendingAudioStream();
 
         await foreach (var sourceChunk in channel.Reader.ReadAllAsync(cancellationToken))
         {
@@ -52,22 +70,32 @@ public sealed class CompositeMeetingAudioSource : IMeetingAudioSource
                 sourceChunk.Kind);
             if (sourceChunk.Kind == AudioSourceKind.Microphone)
             {
-                pendingMicrophone = sourceChunk.Chunk;
+                pendingMicrophone.Enqueue(sourceChunk.Chunk);
             }
             else
             {
-                pendingSystem = sourceChunk.Chunk;
+                pendingSystem.Enqueue(sourceChunk.Chunk);
             }
 
-            if (pendingMicrophone is not null && pendingSystem is not null)
+            while (TryReadAlignedPair(pendingMicrophone, pendingSystem, out var microphoneChunk, out var systemChunk))
             {
-                yield return Mix(pendingMicrophone, pendingSystem);
-                pendingMicrophone = null;
-                pendingSystem = null;
-                continue;
+                yield return CleanAndMix(microphoneChunk, systemChunk, echoCanceller, microphoneGain, systemAudioGain);
             }
 
-            using var timeout = new CancellationTokenSource(TimeSpan.FromMilliseconds(50));
+            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
             {
@@ -80,39 +108,33 @@ public sealed class CompositeMeetingAudioSource : IMeetingAudioSource
             {
             }
 
-            if (pendingMicrophone is not null)
+            while (TryReadRemainder(pendingMicrophone, pendingSystem, out var timeoutMicrophoneChunk, out var timeoutSystemChunk))
             {
-                yield return pendingMicrophone;
-                pendingMicrophone = null;
-            }
-
-            if (pendingSystem is not null)
-            {
-                yield return pendingSystem;
-                pendingSystem = null;
+                yield return CleanAndMix(
+                    timeoutMicrophoneChunk,
+                    timeoutSystemChunk,
+                    echoCanceller,
+                    microphoneGain,
+                    systemAudioGain);
             }
         }
 
-        if (pendingMicrophone is not null)
+        while (TryReadRemainder(pendingMicrophone, pendingSystem, out var microphoneChunk, out var systemChunk))
         {
-            yield return pendingMicrophone;
-        }
-
-        if (pendingSystem is not null)
-        {
-            yield return pendingSystem;
+            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(cancellationToken))
+            await foreach (var chunk in source.CaptureAsync(options, cancellationToken))
             {
                 await writer.WriteAsync(new SourceChunk(kind, chunk), cancellationToken);
             }
@@ -128,32 +150,49 @@ public sealed class CompositeMeetingAudioSource : IMeetingAudioSource
         }
     }
 
-    private AudioChunk Mix(AudioChunk left, AudioChunk right)
+    private AudioChunk CleanAndMix(
+        AudioChunk microphoneChunk,
+        AudioChunk systemChunk,
+        IAcousticEchoCanceller echoCanceller,
+        double microphoneGain,
+        double systemAudioGain)
     {
-        if (left.SampleRate != right.SampleRate || left.Channels != right.Channels)
+        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}.",
-                left.SampleRate,
-                left.Channels,
-                right.SampleRate,
-                right.Channels);
-            return left;
+                microphoneChunk.SampleRate,
+                microphoneChunk.Channels,
+                systemChunk.SampleRate,
+                systemChunk.Channels);
+            return microphoneChunk;
         }
 
-        var mixed = new byte[Math.Max(left.Pcm.Length, right.Pcm.Length)];
+        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(left.Pcm, byteIndex);
-            var rightSample = ReadSample(right.Pcm, byteIndex);
-            var sample = Math.Clamp(leftSample + rightSample, short.MinValue, short.MaxValue);
+            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, left.SampleRate, left.Channels);
+        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)
@@ -161,6 +200,144 @@ public sealed class CompositeMeetingAudioSource : IMeetingAudioSource
         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