improve downloading rate limiter

Scighost · Aug 5, 2024 · bf071c4 · bf071c4
1 parent 96d9493
commit bf071c4
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 41 deletions.
diff --git a/src/Starward/Controls/InstallGameController.xaml.cs b/src/Starward/Controls/InstallGameController.xaml.cs
@@ -134,7 +134,6 @@ private void UpdateSpeedState()
         long totalBytes = 0;
         long finishedBytes = 0;
         bool determinate = false;
-        _installGameManager.UpdateSpeedState();
         foreach (var model in InstallServices)
         {
             model.UpdateState();

diff --git a/src/Starward/Pages/Setting/DownloadSettingPage.xaml.cs b/src/Starward/Pages/Setting/DownloadSettingPage.xaml.cs
@@ -69,7 +69,7 @@ partial void OnDefaultInstallPathChanged(string? value)
     private int speedLimit = AppConfig.SpeedLimitKBPerSecond;
     partial void OnSpeedLimitChanged(int value)
     {
-        InstallGameManager.SpeedLimitBytesPerSecond = value == 0 ? long.MaxValue : value * 1024;
+        InstallGameManager.SetRateLimit(value * 1024);
         AppConfig.SpeedLimitKBPerSecond = value;
     }
 

diff --git a/src/Starward/Services/Download/InstallGameManager.cs b/src/Starward/Services/Download/InstallGameManager.cs
@@ -4,9 +4,8 @@
 using Starward.Messages;
 using System;
 using System.Collections.Concurrent;
-using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
-using System.Threading;
+using System.Threading.RateLimiting;
 
 namespace Starward.Services.Download;
 
@@ -20,8 +19,7 @@ internal class InstallGameManager
     private InstallGameManager()
     {
         _services = new();
-        long speed = AppConfig.SpeedLimitKBPerSecond * 1024;
-        SpeedLimitBytesPerSecond = speed == 0 ? long.MaxValue : speed;
+        SetRateLimit(AppConfig.SpeedLimitKBPerSecond * 1024);
     }
 
 
@@ -30,29 +28,34 @@ private InstallGameManager()
 
 
 
-    public static long DownloadBytesInSecond;
+    public static TokenBucketRateLimiter RateLimiter { get; private set; }
 
 
-    public static long SpeedLimitBytesPerSecond { get; set; }
 
-
-    public static bool IsExceedSpeedLimit => Interlocked.Read(ref DownloadBytesInSecond) >= SpeedLimitBytesPerSecond;
-
-
-    private long _lastTimeStamp;
-
-
-    public void UpdateSpeedState()
+    public static void SetRateLimit(int bytesPerSecond)
     {
-        long ts = Stopwatch.GetTimestamp();
-        if (ts - _lastTimeStamp >= Stopwatch.Frequency)
+        if (bytesPerSecond <= 0)
         {
-            DownloadBytesInSecond = 0;
+            bytesPerSecond = int.MaxValue;
         }
+        else if (bytesPerSecond < (1 << 14))
+        {
+            bytesPerSecond = 1 << 14;
+        }
+        RateLimiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions
+        {
+            TokenLimit = bytesPerSecond,
+            ReplenishmentPeriod = TimeSpan.FromSeconds(1),
+            TokensPerPeriod = bytesPerSecond,
+            QueueProcessingOrder = QueueProcessingOrder.OldestFirst,
+            AutoReplenishment = true
+        });
     }
 
 
 
+
+
     public event EventHandler<InstallGameStateModel> InstallTaskAdded;
 
 

diff --git a/src/Starward/Services/Download/InstallGameService.cs b/src/Starward/Services/Download/InstallGameService.cs
@@ -21,6 +21,7 @@
 using System.Text;
 using System.Text.Json.Nodes;
 using System.Threading;
+using System.Threading.RateLimiting;
 using System.Threading.Tasks;
 using Vanara.PInvoke;
 
