Return shared ArrayPool (#459)

* Update OpenAIEmbedding.cs

* Update AsyncWebsocketMessageEnumerator.cs

* Update RealtimeConversationSession.cs

* Update ConversationTests.cs

Not used

* Move to finally

Author: shargon

src/Custom/Embeddings/OpenAIEmbedding.cs

@@ -118,32 +118,36 @@ private static ReadOnlyMemory<float> ConvertToVectorOfFloats(BinaryData binaryDa
 
         // Decode base64 string to bytes.
         byte[] bytes = ArrayPool<byte>.Shared.Rent(Base64.GetMaxDecodedFromUtf8Length(base64.Length));
-        OperationStatus status = Base64.DecodeFromUtf8(base64, bytes.AsSpan(), out int bytesConsumed, out int bytesWritten);
-        if (status != OperationStatus.Done || bytesWritten % sizeof(float) != 0)
-        {
-            ThrowInvalidData();
-        }
-
-        // Interpret bytes as floats
-        float[] vector = new float[bytesWritten / sizeof(float)];
-        bytes.AsSpan(0, bytesWritten).CopyTo(MemoryMarshal.AsBytes(vector.AsSpan()));
-        if (!BitConverter.IsLittleEndian)
-        {
-            Span<int> ints = MemoryMarshal.Cast<float, int>(vector.AsSpan());
+        try
+        {
+            OperationStatus status = Base64.DecodeFromUtf8(base64, bytes.AsSpan(), out int bytesConsumed, out int bytesWritten);
+            if (status != OperationStatus.Done || bytesWritten % sizeof(float) != 0)
+            {
+                ThrowInvalidData();
+            }
+
+            // Interpret bytes as floats
+            float[] vector = new float[bytesWritten / sizeof(float)];
+            bytes.AsSpan(0, bytesWritten).CopyTo(MemoryMarshal.AsBytes(vector.AsSpan()));
+            if (!BitConverter.IsLittleEndian)
+            {
+                Span<int> ints = MemoryMarshal.Cast<float, int>(vector.AsSpan());
 #if NET8_0_OR_GREATER
-            BinaryPrimitives.ReverseEndianness(ints, ints);
+                BinaryPrimitives.ReverseEndianness(ints, ints);
 #else
             for (int i = 0; i < ints.Length; i++)
             {
                 ints[i] = BinaryPrimitives.ReverseEndianness(ints[i]);
             }
 #endif
+            }
+            return new ReadOnlyMemory<float>(vector);
+        }
+        finally
+        {
+            ArrayPool<byte>.Shared.Return(bytes);
         }
-
-        ArrayPool<byte>.Shared.Return(bytes);
-        return new ReadOnlyMemory<float>(vector);
-
-        static void ThrowInvalidData() =>
+    }
+    static void ThrowInvalidData() =>
             throw new FormatException("The input is not a valid Base64 string of encoded floats.");
-    }
 }

src/Custom/RealtimeConversation/Internal/AsyncWebsocketMessageEnumerator.cs

@@ -26,6 +26,7 @@ public AsyncWebsocketMessageResultEnumerator(WebSocket webSocket, CancellationTo
 
     public ValueTask DisposeAsync()
     {
+        ArrayPool<byte>.Shared.Return(_receiveBuffer);
         _webSocket?.Dispose();
         return new ValueTask(Task.CompletedTask);
     }
@@ -50,4 +51,4 @@ public async ValueTask<bool> MoveNextAsync()
         Current = ClientResult.FromResponse(websocketPipelineResponse);
         return true;
     }
-}
+}

src/Custom/RealtimeConversation/RealtimeConversationSession.cs

@@ -55,9 +55,9 @@ public virtual async Task SendInputAudioAsync(Stream audio, CancellationToken ca
             }
             _isSendingAudioStream = true;
         }
+        byte[] buffer = ArrayPool<byte>.Shared.Rent(1024 * 16);
         try
         {
-            byte[] buffer = ArrayPool<byte>.Shared.Rent(1024 * 16);
             while (true)
             {
                 int bytesRead = await audio.ReadAsync(buffer, 0, buffer.Length, cancellationToken).ConfigureAwait(false);
@@ -75,6 +75,7 @@ public virtual async Task SendInputAudioAsync(Stream audio, CancellationToken ca
         }
         finally
         {
+            ArrayPool<byte>.Shared.Return(buffer);
             using (await _audioSendSemaphore.AutoReleaseWaitAsync(cancellationToken).ConfigureAwait(false))
             {
                 _isSendingAudioStream = false;
@@ -93,9 +94,9 @@ public virtual void SendInputAudio(Stream audio, CancellationToken cancellationT
             }
             _isSendingAudioStream = true;
         }
+        byte[] buffer = ArrayPool<byte>.Shared.Rent(1024 * 16);
         try
         {
-            byte[] buffer = ArrayPool<byte>.Shared.Rent(1024 * 16);
             while (true)
             {
                 int bytesRead = audio.Read(buffer, 0, buffer.Length);
@@ -113,6 +114,7 @@ public virtual void SendInputAudio(Stream audio, CancellationToken cancellationT
         }
         finally
         {
+            ArrayPool<byte>.Shared.Return(buffer);
             using (_audioSendSemaphore.AutoReleaseWait(cancellationToken))
             {
                 _isSendingAudioStream = false;
@@ -349,4 +351,4 @@ public void Dispose()
     {
         WebSocket?.Dispose();
     }
-}
+}

tests/RealtimeConversation/ConversationTests.cs

@@ -341,7 +341,6 @@ public async Task AudioWithToolsWorks(TestAudioSendType audioSendType)
             {
                 byte[] allAudioBytes = await File.ReadAllBytesAsync(inputAudioFilePath, CancellationToken);
                 const int audioSendBufferLength = 8 * 1024;
-                byte[] audioSendBuffer = ArrayPool<byte>.Shared.Rent(audioSendBufferLength);
                 for (int readPos = 0; readPos < allAudioBytes.Length; readPos += audioSendBufferLength)
                 {
                     int nextSegmentLength = Math.Min(audioSendBufferLength, allAudioBytes.Length - readPos);