Adjust code changes to the main launcher

Author: neon-nyan

SharedStatic.V1Ext_Update4.cs

@@ -13,8 +13,9 @@ public partial class SharedStaticV1Ext
     /// Registers the function of the Speed Throttler Service from the Main Application.
     /// To disable the speed throttler service, set the <paramref name="addBytesOrWaitAsyncCallback"/> to <see cref="nint.Zero"/>.
     /// </summary>
-    /// <param name="addBytesOrWaitAsyncCallback">The address of the speed throttler service callback</param>
-    internal delegate HResult RegisterSpeedThrottlerServiceDelegate(nint addBytesOrWaitAsyncCallback);
+    /// <param name="addBytesOrWaitAsyncCallback">The address of the adder wait callback</param>
+    /// <param name="addBytesOrWaitAsyncCallback">The address of the current speed getter callback</param>
+    internal delegate HResult RegisterSpeedThrottlerServiceDelegate(nint addBytesOrWaitAsyncCallback, nint getSharedThrottleBytesCallback);
 }
 
 public partial class SharedStaticV1Ext<T>
@@ -37,53 +38,61 @@ private static void InitExtension_Update4Exports()
     /// <summary>
     /// This method is an ABI proxy and installer for the Speed Throttler Service functionality. To use Speed Throttler Service functionalities, Use <see cref="SpeedLimiterService"/> instead.
     /// </summary>
-    private static unsafe HResult RegisterSpeedThrottlerService(nint addBytesOrWaitAsyncCallback)
+    private static unsafe HResult RegisterSpeedThrottlerService(nint addBytesOrWaitAsyncCallback, nint getSharedThrottleBytesCallback)
     {
-        SpeedLimiterService.AddBytesOrWaitAsyncCallback = (delegate* unmanaged[Stdcall]<nint, long, nint, out nint, int>)addBytesOrWaitAsyncCallback;
-        if (addBytesOrWaitAsyncCallback == nint.Zero)
+        if (addBytesOrWaitAsyncCallback == nint.Zero && getSharedThrottleBytesCallback == nint.Zero)
         {
             SpeedLimiterService.AddBytesOrWaitAsyncCallback = null;
+            SpeedLimiterService.GetSharedThrottleBytesCallback = null;
             InstanceLogger.LogTrace("[RegisterSpeedThrottlerService] Speed Throttler Service has been uninstalled");
 
             return HResult.Ok;
         }
 
-        // Test the delegate first before registering it.
-        nint contextP = SpeedLimiterService.CreateServiceContext();
-        try
+        if (addBytesOrWaitAsyncCallback != nint.Zero && getSharedThrottleBytesCallback != nint.Zero)
         {
-            // Try call the increment function with 16 bytes of load, then check for exception
+            SpeedLimiterService.AddBytesOrWaitAsyncCallback = (delegate* unmanaged[Cdecl]<nint, long, nint, out nint, int>)addBytesOrWaitAsyncCallback;
+
+            // Test the delegate first before registering it.
+            nint contextP = SpeedLimiterService.CreateServiceContext();
+            try
+            {
+                // Try call the increment function with 16 bytes of load, then check for exception
 #pragma warning disable CA2012
-            ValueTask task = SpeedLimiterService.AddBytesOrWaitAsync(contextP, 16);
+                ValueTask task = SpeedLimiterService.AddBytesOrWaitAsync(contextP, 16);
 #pragma warning restore CA2012
-            if (task is { IsCompleted: true, IsFaulted: false } ||
-                task.IsCompletedSuccessfully)
-            {
-                LogSuccess();
-                return HResult.Ok;
-            }
+                if (task is { IsCompleted: true, IsFaulted: false } ||
+                    task.IsCompletedSuccessfully)
+                {
+                    LogSuccess();
+                    return HResult.Ok;
+                }
 
-            // Try block and await if task is still going.
-            task.GetAwaiter().GetResult();
+                // Try block and await if task is still going.
+                task.GetAwaiter().GetResult();
 
-            // If nothing blown up, return OK.
-            LogSuccess();
-            return HResult.Ok;
+                // If nothing blown up, return OK.
+                LogSuccess();
+                return HResult.Ok;
 
-            void LogSuccess()
+                void LogSuccess()
+                {
+                    InstanceLogger.LogTrace("[RegisterSpeedThrottlerService] Speed Throttler Service has been installed. Service's callback is located at address: 0x{Ptr:x8}", addBytesOrWaitAsyncCallback);
+                }
+            }
+            catch (Exception ex)
             {
-                InstanceLogger.LogTrace("[RegisterSpeedThrottlerService] Speed Throttler Service has been installed. Service's callback is located at address: 0x{Ptr:x8}", addBytesOrWaitAsyncCallback);
+                SpeedLimiterService.AddBytesOrWaitAsyncCallback = null; // Reset the callback
+                return Marshal.GetHRForException(ex);
+            }
+            finally
+            {
+                SpeedLimiterService.FreeServiceContext(contextP);
             }
         }
-        catch (Exception ex)
-        {
-            SpeedLimiterService.AddBytesOrWaitAsyncCallback = null; // Reset the callback
-            return Marshal.GetHRForException(ex);
-        }
-        finally
-        {
-            SpeedLimiterService.FreeServiceContext(contextP);
-        }
+
+        InstanceLogger.LogError("[RegisterSpeedThrottlerService] Failed to install/uninstall Speed Throttler Service. You must provide both arguments either all null or not-null!");
+        return 0x80070057; // ERROR_INVALID_PARAMETER
     }
     #endregion
 }

