Better usage of Utils in PushQueue

aneoconsulting · Jul 11, 2023 · 49b2671 · 49b2671
1 parent 4c8418e
commit 49b2671
Showing 1 changed file with 30 additions and 54 deletions.
diff --git a/Adaptors/Amqp/src/PushQueueStorage.cs b/Adaptors/Amqp/src/PushQueueStorage.cs
@@ -27,35 +27,12 @@
 
 using ArmoniK.Core.Base;
 using ArmoniK.Core.Base.DataStructures;
+using ArmoniK.Utils;
 
 using Microsoft.Extensions.Logging;
-using Microsoft.Extensions.ObjectPool;
 
 namespace ArmoniK.Core.Adapters.Amqp;
 
-/// <summary>
-///   Policy for creating a <see cref="Session" /> for the <see cref="ObjectPool{Session}" />
-/// </summary>
-internal sealed class SessionPooledObjectPolicy : IPooledObjectPolicy<Session>
-{
-  private readonly IConnectionAmqp connectionAmqp_;
-
-  /// <summary>
-  ///   Initializes a <see cref="SessionPooledObjectPolicy" />
-  /// </summary>
-  /// <param name="connectionAmqp">AMQP connection that will be used to create new sessions</param>
-  public SessionPooledObjectPolicy(IConnectionAmqp connectionAmqp)
-    => connectionAmqp_ = connectionAmqp;
-
-  /// <inheritdoc />
-  public Session Create()
-    => new(connectionAmqp_.Connection);
-
-  /// <inheritdoc />
-  public bool Return(Session obj)
-    => !obj.IsClosed;
-}
-
 public class PushQueueStorage : QueueStorage, IPushQueueStorage
 {
   private const int MaxInternalQueuePriority = 10;
@@ -70,8 +47,10 @@ public PushQueueStorage(QueueCommon.Amqp          options,
     : base(options,
            connectionAmqp)
   {
-    logger_      = logger;
-    sessionPool_ = new DefaultObjectPool<Session>(new SessionPooledObjectPolicy(ConnectionAmqp));
+    logger_ = logger;
+    sessionPool_ = new ObjectPool<Session>(200,
+                                           () => new Session(connectionAmqp.Connection),
+                                           session => !session.IsClosed);
   }
 
   /// <inheritdoc />
@@ -80,11 +59,12 @@ public async Task PushMessagesAsync(IEnumerable<MessageData> messages,
                                       CancellationToken        cancellationToken = default)
   {
     var priorityGroups = messages.GroupBy(msgData => msgData.Options.Priority);
-    await Task.WhenAll(priorityGroups.Select(group => PushMessagesAsync(group,
-                                                                        partitionId,
-                                                                        group.Key,
-                                                                        cancellationToken)))
-              .ConfigureAwait(false);
+    await priorityGroups.ParallelForEach(new ParallelTaskOptions(cancellationToken),
+                                         group => PushMessagesAsync(group,
+                                                                    partitionId,
+                                                                    group.Key,
+                                                                    cancellationToken))
+                        .ConfigureAwait(false);
   }
 
   private async Task PushMessagesAsync(IEnumerable<MessageData> messages,
@@ -114,29 +94,25 @@ private async Task PushMessagesAsync(IEnumerable<MessageData> messages,
                      whichQueue,
                      internalPriority);
 
-    var session = sessionPool_.Get();
-    try
-    {
-      var sender = new SenderLink(session,
-                                  $"{partitionId}###SenderLink{whichQueue}",
-                                  $"{partitionId}###q{whichQueue}");
-
-      await Task.WhenAll(messages.Select(msgData => sender.SendAsync(new Message(Encoding.UTF8.GetBytes(msgData.TaskId))
-                                                                     {
-                                                                       Header = new Header
-                                                                                {
-                                                                                  Priority = (byte)internalPriority,
-                                                                                },
-                                                                       Properties = new Properties(),
-                                                                     })))
-                .ConfigureAwait(false);
-
-      await sender.CloseAsync()
+    await using var session = await sessionPool_.GetAsync(cancellationToken)
+                                                .ConfigureAwait(false);
+
+    var sender = new SenderLink(session,
+                                $"{partitionId}###SenderLink{whichQueue}",
+                                $"{partitionId}###q{whichQueue}");
+
+    await messages.ParallelForEach(new ParallelTaskOptions(cancellationToken),
+                                   msgData => sender.SendAsync(new Message(Encoding.UTF8.GetBytes(msgData.TaskId))
+                                                               {
+                                                                 Header = new Header
+                                                                          {
+                                                                            Priority = (byte)internalPriority,
+                                                                          },
+                                                                 Properties = new Properties(),
+                                                               }))
                   .ConfigureAwait(false);
-    }
-    finally
-    {
-      sessionPool_.Return(session);
-    }
+
+    await sender.CloseAsync()
+                .ConfigureAwait(false);
   }
 }