@@ -245,7 +246,7 @@ await Task.Run(() =>
             {
                 File.SetAttributes(file, FileAttributes.Normal);
             }
-        });
+        }).ConfigureAwait(false);
         CurrentGameBiz = gameBiz;
         _installPath = installPath;
         _initialized = true;
@@ -714,7 +715,7 @@ protected void StartTask(InstallGameState state)
         _cancellationTokenSource = new CancellationTokenSource();
         for (int i = 0; i < Environment.ProcessorCount; i++)
         {
-            _ = ExecuteTaskItemAsync(_cancellationTokenSource.Token);
+            _ = ExecuteTaskItemAsync(_cancellationTokenSource.Token).ConfigureAwait(false);
         }
     }
 
@@ -981,15 +982,15 @@ protected async Task CleanGameDeprecatedFilesAsync()
         {
             File.Delete(file);
         }
-        foreach (var file in await _hoYoPlayService.GetGameDeprecatedFilesAsync(CurrentGameBiz))
+        foreach (var file in await _hoYoPlayService.GetGameDeprecatedFilesAsync(CurrentGameBiz).ConfigureAwait(false))
         {
             var path = Path.Combine(_installPath, file.Name);
             if (File.Exists(path))
             {
                 File.Delete(path);
             }
         }
-        await WriteConfigFileAsync();
+        await WriteConfigFileAsync().ConfigureAwait(false);
         CurrentTaskFinished();
     }
 
