From 77d9f4d174df52edc15d19e6a720557a5ae24d07 Mon Sep 17 00:00:00 2001
From: Alex Li <github@alexv525.com>
Date: Thu, 20 Jun 2024 09:08:14 +0800
Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Refactor=20using=20`WeakRe?=
 =?UTF-8?q?ference`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 dio/lib/src/adapters/io_adapter.dart          | 45 +++++--------------
 .../http2_adapter/lib/src/http2_adapter.dart  | 43 +++++-------------
 2 files changed, 20 insertions(+), 68 deletions(-)

diff --git a/dio/lib/src/adapters/io_adapter.dart b/dio/lib/src/adapters/io_adapter.dart
index 62837e140..cfdd26ae2 100644
--- a/dio/lib/src/adapters/io_adapter.dart
+++ b/dio/lib/src/adapters/io_adapter.dart
@@ -59,10 +59,6 @@ class IOHttpClientAdapter implements HttpClientAdapter {
   HttpClient? _cachedHttpClient;
   bool _closed = false;
 
-  final cancelFutureOperationPool =
-      <Future<void>, Set<CancelableOperation<ResponseBody>>>{};
-  final cancelFutureRequestPool = <Future<void>, Set<HttpClientRequest>>{};
-
   @override
   Future<ResponseBody> fetch(
     RequestOptions options,
@@ -77,24 +73,11 @@ class IOHttpClientAdapter implements HttpClientAdapter {
     final operation = CancelableOperation.fromFuture(
       _fetch(options, requestStream, cancelFuture),
     );
-
-    if (cancelFuture != null) {
-      cancelFutureOperationPool.putIfAbsent(cancelFuture, () => {});
-      cancelFutureRequestPool.putIfAbsent(cancelFuture, () => {});
-
-      cancelFuture.whenComplete(() {
-        cancelFutureOperationPool[cancelFuture]?.forEach((e) => e.cancel());
-        cancelFutureOperationPool.remove(cancelFuture);
-        cancelFutureRequestPool[cancelFuture]?.forEach((e) => e.abort());
-        cancelFutureRequestPool.remove(cancelFuture);
-      });
-
-      cancelFutureOperationPool[cancelFuture]!.add(operation);
-    }
-
-    return operation.value.whenComplete(
-      () => cancelFutureOperationPool[cancelFuture]?.remove(operation),
-    );
+    final wr = WeakReference<CancelableOperation<ResponseBody>>(operation);
+    cancelFuture?.whenComplete(() {
+      wr.target?.cancel();
+    });
+    return operation.value;
   }
 
   Future<ResponseBody> _fetch(
@@ -121,9 +104,10 @@ class IOHttpClientAdapter implements HttpClientAdapter {
         request = await reqFuture;
       }
 
-      if (cancelFuture != null) {
-        cancelFutureRequestPool[cancelFuture]!.add(request);
-      }
+      final wr = WeakReference<HttpClientRequest>(request);
+      cancelFuture?.whenComplete(() {
+        wr.target?.abort();
+      });
 
       // Set Headers
       options.headers.forEach((key, value) {
@@ -164,13 +148,6 @@ class IOHttpClientAdapter implements HttpClientAdapter {
     request.maxRedirects = options.maxRedirects;
     request.persistentConnection = options.persistentConnection;
 
-    void removeRequestInPool() {
-      if (cancelFuture == null) {
-        return;
-      }
-      cancelFutureRequestPool[cancelFuture]?.remove(request);
-    }
-
     if (requestStream != null) {
       // Transform the request data.
       Future<dynamic> future = request.addStream(requestStream);
@@ -180,7 +157,6 @@ class IOHttpClientAdapter implements HttpClientAdapter {
           sendTimeout,
           onTimeout: () {
             request.abort();
-            removeRequestInPool();
             throw DioException.sendTimeout(
               timeout: sendTimeout,
               requestOptions: options,
@@ -198,7 +174,6 @@ class IOHttpClientAdapter implements HttpClientAdapter {
         receiveTimeout,
         onTimeout: () {
           request.abort();
-          removeRequestInPool();
           throw DioException.receiveTimeout(
             timeout: receiveTimeout,
             requestOptions: options,
@@ -206,7 +181,7 @@ class IOHttpClientAdapter implements HttpClientAdapter {
         },
       );
     }
-    final responseStream = await future.whenComplete(removeRequestInPool);
+    final responseStream = await future;
 
     if (validateCertificate != null) {
       final host = options.uri.host;
diff --git a/plugins/http2_adapter/lib/src/http2_adapter.dart b/plugins/http2_adapter/lib/src/http2_adapter.dart
index ad059ba25..fc08ceb8d 100644
--- a/plugins/http2_adapter/lib/src/http2_adapter.dart
+++ b/plugins/http2_adapter/lib/src/http2_adapter.dart
@@ -41,10 +41,6 @@ class Http2Adapter implements HttpClientAdapter {
   /// Handles [DioH2NotSupportedException] and returns a [ResponseBody].
   H2NotSupportedCallback? onNotSupported;
 
-  final cancelFutureOperationPool =
-      <Future<void>, Set<CancelableOperation<ResponseBody>>>{};
-  final cancelFutureRequestPool = <Future<void>, Set<ClientTransportStream>>{};
-
   @override
   Future<ResponseBody> fetch(
     RequestOptions options,
@@ -56,25 +52,11 @@ class Http2Adapter implements HttpClientAdapter {
     final operation = CancelableOperation.fromFuture(
       _fetch(options, requestStream, cancelFuture, redirects),
     );
-
-    if (cancelFuture != null) {
-      cancelFutureOperationPool.putIfAbsent(cancelFuture, () => {});
-      cancelFutureRequestPool.putIfAbsent(cancelFuture, () => {});
-
-      cancelFuture.whenComplete(() {
-        cancelFutureOperationPool[cancelFuture]?.forEach((e) => e.cancel());
-        cancelFutureOperationPool.remove(cancelFuture);
-        cancelFutureRequestPool[cancelFuture]
-            ?.forEach((e) => e.outgoingMessages.close());
-        cancelFutureRequestPool.remove(cancelFuture);
-      });
-
-      cancelFutureOperationPool[cancelFuture]!.add(operation);
-    }
-
-    return operation.value.whenComplete(
-      () => cancelFutureOperationPool[cancelFuture]?.remove(operation),
-    );
+    final wr = WeakReference<CancelableOperation<ResponseBody>>(operation);
+    cancelFuture?.whenComplete(() {
+      wr.target?.cancel();
+    });
+    return operation.value;
   }
 
   Future<ResponseBody> _fetch(
@@ -150,17 +132,13 @@ class Http2Adapter implements HttpClientAdapter {
 
     // Creates a new outgoing stream.
     final stream = transport.makeRequest(headers);
+    final streamWR = WeakReference<ClientTransportStream>(stream);
 
     final hasRequestData = requestStream != null;
     if (hasRequestData && cancelFuture != null) {
-      cancelFutureRequestPool[cancelFuture]!.add(stream);
-    }
-
-    void removeRequestInPool() {
-      if (cancelFuture == null) {
-        return;
-      }
-      cancelFutureRequestPool[cancelFuture]?.remove(stream);
+      cancelFuture.whenComplete(() {
+        streamWR.target?.outgoingMessages.close();
+      });
     }
 
     List<Uint8List>? list;
@@ -179,7 +157,6 @@ class Http2Adapter implements HttpClientAdapter {
           sendTimeout,
           onTimeout: () {
             stream.outgoingMessages.close().catchError((_) {});
-            removeRequestInPool();
             throw DioException.sendTimeout(
               timeout: sendTimeout,
               requestOptions: options,
@@ -189,7 +166,7 @@ class Http2Adapter implements HttpClientAdapter {
       }
       await requestStreamFuture;
     }
-    await stream.outgoingMessages.close().whenComplete(removeRequestInPool);
+    await stream.outgoingMessages.close();
 
     final responseSink = StreamController<Uint8List>();
     final responseHeaders = Headers();