rethrow SocketException as ConnectionFailedTemporarilyException

edgedb · Apr 19, 2024 · e20c36d · e20c36d
1 parent de0bd13
commit e20c36d
Showing 1 changed file with 39 additions and 10 deletions.
diff --git a/src/EdgeDB.Net.Driver/Binary/Duplexers/StreamDuplexer.cs b/src/EdgeDB.Net.Driver/Binary/Duplexers/StreamDuplexer.cs
@@ -4,6 +4,7 @@
 using Microsoft.Extensions.Logging;
 using System.Buffers;
 using System.Diagnostics;
+using System.Net.Sockets;
 using System.Runtime.CompilerServices;
 
 namespace EdgeDB.Binary;
@@ -111,7 +112,7 @@ public async ValueTask DisconnectAsync(CancellationToken token = default)
 
             if (packetFactory is null)
             {
-                // unknow/unsupported packet
+                // unknown/unsupported packet
                 _client.Logger.UnknownPacket($"{header.Type}{{0x{(byte)header.Type}}}:{header.Length}");
 
                 await DisconnectInternalAsync();
@@ -121,16 +122,27 @@ public async ValueTask DisconnectAsync(CancellationToken token = default)
             var packet = PacketSerializer.DeserializePacket(in packetFactory, in buffer);
 
             // check for idle timeout
-            if (packet is IProtocolError err && err.ErrorCode == ServerErrorCodes.IdleSessionTimeoutError)
-            {
-                // all connection state needs to be reset for the client here.
-                _client.Logger.IdleDisconnect();
+            if (packet is not IProtocolError {ErrorCode: ServerErrorCodes.IdleSessionTimeoutError} err) return packet;
 
-                await DisconnectInternalAsync();
-                throw new EdgeDBErrorException(err);
-            }
+            // all connection state needs to be reset for the client here.
+            _client.Logger.IdleDisconnect();
 
-            return packet;
+            await DisconnectInternalAsync();
+            throw new EdgeDBErrorException(err);
+        }
+        catch (IOException ioException) when (ioException.InnerException is SocketException socketException)
+        {
+            switch (socketException.SocketErrorCode)
+            {
+                case SocketError.ConnectionRefused
+                    or SocketError.ConnectionAborted
+                    or SocketError.ConnectionReset
+                    or SocketError.HostNotFound
+                    or SocketError.NotInitialized:
+                    throw new ConnectionFailedTemporarilyException(socketException.SocketErrorCode);
+                default:
+                    throw;
+            }
         }
         catch (EndOfStreamException)
         {
@@ -168,7 +180,24 @@ public async ValueTask SendAsync(CancellationToken token = default, params Senda
 
         using var linkedToken = CancellationTokenSource.CreateLinkedTokenSource(token, _disconnectTokenSource.Token);
 
-        await _stream.WriteAsync(BinaryUtils.BuildPackets(packets), linkedToken.Token).ConfigureAwait(false);
+        try
+        {
+            await _stream.WriteAsync(BinaryUtils.BuildPackets(packets), linkedToken.Token).ConfigureAwait(false);
+        }
+        catch (IOException ioException) when (ioException.InnerException is SocketException socketException)
+        {
+            switch (socketException.SocketErrorCode)
+            {
+                case SocketError.ConnectionRefused
+                    or SocketError.ConnectionAborted
+                    or SocketError.ConnectionReset
+                    or SocketError.HostNotFound
+                    or SocketError.NotInitialized:
+                    throw new ConnectionFailedTemporarilyException(socketException.SocketErrorCode);
+                default:
+                    throw;
+            }
+        }
 
         // only perform second iteration if debug log enabled.
         if (_client.Logger.IsEnabled(LogLevel.Debug))