feat: add tests for socket pool

Author: ikesnowy

sample/Cnblogs.DashScope.Sample/Program.cs

@@ -80,7 +80,17 @@
             }
 
             writer.Close();
-            Console.WriteLine($"audio saved to {file.FullName}");
+
+            var tokenUsage = 0;
+            await foreach (var message in tts.GetMessagesAsync())
+            {
+                if (message.Payload.Usage?.Characters > tokenUsage)
+                {
+                    tokenUsage = message.Payload.Usage.Characters;
+                }
+            }
+
+            Console.WriteLine($"audio saved to {file.FullName}, token usage: {tokenUsage}");
             break;
         }
 }

src/Cnblogs.DashScope.AspNetCore/ServiceCollectionInjector.cs

@@ -72,7 +72,7 @@ public static IHttpClientBuilder AddDashScopeClient(
 
             if (baseWebsocketAddress != null)
             {
-                o.BaseWebsocketAddress = baseWebsocketAddress;
+                o.WebsocketBaseAddress = baseWebsocketAddress;
             }
 
             o.WorkspaceId = workspaceId;
@@ -87,8 +87,11 @@ private static IHttpClientBuilder AddDashScopeHttpClient(
         string? baseAddress,
         string? workspaceId)
     {
+        services.AddSingleton<IDashScopeClientWebSocketFactory, DashScopeClientWebSocketFactory>();
         services.AddSingleton<DashScopeClientWebSocketPool>(sp
-            => new DashScopeClientWebSocketPool(sp.GetRequiredService<IOptions<DashScopeOptions>>().Value));
+            => new DashScopeClientWebSocketPool(
+                sp.GetRequiredService<IDashScopeClientWebSocketFactory>(),
+                sp.GetRequiredService<IOptions<DashScopeOptions>>().Value));
         services.AddScoped<IDashScopeClient, DashScopeClientAspNetCore>();
         return services.AddHttpClient(
             DashScopeAspNetCoreDefaults.DefaultHttpClientName,

src/Cnblogs.DashScope.Core/DashScopeClient.cs

@@ -49,10 +49,11 @@ private static DashScopeClientWebSocketPool GetConfiguredSocketPool(
         if (pool is null)
         {
             pool = new DashScopeClientWebSocketPool(
+                new DashScopeClientWebSocketFactory(),
                 new DashScopeOptions
                 {
                     ApiKey = apiKey,
-                    BaseWebsocketAddress = baseAddress,
+                    WebsocketBaseAddress = baseAddress,
                     SocketPoolSize = socketPoolSize,
                     WorkspaceId = workspaceId
                 });

src/Cnblogs.DashScope.Core/DashScopeClientCore.cs

@@ -283,7 +283,7 @@ public async Task<SpeechSynthesizerSocketSession> CreateSpeechSynthesizerSocketS
         string modelId,
         CancellationToken cancellationToken = default)
     {
-        var socket = await _socketPool.RentSocketAsync<SpeechSynthesizerOutput>(cancellationToken);
+        var socket = await _socketPool.RentSocketAsync(cancellationToken);
         return new SpeechSynthesizerSocketSession(socket, modelId);
     }
 

src/Cnblogs.DashScope.Core/DashScopeClientWebSocket.cs

@@ -23,6 +23,12 @@ public sealed class DashScopeClientWebSocket : IDisposable
     private Task? _receiveTask;
     private TaskCompletionSource<bool> _taskStartedSignal = new();
     private Channel<byte>? _binaryOutput;
+    private Channel<DashScopeWebSocketResponse<JsonElement>>? _jsonOutput;
+
+    /// <summary>
+    /// Unique id of this socket.
+    /// </summary>
+    internal Guid Id { get; } = Guid.NewGuid();
 
     /// <summary>
     /// The binary output.
@@ -31,6 +37,14 @@ public ChannelReader<byte> BinaryOutput
         => _binaryOutput?.Reader
            ?? throw new InvalidOperationException("Please call ResetOutput() before accessing output");
 
+    /// <summary>
+    /// The json output.
+    /// </summary>
+    /// <exception cref="InvalidOperationException">Throws when ResetOutput is not called.</exception>
+    public ChannelReader<DashScopeWebSocketResponse<JsonElement>> JsonOutput
+        => _jsonOutput?.Reader
+           ?? throw new InvalidOperationException("Please call ResetOutput() before accessing output");
+
     /// <summary>
     /// A task that completed when received task-started event.
     /// </summary>
@@ -72,13 +86,11 @@ internal DashScopeClientWebSocket(IClientWebSocket socket)
     /// <param name="uri">Websocket API uri.</param>
     /// <param name="cancellationToken">The cancellation token to use.</param>
     /// <returns></returns>
-    /// <typeparam name="TOutput">The type of the response content.</typeparam>
     /// <exception cref="OperationCanceledException">When <paramref name="cancellationToken"/> was request.</exception>
-    public async Task ConnectAsync<TOutput>(Uri uri, CancellationToken cancellationToken = default)
-        where TOutput : class
+    public async Task ConnectAsync(Uri uri, CancellationToken cancellationToken = default)
     {
         await _socket.ConnectAsync(uri, cancellationToken);
-        _receiveTask = ReceiveMessagesAsync<TOutput>(cancellationToken);
+        _receiveTask = ReceiveMessagesAsync(cancellationToken);
         State = DashScopeWebSocketState.Ready;
     }
 
@@ -89,6 +101,8 @@ public void ResetOutput()
     {
         _binaryOutput?.Writer.TryComplete();
         _binaryOutput = Channel.CreateUnbounded<byte>(UnboundedChannelOptions);
+        _jsonOutput?.Writer.TryComplete();
+        _jsonOutput = Channel.CreateUnbounded<DashScopeWebSocketResponse<JsonElement>>(UnboundedChannelOptions);
         _taskStartedSignal.TrySetResult(false);
         _taskStartedSignal = new TaskCompletionSource<bool>();
     }
@@ -125,7 +139,6 @@ public Task SendMessageAsync<TInput, TParameter>(
 
     private async Task<DashScopeWebSocketResponse<TOutput>?> ReceiveMessageAsync<TOutput>(
         CancellationToken cancellationToken = default)
-        where TOutput : class
     {
         var buffer = new byte[1024 * 4];
         var segment = new ArraySegment<byte>(buffer);
@@ -169,19 +182,22 @@ public Task SendMessageAsync<TInput, TParameter>(
     /// Wait for server response.
     /// </summary>
     /// <param name="cancellationToken">A cancellation token used to propagate notification that this operation should be canceled.</param>
-    /// <typeparam name="TOutput">Type of the response content.</typeparam>
     /// <exception cref="DashScopeException">The task was failed.</exception>
-    private async Task ReceiveMessagesAsync<TOutput>(CancellationToken cancellationToken = default)
-        where TOutput : class
+    private async Task ReceiveMessagesAsync(CancellationToken cancellationToken = default)
     {
         while (State != DashScopeWebSocketState.Closed && _socket.CloseStatus == null)
         {
-            var json = await ReceiveMessageAsync<TOutput>(cancellationToken);
+            var json = await ReceiveMessageAsync<JsonElement>(cancellationToken);
             if (json == null)
             {
                 continue;
             }
 
+            if (_jsonOutput is not null)
+            {
+                await _jsonOutput.Writer.WriteAsync(json, cancellationToken);
+            }
+
             var eventStr = json.Header.Event;
             switch (eventStr)
             {
@@ -192,10 +208,12 @@ private async Task ReceiveMessagesAsync<TOutput>(CancellationToken cancellationT
                 case "task-finished":
                     State = DashScopeWebSocketState.Ready;
                     _binaryOutput?.Writer.Complete();
+                    _jsonOutput?.Writer.Complete();
                     break;
                 case "task-failed":
                     await CloseAsync(cancellationToken);
                     _binaryOutput?.Writer.Complete();
+                    _jsonOutput?.Writer.Complete();
                     throw new DashScopeException(
                         null,
                         400,
@@ -230,6 +248,7 @@ private void Dispose(bool disposing)
             // Dispose managed resources.
             _socket.Dispose();
             _binaryOutput?.Writer.TryComplete();
+            _jsonOutput?.Writer.TryComplete();
         }
     }
 

src/Cnblogs.DashScope.Core/DashScopeClientWebSocketFactory.cs

@@ -0,0 +1,13 @@
+namespace Cnblogs.DashScope.Core;
+
+/// <summary>
+/// Default implementation for <see cref="IDashScopeClientWebSocketFactory"/>.
+/// </summary>
+public class DashScopeClientWebSocketFactory : IDashScopeClientWebSocketFactory
+{
+    /// <inheritdoc />
+    public DashScopeClientWebSocket GetClientWebSocket(string apiKey, string? workspaceId = null)
+    {
+        return new DashScopeClientWebSocket(apiKey, workspaceId);
+    }
+}

src/Cnblogs.DashScope.Core/DashScopeClientWebSocketPool.cs

@@ -8,29 +8,50 @@ namespace Cnblogs.DashScope.Core;
 public sealed class DashScopeClientWebSocketPool : IDisposable
 {
     private readonly ConcurrentBag<DashScopeClientWebSocket> _available = new();
-    private readonly ConcurrentBag<DashScopeClientWebSocket> _active = new();
+    private readonly ConcurrentDictionary<Guid, DashScopeClientWebSocket> _active = new();
     private readonly DashScopeOptions _options;
+    private readonly IDashScopeClientWebSocketFactory _dashScopeClientWebSocketFactory;
 
     /// <summary>
     /// Socket pool for DashScope API.
     /// </summary>
+    /// <param name="dashScopeClientWebSocketFactory"></param>
     /// <param name="options">Options for DashScope sdk.</param>
-    public DashScopeClientWebSocketPool(DashScopeOptions options)
+    public DashScopeClientWebSocketPool(
+        IDashScopeClientWebSocketFactory dashScopeClientWebSocketFactory,
+        DashScopeOptions options)
     {
+        _dashScopeClientWebSocketFactory = dashScopeClientWebSocketFactory;
         _options = options;
     }
 
-    internal DashScopeClientWebSocketPool(IEnumerable<DashScopeClientWebSocket> sockets)
+    /// <summary>
+    /// Get available connection count.
+    /// </summary>
+    internal int AvailableSocketCount => _available.Count;
+
+    /// <summary>
+    /// Get active connection count.
+    /// </summary>
+    internal int ActiveSocketCount => _active.Count;
+
+    internal DashScopeClientWebSocketPool(
+        IEnumerable<DashScopeClientWebSocket> sockets,
+        IDashScopeClientWebSocketFactory dashScopeClientWebSocketFactory)
     {
         _options = new DashScopeOptions();
         foreach (var socket in sockets)
         {
             _available.Add(socket);
         }
+
+        _dashScopeClientWebSocketFactory = dashScopeClientWebSocketFactory;
     }
 
-    internal void ReturnSocketAsync(DashScopeClientWebSocket socket)
+    internal void ReturnSocket(DashScopeClientWebSocket socket)
     {
+        _active.Remove(socket.Id, out _);
+
         if (socket.State != DashScopeWebSocketState.Ready)
         {
             // not returnable, disposing.
@@ -45,11 +66,8 @@ internal void ReturnSocketAsync(DashScopeClientWebSocket socket)
     /// Rent or create a socket connection from pool.
     /// </summary>
     /// <param name="cancellationToken"></param>
-    /// <typeparam name="TOutput">The output type.</typeparam>
     /// <returns></returns>
-    public async Task<DashScopeClientWebSocketWrapper> RentSocketAsync<TOutput>(
-        CancellationToken cancellationToken = default)
-        where TOutput : class
+    public async Task<DashScopeClientWebSocketWrapper> RentSocketAsync(CancellationToken cancellationToken = default)
     {
         var found = false;
         DashScopeClientWebSocket? socket = null;
@@ -67,7 +85,7 @@ public async Task<DashScopeClientWebSocketWrapper> RentSocketAsync<TOutput>(
             }
             else
             {
-                socket = await InitializeNewSocketAsync<TOutput>(_options.BaseWebsocketAddress, cancellationToken);
+                socket = await InitializeNewSocketAsync(_options.WebsocketBaseAddress, cancellationToken);
                 found = true;
             }
         }
@@ -77,22 +95,21 @@ public async Task<DashScopeClientWebSocketWrapper> RentSocketAsync<TOutput>(
 
     private DashScopeClientWebSocketWrapper ActivateSocket(DashScopeClientWebSocket socket)
     {
-        _active.Add(socket);
+        _active.TryAdd(socket.Id, socket);
         return new DashScopeClientWebSocketWrapper(socket, this);
     }
 
-    private async Task<DashScopeClientWebSocket> InitializeNewSocketAsync<TOutput>(
+    private async Task<DashScopeClientWebSocket> InitializeNewSocketAsync(
         string url,
         CancellationToken cancellationToken = default)
-        where TOutput : class
     {
         if (_available.Count + _active.Count >= _options.SocketPoolSize)
         {
             throw new InvalidOperationException("[DashScopeSDK] Socket pool is full");
         }
 
-        var socket = new DashScopeClientWebSocket(_options.ApiKey, _options.WorkspaceId);
-        await socket.ConnectAsync<TOutput>(new Uri(url), cancellationToken);
+        var socket = _dashScopeClientWebSocketFactory.GetClientWebSocket(_options.ApiKey, _options.WorkspaceId);
+        await socket.ConnectAsync(new Uri(url), cancellationToken);
         return socket;
     }
 
@@ -107,10 +124,10 @@ private void Dispose(bool disposing)
                 socket?.Dispose();
             }
 
-            while (_active.IsEmpty == false)
+            var activeSockets = _active.Values;
+            foreach (var activeSocket in activeSockets)
             {
-                _active.TryTake(out var socket);
-                socket?.Dispose();
+                activeSocket.Dispose();
             }
         }
     }

src/Cnblogs.DashScope.Core/DashScopeClientWebSocketWrapper.cs

@@ -1,4 +1,6 @@
-namespace Cnblogs.DashScope.Core;
+using System.Text.Json;
+
+namespace Cnblogs.DashScope.Core;
 
 /// <summary>
 /// Represents a transient wrapper for rented websocket, should be transient.
@@ -14,15 +16,20 @@ public sealed record DashScopeClientWebSocketWrapper(DashScopeClientWebSocket So
     public IAsyncEnumerable<byte> BinaryOutput => Socket.BinaryOutput.ReadAllAsync();
 
     /// <summary>
-    /// The task that completes when received task-started event from server.
+    /// The json message output.
     /// </summary>
-    public Task TaskStarted => Socket.TaskStarted;
+    public IAsyncEnumerable<DashScopeWebSocketResponse<JsonElement>> JsonOutput => Socket.JsonOutput.ReadAllAsync();
 
     /// <summary>
     /// Reset task signal and output cannel.
     /// </summary>
     public void ResetTask() => Socket.ResetOutput();
 
+    /// <summary>
+    /// The task that completes when received task-started event from server.
+    /// </summary>
+    public Task TaskStarted => Socket.TaskStarted;
+
     /// <summary>
     /// Send message to server.
     /// </summary>
@@ -44,7 +51,7 @@ public Task SendMessageAsync<TInput, TParameter>(
     /// <inheritdoc />
     public void Dispose()
     {
-        Pool.ReturnSocketAsync(Socket);
+        Pool.ReturnSocket(Socket);
         GC.SuppressFinalize(this);
     }
 

src/Cnblogs.DashScope.Core/DashScopeOptions.cs

@@ -20,7 +20,7 @@ public class DashScopeOptions
     /// <summary>
     /// Base address for DashScope websocket API.
     /// </summary>
-    public string BaseWebsocketAddress { get; set; } = DashScopeDefaults.WebsocketApiBaseAddress;
+    public string WebsocketBaseAddress { get; set; } = DashScopeDefaults.WebsocketApiBaseAddress;
 
     /// <summary>
     /// Default workspace Id.

src/Cnblogs.DashScope.Core/DashScopeWebSocketResponse.cs

@@ -8,5 +8,4 @@
 /// <typeparam name="TOutput">Output type of the response.</typeparam>
 public record DashScopeWebSocketResponse<TOutput>(
     DashScopeWebSocketResponseHeader Header,
-    DashScopeWebSocketResponsePayload<TOutput> Payload)
-    where TOutput : class;
+    DashScopeWebSocketResponsePayload<TOutput> Payload);