@@ -1118,19 +1119,19 @@ protected async Task ExecuteTaskItemAsync(CancellationToken cancellationToken =
                         case InstallGameItemType.None:
                             break;
                         case InstallGameItemType.Download:
-                            await DownloadItemAsync(item, cancellationToken);
+                            await DownloadItemAsync(item, cancellationToken).ConfigureAwait(false);
                             Interlocked.Increment(ref _finishCount);
                             break;
                         case InstallGameItemType.Verify:
-                            await VerifyItemAsync(item, cancellationToken);
+                            await VerifyItemAsync(item, cancellationToken).ConfigureAwait(false);
                             Interlocked.Increment(ref _finishCount);
                             break;
                         case InstallGameItemType.Decompress:
-                            await DecompressItemAsync(item, cancellationToken);
+                            await DecompressItemAsync(item, cancellationToken).ConfigureAwait(false);
                             Interlocked.Increment(ref _finishCount);
                             break;
                         case InstallGameItemType.HardLink:
-                            await HardLinkItemAsync(item, cancellationToken);
+                            await HardLinkItemAsync(item, cancellationToken).ConfigureAwait(false);
                             Interlocked.Increment(ref _finishCount);
                             break;
                         default:
@@ -1141,7 +1142,7 @@ protected async Task ExecuteTaskItemAsync(CancellationToken cancellationToken =
                 {
                     // network error
                     _installItemQueue.Enqueue(item);
-                    await Task.Delay(1000, CancellationToken.None);
+                    await Task.Delay(1000, CancellationToken.None).ConfigureAwait(false);
                 }
                 catch (Exception ex) when (ex is OperationCanceledException or TaskCanceledException)
                 {
@@ -1171,7 +1172,7 @@ protected async Task ExecuteTaskItemAsync(CancellationToken cancellationToken =
 
     protected async Task DownloadItemAsync(InstallGameItem item, CancellationToken cancellationToken = default)
     {
-        const int BUFFER_SIZE = 1 << 14;
+        const int BUFFER_SIZE = 1 << 10;
         string file = item.Path;
         string file_tmp = item.Path + "_tmp";
         string file_target;
@@ -1203,14 +1204,14 @@ protected async Task DownloadItemAsync(InstallGameItem item, CancellationToken c
                 int length;
                 while ((length = await hs.ReadAsync(buffer, cancellationToken).ConfigureAwait(false)) != 0)
                 {
-                    await fs.WriteAsync(buffer.AsMemory(0, length), cancellationToken).ConfigureAwait(false);
-                    Interlocked.Add(ref _finishBytes, length);
-                    Interlocked.Add(ref InstallGameManager.DownloadBytesInSecond, length);
-                    if (InstallGameManager.IsExceedSpeedLimit)
+                    RateLimitLease lease = await InstallGameManager.RateLimiter.AcquireAsync(length, cancellationToken).ConfigureAwait(false);
+                    while (!lease.IsAcquired)
                     {
-                        long t = Stopwatch.GetTimestamp() / (Stopwatch.Frequency / 1000) % 1000;
-                        await Task.Delay((int)(1000 - t), cancellationToken);
+                        await Task.Delay(1, cancellationToken).ConfigureAwait(false);
+                        lease = await InstallGameManager.RateLimiter.AcquireAsync(length, cancellationToken).ConfigureAwait(false);
                     }
+                    await fs.WriteAsync(buffer.AsMemory(0, length), cancellationToken).ConfigureAwait(false);
+                    Interlocked.Add(ref _finishBytes, length);
                 }
             }
         }
@@ -1246,7 +1247,7 @@ protected async Task VerifyItemAsync(InstallGameItem item, CancellationToken can
         var buffer = ArrayPool<byte>.Shared.Rent(BUFFER_SIZE);
         try
         {
-            await _verifyGlobalSemaphore.WaitAsync(cancellationToken);
+            await _verifyGlobalSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
             using var fs = File.OpenRead(file_target);
             if (fs.Length != item.Size)
             {
@@ -1293,7 +1294,7 @@ protected async Task DecompressItemAsync(InstallGameItem item, CancellationToken
     {
         try
         {
-            await _decompressSemaphore.WaitAsync();
+            await _decompressSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
             using var fs = new FileSliceStream(item.DecompressPackageFiles);
             if (item.DecompressPackageFiles[0].Contains(".7z", StringComparison.CurrentCultureIgnoreCase))
             {
@@ -1310,7 +1311,7 @@ await Task.Run(() =>
                     };
                     extra.Extract(item.DecompressPath, true);
                     _finishBytes += fs.Length - sum;
-                }).ConfigureAwait(false);
+                }, cancellationToken).ConfigureAwait(false);
             }
             else
             {
@@ -1335,8 +1336,8 @@ await Task.Run(async () =>
                         }
                     }
                     _finishBytes += fs.Length - sum;
-                    await ApplyDiffFilesAsync(item.DecompressPath);
-                }).ConfigureAwait(false);
+                    await ApplyDiffFilesAsync(item.DecompressPath).ConfigureAwait(false);
+                }, cancellationToken).ConfigureAwait(false);
             }
             fs.Dispose();
             foreach (var file in item.DecompressPackageFiles)
@@ -1440,7 +1441,7 @@ protected async Task HardLinkItemAsync(InstallGameItem item, CancellationToken c
         var buffer = ArrayPool<byte>.Shared.Rent(BUFFER_SIZE);
         try
         {
-            await _verifyGlobalSemaphore.WaitAsync(cancellationToken);
+            await _verifyGlobalSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
             using var fs = File.OpenRead(file_source);
             if (fs.Length != item.Size)
             {

diff --git a/src/Starward/Starward.csproj b/src/Starward/Starward.csproj
@@ -43,6 +43,7 @@
 		<PackageReference Include="Serilog.Sinks.File" Version="6.0.0" />
 		<PackageReference Include="Starward.Assets" Version="0.4.13" />
 		<PackageReference Include="Starward.NativeLib" Version="0.2.1" />
+		<PackageReference Include="System.Threading.RateLimiting" Version="8.0.0" />
 		<PackageReference Include="Vanara.PInvoke.ComCtl32" Version="4.0.2" />
 		<PackageReference Include="Vanara.PInvoke.DwmApi" Version="4.0.2" />
 		<PackageReference Include="Vanara.PInvoke.Ole" Version="4.0.2" />