From da5880ca7382c7ea676f19a00a73d1a1efa59f90 Mon Sep 17 00:00:00 2001
From: KirillKurdyukov <kurdyukov-kir@ydb.tech>
Date: Thu, 15 Aug 2024 11:53:41 +0300
Subject: [PATCH] fixes

---
 slo/src/Internal/SloContext.cs | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/slo/src/Internal/SloContext.cs b/slo/src/Internal/SloContext.cs
index 7766eaa..cdc4005 100644
--- a/slo/src/Internal/SloContext.cs
+++ b/slo/src/Internal/SloContext.cs
@@ -135,7 +135,7 @@ public async Task Run(RunConfig runConfig)
             { Window = TimeSpan.FromSeconds(1), PermitLimit = runConfig.ReadRps, QueueLimit = int.MaxValue });
 
         var cancellationTokenSource = new CancellationTokenSource();
-        cancellationTokenSource.CancelAfter(TimeSpan.FromSeconds(runConfig.ShutdownTime));
+        cancellationTokenSource.CancelAfter(TimeSpan.FromSeconds(runConfig.Time));
 
         var writeTask = ShootingTask(writeLimiter, "write", Upsert);
         var readTask = ShootingTask(readLimiter, "read", Select);
@@ -153,22 +153,18 @@ public async Task Run(RunConfig runConfig)
         Task ShootingTask(RateLimiter rateLimitPolicy, string shootingName,
             Func<T, RunConfig, Gauge?, Task<(int, StatusCode)>> action)
         {
+            // ReSharper disable once MethodSupportsCancellation
             return Task.Run(async () =>
             {
                 var tasks = new List<Task>();
 
-                long activeTasks = 0;
-
                 while (!cancellationTokenSource.Token.IsCancellationRequested)
                 {
                     using var lease = await rateLimitPolicy
                         .AcquireAsync(cancellationToken: cancellationTokenSource.Token);
 
-                    tasks.Add(Task.Run(async () =>
+                    _ = Task.Run(async () =>
                     {
-                        // ReSharper disable once AccessToModifiedClosure
-                        Interlocked.Increment(ref activeTasks);
-
                         var sw = Stopwatch.StartNew();
                         var (attempts, statusCode) = await action(client, runConfig, errorsGauge);
                         string label;
@@ -186,16 +182,16 @@ Task ShootingTask(RateLimiter rateLimitPolicy, string shootingName,
                             label = "ok";
                         }
 
-                        Interlocked.Decrement(ref activeTasks);
                         attemptsHistogram.WithLabels(label).Observe(attempts);
                         latencySummary.WithLabels(label).Observe(sw.ElapsedMilliseconds);
-                    }, cancellationTokenSource.Token));
+                    }, cancellationTokenSource.Token);
                 }
 
-                await Task.WhenAll(tasks);
+                // ReSharper disable once MethodSupportsCancellation
+                await Task.Delay(TimeSpan.FromSeconds(runConfig.ShutdownTime));
 
                 _logger.LogInformation("{ShootingName} shooting is stopped", shootingName);
-            }, cancellationTokenSource.Token);
+            });
         }
     }