fix(api): Don't block subscription events during WebSocket reconnect (#7282)

Author: VarshithaPamisetty

packages/api/amplify_api_dart/lib/src/graphql/web_socket/blocs/web_socket_bloc.dart

@@ -108,6 +108,9 @@ class WebSocketBloc with AWSDebuggable, AmplifyLoggerMixin {
 
   late WebSocketState _currentState;
 
+  /// The in-flight reconnection, if any. Ensures at most one runs at a time.
+  Future<void>? _reconnectOperation;
+
   /// OVERRIDES
   ///
   ///
@@ -378,12 +381,25 @@ class WebSocketBloc with AWSDebuggable, AmplifyLoggerMixin {
   /// First establishes there is a connection to AppSync
   /// Then clears web socket connection and restarts init workflow
   /// Sends [NetworkException] when unable to reach AppSync
+  ///
+  /// Runs off the event queue so a slow reconnect doesn't block incoming
+  /// events (e.g. subscription data, keep alives) from being processed.
   Stream<WebSocketState> _reconnect() async* {
     assert(
       _currentState is ReconnectingState,
       'Bloc should be set to connecting before starting reconnection.',
     );
-    final state = _currentState;
+    _reconnectOperation ??= _performReconnect(
+      _currentState,
+    ).whenComplete(() => _reconnectOperation = null);
+
+    // TODO(dnys1): Yield broken on web debug build.
+    yield* const Stream.empty();
+  }
+
+  /// Pings AppSync with retry/back off and reinitializes the connection, or
+  /// shuts down on failure. Runs off the event queue via [_reconnect].
+  Future<void> _performReconnect(WebSocketState state) async {
     try {
       // Begin reconnection with retry/back off on ping endpoint
       final res = await state.options.retryOptions.retry(
@@ -395,16 +411,20 @@ class WebSocketBloc with AWSDebuggable, AmplifyLoggerMixin {
 
       // **Ping succeeded**
 
+      // Bloc may have shut down during the ping, don't act on a closed bloc
+      if (_isShuttingDown) return;
+
       // Prep new connection
       await state.service.close();
       for (final bloc in state.subscriptionBlocs.values) {
         bloc.add(SubscriptionPendingEvent(bloc.currentState.request.id));
       }
 
       // Init new connection
-      add(const InitEvent());
+      _safeAdd(const InitEvent());
     } on Exception catch (e, st) {
-      // Ping failed, close down
+      // Ping failed, nothing to do if already shutting down
+      if (_isShuttingDown) return;
       _shutdownWithException(
         NetworkException(
           'Unable to recover network connection, web socket will close.',
@@ -414,11 +434,15 @@ class WebSocketBloc with AWSDebuggable, AmplifyLoggerMixin {
         st,
       );
     }
-
-    // TODO(dnys1): Yield broken on web debug build.
-    yield* const Stream.empty();
   }
 
+  /// Whether the bloc is closed or shutting down.
+  bool get _isShuttingDown =>
+      _wsEventController.isClosed ||
+      _currentState is PendingDisconnect ||
+      _currentState is DisconnectedState ||
+      _currentState is FailureState;
+
   /// Sends registration message on ws channel when connected
   void _registerSubscriptionRequest(GraphQLRequest<Object?> request) {
     final currentState = _currentState;

packages/api/amplify_api_dart/test/web_socket/web_socket_bloc_test.dart

@@ -181,6 +181,71 @@ void main() {
       await dataCompleter.future;
     });
 
+    test(
+      'delivers incoming subscription data while a reconnect is in progress',
+      () async {
+        final firstData = Completer<String>();
+        final dataDuringReconnect = Completer<String>();
+
+        final subscribeEvent = SubscribeEvent(subscriptionRequest, () {
+          service!.channel.sink.add(mockDataString);
+        });
+
+        final bloc = getWebSocketBloc();
+        bloc.subscribe(subscribeEvent).listen((event) {
+          if (!firstData.isCompleted) {
+            firstData.complete(event.data);
+          } else if (!dataDuringReconnect.isCompleted) {
+            dataDuringReconnect.complete(event.data);
+          }
+        });
+
+        expect(await firstData.future, json.encode(mockSubscriptionData));
+
+        mockPollClient.induceTimeout = true;
+        mockNetworkStreamController.add(ConnectivityStatus.disconnected);
+        await expectLater(bloc.stream, emitsThrough(isA<ReconnectingState>()));
+
+        service!.channel.sink.add(mockDataString);
+
+        final data = await dataDuringReconnect.future.timeout(
+          const Duration(seconds: 3),
+          onTimeout: () => throw TimeoutException(
+            'Subscription data was not delivered while a reconnect was in '
+            'progress.',
+          ),
+        );
+        expect(data, json.encode(mockSubscriptionData));
+
+        mockPollClient.induceTimeout = false;
+      },
+    );
+
+    test('shuts down cleanly while a reconnect is in progress', () async {
+      final established = Completer<void>();
+      final bloc = getWebSocketBloc();
+      bloc
+          .subscribe(SubscribeEvent(subscriptionRequest, established.complete))
+          .listen((_) {}, onError: (_) {});
+      await established.future;
+
+      mockPollClient.induceTimeout = true;
+      mockNetworkStreamController.add(ConnectivityStatus.disconnected);
+      await expectLater(bloc.stream, emitsThrough(isA<ReconnectingState>()));
+
+      bloc.add(const ShutdownEvent());
+      await bloc.done.future.timeout(
+        const Duration(seconds: 3),
+        onTimeout: () => throw TimeoutException(
+          'Shutdown was blocked by an in-progress reconnect.',
+        ),
+      );
+
+      mockPollClient.induceTimeout = false;
+      await Future<void>.delayed(const Duration(seconds: 7));
+      expect(bloc.done.isCompleted, isTrue);
+    });
+
     test('should throttle reconnect after repeated wifi toggling', () async {
       final blocReady = Completer<void>();
       final subscribeEvent = SubscribeEvent(

packages/authenticator/amplify_authenticator/example/test_driver/integration_test.dart

@@ -5,4 +5,4 @@ import 'package:integration_test/integration_test_driver.dart';
 
 // Required for running integration tests in the browser:
 // https://docs.flutter.dev/cookbook/testing/integration/introduction#5b-web
-Future<void> main() => integrationDriver();
+Future<void> main() => integrationDriver(timeout: const Duration(minutes: 60));