Put two Task.Yield calls into place to make sure the receive loop in …

…WebSocketGremlinqClient won't get stuck.
Gremlinq · Jan 9, 2025 · 1992ff2 · 1992ff2
1 parent 6c57f81
commit 1992ff2
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/src/Providers.Core/Factory/WebSocketGremlinqClientFactory.cs b/src/Providers.Core/Factory/WebSocketGremlinqClientFactory.cs
@@ -128,12 +128,21 @@ public async IAsyncEnumerator<ResponseMessage<T>> GetAsyncEnumerator(Cancellatio
                             if (await new ValueTask<ResponseAndQueueUnion<T>?>(this, 0).ConfigureAwait(false) is { } union)
                             {
                                 if (union.TryGetResponse(out var response))
+                                {
+                                    // Since the below yield return is what effectively yields control back to user code,
+                                    // the receive loop may stall if user code blocks. Although technically not Gremlinq's
+                                    // fault, we take this measure.
+                                    await Task.Yield();
+
                                     yield return response;
+                                }
                                 else if (union.TryGetQueue(out var semaphore, out var queue))
                                 {
                                     while (true)
                                     {
-                                        await semaphore.WaitAsync(ct).ConfigureAwait(false);
+                                        await semaphore
+                                            .WaitAsync(ct)
+                                            .ConfigureAwait(false);
 
                                         if (queue.TryDequeue(out var queuedResponse))
                                         {
@@ -155,6 +164,9 @@ await _client
                                             }
                                             else
                                             {
+                                                // See above.
+                                                await Task.Yield();
+
                                                 yield return queuedResponse;
 
                                                 if (queuedResponse.Status.Code != PartialContent)