Utility/RetryableCopyToStreamTask.cs

@@ -75,7 +75,9 @@ public async ValueTask DisposeAsync()
         GC.SuppressFinalize(this);
     }
 
-    private RetryableCopyToStreamTask(SourceStreamFactory sourceStreamFactory, Stream targetStream, RetryableCopyToStreamTaskOptions options)
+    private RetryableCopyToStreamTask(SourceStreamFactory              sourceStreamFactory,
+                                      Stream                           targetStream,
+                                      RetryableCopyToStreamTaskOptions options)
     {
         _sourceStreamFactory = sourceStreamFactory;
         _targetStream        = targetStream;
@@ -135,7 +137,9 @@ private async ValueTask WriteTaskCore(ReadDelegate? readDelegate, byte[] buffer,
 
         while (retryAttemptLeft > 0)
         {
-            var (timedOutCts, coopCts) = RenewTimeOutCancelToken(in timeoutSpan, in token);
+            (CancellationTokenSource timedOutCts, CancellationTokenSource coopCts) =
+                RenewTimeOutCancelToken(in timeoutSpan, in token);
+
             _sourceStream = await _sourceStreamFactory(lastBytesPosition, coopCts.Token);
             if (_sourceStream == null)
             {
@@ -150,8 +154,13 @@ private async ValueTask WriteTaskCore(ReadDelegate? readDelegate, byte[] buffer,
                     .ReadAsync(buffer, coopCts.Token)
                     .ConfigureAwait(false)) > 0)
                 {
+                    await SpeedLimiterService.AddBytesOrWaitAsync(_options.SpeedLimiterServiceContext,
+                                                                  read,
+                                                                  coopCts.Token);
+
                     lastBytesPosition += read;
                     readDelegate?.Invoke(read);
+
                     await _targetStream
                         .WriteAsync(buffer.AsMemory(0, read), coopCts.Token)
                         .ConfigureAwait(false);
@@ -213,7 +222,7 @@ private static bool IsRetryableHttpStatusCode(HttpStatusCode? statusCode)
     private static (CancellationTokenSource TimedOutCts, CancellationTokenSource CoopCts)
         RenewTimeOutCancelToken(in TimeSpan timeoutSpan, in CancellationToken innerToken)
     {
-        CancellationTokenSource timedOutCts = new CancellationTokenSource(timeoutSpan);
+        CancellationTokenSource timedOutCts = new(timeoutSpan);
         CancellationTokenSource coopCts     = CancellationTokenSource.CreateLinkedTokenSource(timedOutCts.Token, innerToken);
 
         return (timedOutCts, coopCts);

Utility/RetryableCopyToStreamTaskOptions.cs

@@ -47,7 +47,15 @@ public class RetryableCopyToStreamTaskOptions
     /// Whether to dispose the target <see cref="Stream"/> once <see cref="IDisposable.Dispose()"/> is being called.
     /// </summary>
     /// <remarks>
-    /// Default: false
+    /// Default: <see langword="false"/>
     /// </remarks>
     public bool IsDisposeTargetStream { get; init; }
+
+    /// <summary>
+    /// The context of speed limiter service to be used for download speed throttling.
+    /// </summary>
+    /// <remarks>
+    /// Default: <see cref="nint.Zero"/>
+    /// </remarks>
+    public nint SpeedLimiterServiceContext { get; init; }
 }

Utility/SpeedLimiterService.cs

@@ -1,9 +1,10 @@
-using System.Runtime.CompilerServices;
+using Microsoft.Win32.SafeHandles;
+using System;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
-using Hi3Helper.Plugin.Core.Utility.Windows;
-using Microsoft.Win32.SafeHandles;
+using System.Threading.Tasks.Sources;
 
 namespace Hi3Helper.Plugin.Core.Utility;
 
@@ -35,15 +36,46 @@ namespace Hi3Helper.Plugin.Core.Utility;
 /// </remarks>
 public static class SpeedLimiterService
 {
-    internal static unsafe delegate* unmanaged[Stdcall]<nint, long, nint, out nint, int> AddBytesOrWaitAsyncCallback =
+    internal static unsafe delegate* unmanaged[Cdecl]<nint, long, nint, out nint, int> AddBytesOrWaitAsyncCallback =
+        null;
+
+    internal static unsafe delegate* unmanaged[Cdecl]<ref long, ref long, void> GetSharedThrottleBytesCallback =
         null;
 
     /// <summary>
     /// Creates a context to be used for the speed limiter service. This context can be used into multiple instances or threads of your downloader.
     /// </summary>
     /// <returns></returns>
     public static unsafe nint CreateServiceContext()
-        => (nint)Mem.Alloc<long>(2); // Context struct is 16 bytes in size.
+    {
+        ThrottleServiceContext* alloc = Mem.Alloc<ThrottleServiceContext>();
+        alloc->AvailableTokens = 0;
+        alloc->LastTimestamp   = Environment.TickCount64;
+
+        if (GetSharedThrottleBytesCallback == null)
+            return (nint)alloc;
+
+        try
+        {
+            long bytesPerSecond = 0;
+            long burstBytes     = 0;
+
+            GetSharedThrottleBytesCallback(ref bytesPerSecond, ref burstBytes);
+
+            if (bytesPerSecond < burstBytes)
+            {
+                bytesPerSecond = burstBytes;
+            }
+
+            alloc->AvailableTokens = bytesPerSecond;
+        }
+        catch
+        {
+            // ignored
+        }
+
+        return (nint)alloc;
+    }
 
     /// <summary>
     /// Free the speed limiter service context.
@@ -65,49 +97,126 @@ public static unsafe ValueTask AddBytesOrWaitAsync(
         long              readBytes,
         CancellationToken token = default)
     {
-        if (AddBytesOrWaitAsyncCallback == null)
-        {
+        if (context == nint.Zero || AddBytesOrWaitAsyncCallback == null)
             return ValueTask.CompletedTask;
-        }
 
-        nint tokenHandle = token.WaitHandle.SafeWaitHandle.DangerousGetHandle();
         int hr = AddBytesOrWaitAsyncCallback(context,
-                                         readBytes,
-                                         tokenHandle,
-                                         out nint asyncWaitHandle);
+                                             readBytes,
+                                             nint.Zero,
+                                             out nint completionHandle);
 
-        AsyncValueTaskMethodBuilder valueTaskCs = new();
-        if (Marshal.GetExceptionForHR(hr) is { } exception)
-        {
-            valueTaskCs.SetException(exception);
-            return valueTaskCs.Task;
-        }
+        if (Marshal.GetExceptionForHR(hr) is { } ex)
+            return ValueTask.FromException(ex);
+
+        NativeThrottleOperation op = new();
+        op.Initialize(completionHandle, token);
 
-        SafeWaitHandle safeHandle = new(asyncWaitHandle, false);
-        WaitHandle waitHandle = new EventWaitHandle(false, EventResetMode.ManualReset)
+        return op.AsValueTask();
+    }
+
+    [StructLayout(LayoutKind.Sequential, Pack = 8)] // Pack to 8 bytes to ensure aligning
+    private struct ThrottleServiceContext
+    {
+        public long AvailableTokens;
+        public long LastTimestamp;
+    }
+
+    private sealed class NativeThrottleOperation : IValueTaskSource
+    {
+        private ManualResetValueTaskSourceCore<bool> _core = new()
         {
-            SafeWaitHandle = safeHandle
+            RunContinuationsAsynchronously = true
         };
 
-        ThreadPool.UnsafeRegisterWaitForSingleObject(waitHandle,
-            DisposeWaitHandleCallback,
-            null,
-            -1,
-            true);
+        private int _isCompleted;
+        private EventWaitHandle? _completionWait;
+        private SafeWaitHandle? _completionSafe;
+        private RegisteredWaitHandle? _registeredWait;
+        private CancellationTokenRegistration _ctr;
 
-        return valueTaskCs.Task;
+        public ValueTask AsValueTask()
+            => new(this, _core.Version);
 
-        void DisposeWaitHandleCallback(object? state, bool isTimedOut)
+        public void Initialize(
+            nint completionHandle,
+            CancellationToken token)
         {
-            safeHandle.Dispose();
-            waitHandle.Dispose();
+            _completionSafe = new SafeWaitHandle(completionHandle, true);
+            _completionWait = new EventWaitHandle(false, EventResetMode.ManualReset)
+            {
+                SafeWaitHandle = _completionSafe
+            };
+
+            _registeredWait =
+                ThreadPool.RegisterWaitForSingleObject(_completionWait,
+                                                       OnWaitSingleCompleted,
+                                                       this,
+                                                       -1,
+                                                       true);
+
+            if (token.CanBeCanceled)
+            {
+                _ctr = token.Register(OnCancellationRequested, this);
+            }
+        }
+
+        private static void OnWaitSingleCompleted(object? state, bool isTimedOut)
+        {
+            NativeThrottleOperation op = (NativeThrottleOperation)state!;
+            op.Complete();
+        }
 
-            if (asyncWaitHandle != nint.Zero)
+        private static void OnCancellationRequested(object? state)
+        {
+            NativeThrottleOperation op = (NativeThrottleOperation)state!;
+            op.Cancel();
+        }
+
+        private void Complete()
+        {
+            if (Interlocked.Exchange(ref _isCompleted, 1) == 1)
             {
-                _ = PInvoke.CloseHandle(asyncWaitHandle);
+                return;
             }
 
-            valueTaskCs.SetResult();
+            Cleanup();
+            _core.SetResult(true);
         }
+
+        private void Cancel()
+        {
+            if (Interlocked.Exchange(ref _isCompleted, 1) == 1)
+            {
+                return;
+            }
+
+            Cleanup();
+            _core.SetException(new OperationCanceledException());
+        }
+
+        private void Cleanup()
+        {
+            _registeredWait?.Unregister(null);
+            _registeredWait = null;
+
+            _ctr.Dispose();
+
+            _completionWait?.Dispose();
+            _completionWait = null;
+            _completionSafe = null;
+        }
+
+        public void GetResult(short token)
+            => _core.GetResult(token);
+
+        public ValueTaskSourceStatus GetStatus(short token)
+            => _core.GetStatus(token);
+
+        public void OnCompleted(
+            Action<object?> continuation,
+            object? state,
+            short token,
+            ValueTaskSourceOnCompletedFlags flags)
+            => _core.OnCompleted(continuation, state, token, flags);
     }
 }