Fix rare broken socket detection issues + tool to delay bytes in test…

…s. (#375)
isoos · Sep 13, 2024 · 2a8d9cf · 2a8d9cf
1 parent 76601a1
commit 2a8d9cf
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 6 deletions.
diff --git a/lib/src/v3/connection.dart b/lib/src/v3/connection.dart
@@ -202,6 +202,8 @@ class PgConnectionImplementation extends _PgSessionBase implements Connection {
   static Future<PgConnectionImplementation> connect(
     Endpoint endpoint, {
     ConnectionSettings? connectionSettings,
+    @visibleForTesting
+    StreamTransformer<Uint8List, Uint8List>? incomingBytesTransformer,
   }) async {
     final settings = connectionSettings is ResolvedConnectionSettings
         ? connectionSettings
@@ -217,6 +219,7 @@ class PgConnectionImplementation extends _PgSessionBase implements Connection {
       endpoint,
       settings,
       codecContext: codecContext,
+      incomingBytesTransformer: incomingBytesTransformer,
     );
 
     if (_debugLog) {
@@ -257,6 +260,7 @@ class PgConnectionImplementation extends _PgSessionBase implements Connection {
     Endpoint endpoint,
     ResolvedConnectionSettings settings, {
     required CodecContext codecContext,
+    StreamTransformer<Uint8List, Uint8List>? incomingBytesTransformer,
   }) async {
     final host = endpoint.host;
     final port = endpoint.port;
@@ -337,6 +341,10 @@ class PgConnectionImplementation extends _PgSessionBase implements Connection {
       adaptedStream = async.SubscriptionStream(subscription);
     }
 
+    if (incomingBytesTransformer != null) {
+      adaptedStream = adaptedStream.transform(incomingBytesTransformer);
+    }
+
     final outgoingSocket = async.StreamSinkExtensions(socket)
         .transform<Uint8List>(
             async.StreamSinkTransformer.fromHandlers(handleDone: (out) {
@@ -373,6 +381,7 @@ class PgConnectionImplementation extends _PgSessionBase implements Connection {
   final bool _channelIsSecure;
   late final StreamSubscription<Message> _serverMessages;
   bool _isClosing = false;
+  bool _socketIsBroken = false;
 
   _PendingOperation? _pending;
   // Errors happening while a transaction is active will roll back the
@@ -558,19 +567,21 @@ class PgConnectionImplementation extends _PgSessionBase implements Connection {
 
   Future<void> _close(bool interruptRunning, PgException? cause,
       {bool socketIsBroken = false}) async {
+    _socketIsBroken = _socketIsBroken || socketIsBroken;
     if (!_isClosing) {
       _isClosing = true;
 
       if (interruptRunning) {
         _pending?.handleConnectionClosed(cause);
-        if (!socketIsBroken) {
+        if (!_socketIsBroken) {
           _channel.sink.add(const TerminateMessage());
         }
       } else {
         // Wait for the previous operation to complete by using the lock
         await _operationLock.withResource(() {
-          // Use lock to await earlier operations
-          _channel.sink.add(const TerminateMessage());
+          if (!_socketIsBroken) {
+            _channel.sink.add(const TerminateMessage());
+          }
         });
       }
 
@@ -580,7 +591,11 @@ class PgConnectionImplementation extends _PgSessionBase implements Connection {
   }
 
   void _closeAfterError([PgException? cause]) {
-    _close(true, cause);
+    _close(
+      true,
+      cause,
+      socketIsBroken: cause?.willAbortConnection ?? false,
+    );
   }
 }
 

diff --git a/test/docker.dart b/test/docker.dart
@@ -1,5 +1,6 @@
 import 'dart:async';
 import 'dart:io';
+import 'dart:typed_data';
 
 import 'package:async/async.dart';
 import 'package:docker_process/containers/postgres.dart';
@@ -8,9 +9,12 @@ import 'package:meta/meta.dart';
 import 'package:path/path.dart' as p;
 import 'package:postgres/messages.dart';
 import 'package:postgres/postgres.dart';
+import 'package:postgres/src/v3/connection.dart';
 import 'package:stream_channel/stream_channel.dart';
 import 'package:test/test.dart';
 
+final _splitAndDelayBytes = false;
+
 // We log all packets sent to and received from the postgres server. This can be
 // used to debug failing tests. To view logs, something like this can be put
 // at the beginning of `main()`:
@@ -64,16 +68,18 @@ class PostgresServer {
     SslMode? sslMode,
     QueryMode? queryMode,
   }) async {
-    return Connection.open(
+    return await PgConnectionImplementation.connect(
       await endpoint(),
-      settings: ConnectionSettings(
+      connectionSettings: ConnectionSettings(
         connectTimeout: Duration(seconds: 3),
         queryTimeout: Duration(seconds: 3),
         replicationMode: replicationMode,
         transformer: loggingTransformer('conn'),
         sslMode: sslMode,
         queryMode: queryMode,
       ),
+      incomingBytesTransformer:
+          _splitAndDelayBytes ? _transformIncomingBytes() : null,
     );
   }
 
@@ -82,6 +88,30 @@ class PostgresServer {
   }
 }
 
+StreamTransformer<Uint8List, Uint8List> _transformIncomingBytes() {
+  return StreamTransformer.fromBind((s) => s.asyncExpand((u) {
+        if (u.length <= 2) {
+          return Stream.value(u);
+        }
+        final hash = u.hashCode.abs();
+        final split = hash % u.length;
+        if (split == 0 || split >= u.length - 1) {
+          return Stream.value(u);
+        }
+
+        final p1 = u.sublist(0, split);
+        final p2 = u.sublist(split);
+
+        return Stream.fromFutures([
+          Future.value(p1),
+          Future.delayed(
+            Duration(milliseconds: 50),
+            () => p2,
+          )
+        ]);
+      }));
+}
+
 @isTestGroup
 void withPostgresServer(
   String name,