Refactor ChunkUploaderWebSocket and add SendChunksAsync method

Etherna · Oct 5, 2024 · 289f8f1 · 289f8f1
1 parent 7095c48
commit 289f8f1
Showing 1 changed file with 29 additions and 4 deletions.
diff --git a/src/BeeNet.Core/Models/ChunkUploaderWebSocket.cs b/src/BeeNet.Core/Models/ChunkUploaderWebSocket.cs
@@ -19,21 +19,29 @@
 
 namespace Etherna.BeeNet.Models
 {
-    public sealed class ChunkUploaderWebSocket(
+    public class ChunkUploaderWebSocket(
         WebSocket webSocket)
         : IDisposable
     {
         // Fields.
         private readonly byte[] responseBuffer = new byte[SwarmHash.HashSize]; //not really used
 
         // Dispose.
+        protected virtual void Dispose(bool disposing)
+        {
+            if (disposing)
+            {
+                webSocket.Dispose();
+            }
+        }
         public void Dispose()
         {
-            webSocket.Dispose();
+            Dispose(true);
+            GC.SuppressFinalize(this);
         }
 
         // Methods.
-        public async Task CloseAsync()
+        public virtual async Task CloseAsync()
         {
             if (webSocket.State == WebSocketState.Open)
             {
@@ -49,7 +57,9 @@ await webSocket.CloseOutputAsync(
             }
         }
 
-        public async Task SendChunkAsync(SwarmChunk chunk, CancellationToken cancellationToken)
+        public virtual async Task SendChunkAsync(
+            SwarmChunk chunk,
+            CancellationToken cancellationToken)
         {
             ArgumentNullException.ThrowIfNull(chunk, nameof(chunk));
 
@@ -62,5 +72,20 @@ public async Task SendChunkAsync(SwarmChunk chunk, CancellationToken cancellatio
                 throw new OperationCanceledException(
                     $"Connection closed by server, message: {response.CloseStatusDescription}");
         }
+
+        public virtual async Task SendChunksAsync(
+            SwarmChunk[] chunks,
+            Action<int>? onChunkBatchSent = null,
+            CancellationToken cancellationToken = default)
+        {
+            ArgumentNullException.ThrowIfNull(chunks, nameof(chunks));
+
+            // Iterate on chunks.
+            for (int i = 0; i < chunks.Length; i++)
+            {
+                await SendChunkAsync(chunks[i], cancellationToken).ConfigureAwait(false);
+                onChunkBatchSent?.Invoke(i);
+            }
+        }
     }
 }