diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,11 @@ +2026-04-21 Laurent P. René de Cotret 0.1.2 + +* Eliminated a rare race condition that allowed the transport to read messages before + marking the connection as open, violating the interface expectations. +* Better handling of lost connections. +* Fixed rare race condition in establishing connections. + 2026-01-01 Laurent P. René de Cotret 0.1.1 * Documentation and packaging improvements. diff --git a/bench/Bench.hs b/bench/Bench.hs --- a/bench/Bench.hs +++ b/bench/Bench.hs @@ -58,7 +58,15 @@ >>= \case Left errmsg -> throwIO $ userError errmsg Right credentials -> - QUIC.createTransport "127.0.0.1" "0" (credentials :| []) + QUIC.createTransport + ( QUIC.QUICTransportConfig + { hostName = "127.0.0.1" + , serviceName = "0" + , credentials = credentials :| [] + , -- credentials are self-signed + validateCredentials = False + } + ) } data BenchParams = BenchParams @@ -115,8 +123,7 @@ connections <- replicateM connectionCount - ( connect senderEP receiverAddr ReliableOrdered defaultConnectHints >>= either throwIO pure - ) + (connect senderEP receiverAddr ReliableOrdered defaultConnectHints >>= either throwIO pure) takeMVar receiverReady diff --git a/network-transport-quic.cabal b/network-transport-quic.cabal --- a/network-transport-quic.cabal +++ b/network-transport-quic.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 Name: network-transport-quic -Version: 0.1.1 +Version: 0.1.2 build-Type: Simple License: BSD-3-Clause License-file: LICENSE @@ -20,7 +20,7 @@ In dense network topologies, using ["Network.Transport.QUIC"] may improve performance by a factor of 2 over other transport implementations. -tested-with: GHC==8.10.7 GHC==9.0.2 GHC==9.2.8 GHC==9.4.8 GHC==9.6.7 GHC==9.8.4 GHC==9.10.3 GHC==9.12.2 +tested-with: GHC==8.10.7 GHC==9.0.2 GHC==9.2.8 GHC==9.4.8 GHC==9.6.7 GHC==9.8.4 GHC==9.10.3 GHC==9.12.2 GHC==9.14.1 Category: Network extra-doc-files: README.md @@ -62,10 +62,10 @@ -- Prior to version 0.2.20, `quic` had issues with handling -- pending data in the stream buffer. This meant that vectored -- message sends did not work correctly at the transport layer - , quic >=0.2.20 && <0.3 + , quic >=0.2.20 && <0.4 , stm >=2.4 && <2.6 - , tls >= 2.1 && < 2.2 - , tls-session-manager >= 0.0.5 && <0.1 + , tls >= 2.1 && < 2.5 + , tls-session-manager >= 0.0.5 && <0.2 exposed-modules: Network.Transport.QUIC Network.Transport.QUIC.Internal other-modules: Network.Transport.QUIC.Internal.Configuration diff --git a/src/Network/Transport/QUIC/Internal.hs b/src/Network/Transport/QUIC/Internal.hs --- a/src/Network/Transport/QUIC/Internal.hs +++ b/src/Network/Transport/QUIC/Internal.hs @@ -22,7 +22,7 @@ where import Control.Concurrent (forkIO, killThread, modifyMVar_, newEmptyMVar, readMVar) -import Control.Concurrent.MVar (modifyMVar, putMVar, takeMVar, withMVar) +import Control.Concurrent.MVar (modifyMVar, putMVar, takeMVar, tryPutMVar, withMVar) import Control.Concurrent.STM (atomically, newTQueueIO) import Control.Concurrent.STM.TQueue ( TQueue, @@ -34,13 +34,11 @@ import Data.Bifunctor (Bifunctor (first)) import Data.Binary qualified as Binary (decodeOrFail) import Data.ByteString (ByteString, fromStrict) -import Data.Foldable (forM_) import Data.Function ((&)) import Data.Functor ((<&>)) import Data.IORef (newIORef, readIORef, writeIORef) import Data.List.NonEmpty (NonEmpty) import Data.Map.Strict qualified as Map -import Data.Maybe (isNothing) import Lens.Micro.Platform ((+~)) import Network.QUIC qualified as QUIC import Network.TLS (Credential) @@ -62,8 +60,7 @@ ) import Network.Transport.QUIC.Internal.Configuration (credentialLoadX509) import Network.Transport.QUIC.Internal.Messaging - ( ClientConnId, - MessageReceived (..), + ( MessageReceived (..), createConnectionId, decodeMessage, encodeMessage, @@ -101,7 +98,6 @@ nextSelfConnOutId, remoteEndPointAddress, remoteEndPointState, - remoteIncoming, remoteServerConnId, remoteStream, transportConfig, @@ -182,21 +178,15 @@ (remoteEndPoint, _) <- either throwIO pure =<< createRemoteEndPoint ourEndPoint remoteAddress Incoming doneMVar <- newEmptyMVar - -- Sending an ack is important, because otherwise - -- the client may start sending messages well before we - -- start being able to receive them - - clientConnId <- either (throwIO . userError) (pure . fromIntegral) =<< recvWord32 stream let serverConnId = remoteServerConnId remoteEndPoint - connectionId = createConnectionId serverConnId clientConnId + -- One logical connection per stream; clientConnId is always 0. + connectionId = createConnectionId serverConnId 0 let st = RemoteEndPointValid $ ValidRemoteEndPointState { _remoteStream = stream, - _remoteStreamIsClosed = doneMVar, - _remoteIncoming = Just clientConnId, - _remoteNextConnOutId = 0 + _remoteStreamIsClosed = doneMVar } modifyMVar_ (remoteEndPoint ^. remoteEndPointState) @@ -205,13 +195,10 @@ _ -> undefined ) - tid <- - forkIO $ - -- If we've reached this stage, the connection handhake succeeded - handleIncomingMessages - ourEndPoint - remoteEndPoint - + -- Enqueue ConnectionOpened before forking the reader thread. + -- Otherwise the reader may dequeue a Message off the stream and + -- enqueue Received/ConnectionClosed before ConnectionOpened, which + -- breaks the ordering guarantee the transport API owes callers. atomically $ writeTQueue (ourEndPoint ^. localQueue) @@ -221,6 +208,19 @@ remoteAddress ) + -- Second handshake ack: only sent after ConnectionOpened is + -- enqueued. The initiator's @connect@ call blocks on this ack, so + -- when it returns, the caller can trust that any peer observing + -- events on this endpoint will see ConnectionOpened before any + -- messages the caller subsequently sends on other connections. + sendAck stream + + tid <- + forkIO $ + handleIncomingMessages + ourEndPoint + remoteEndPoint + takeMVar doneMVar QUIC.shutdownStream stream killThread tid @@ -247,12 +247,8 @@ release (Left err) = closeRemoteEndPoint Incoming remoteEndPoint >> prematureExit err release (Right _) = closeRemoteEndPoint Incoming remoteEndPoint - connectionId = createConnectionId serverConnId - - writeConnectionClosedSTM connId = - writeTQueue - ourQueue - (ConnectionClosed (connectionId connId)) + -- One logical connection per stream; clientConnId is always 0. + connectionId = createConnectionId serverConnId 0 go = either prematureExit loop @@ -262,34 +258,31 @@ Left errmsg -> do -- Throwing will trigger 'prematureExit' throwIO $ userError $ "(handleIncomingMessages) Failed with: " <> errmsg - Right (Message connId bytes) -> handleMessage connId bytes >> loop stream + Right (Message bytes) -> handleMessage bytes >> loop stream Right StreamClosed -> throwIO $ userError "(handleIncomingMessages) Stream closed" - Right (CloseConnection connId) -> do - atomically (writeConnectionClosedSTM connId) + Right CloseConnection -> do + atomically (writeTQueue ourQueue (ConnectionClosed connectionId)) mAct <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case RemoteEndPointInit -> pure (RemoteEndPointClosed, Nothing) RemoteEndPointClosed -> pure (RemoteEndPointClosed, Nothing) - RemoteEndPointValid (ValidRemoteEndPointState _ isClosed _ _) -> do + RemoteEndPointValid (ValidRemoteEndPointState _ isClosed) -> do pure (RemoteEndPointClosed, Just $ putMVar isClosed ()) case mAct of Nothing -> pure () Just cleanup -> cleanup Right CloseEndPoint -> do - connIds <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case - RemoteEndPointValid vst -> do - pure (RemoteEndPointClosed, vst ^. remoteIncoming) - other -> pure (other, Nothing) - unless - (isNothing connIds) - ( atomically $ - forM_ - connIds - (writeTQueue ourQueue . ConnectionClosed . connectionId) - ) + -- handleIncomingMessages only runs on incoming remote endpoints, so if + -- the state was still Valid there is exactly one logical connection to + -- surface as closed. + wasValid <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case + RemoteEndPointValid _ -> pure (RemoteEndPointClosed, True) + other -> pure (other, False) + when wasValid $ + atomically $ writeTQueue ourQueue (ConnectionClosed connectionId) - handleMessage :: ClientConnId -> [ByteString] -> IO () - handleMessage clientConnId payload = - atomically (writeTQueue ourQueue (Received (connectionId clientConnId) payload)) + handleMessage :: [ByteString] -> IO () + handleMessage payload = + atomically (writeTQueue ourQueue (Received connectionId payload)) prematureExit :: IOException -> IO () prematureExit exc = do @@ -357,24 +350,24 @@ else createConnectionTo creds validateCreds ourEndPoint remoteAddress >>= \case Left err -> pure $ Left err - Right (remoteEndPoint, connId) -> do + Right remoteEndPoint -> do connAlive <- newIORef True pure . Right $ Connection - { send = sendConn remoteEndPoint connAlive connId, - close = closeConn remoteEndPoint connAlive connId + { send = sendConn remoteEndPoint connAlive, + close = closeConn remoteEndPoint connAlive } where ourAddress = ourEndPoint ^. localAddress - sendConn remoteEndPoint connAlive connId packets = + sendConn remoteEndPoint connAlive packets = readMVar (remoteEndPoint ^. remoteEndPointState) >>= \case RemoteEndPointInit -> undefined RemoteEndPointValid vst -> readIORef connAlive >>= \case False -> pure . Left $ TransportError SendClosed "Connection closed" True -> - sendMessage (vst ^. remoteStream) connId packets + sendMessage (vst ^. remoteStream) packets <&> first (TransportError SendFailed . show) RemoteEndPointClosed -> do readIORef connAlive >>= \case @@ -384,16 +377,21 @@ -- 'connAlive' IORefs. False -> pure . Left $ TransportError SendClosed "Connection closed" True -> pure . Left $ TransportError SendFailed "Remote endpoint closed" - closeConn remoteEndPoint connAlive connId = do + closeConn remoteEndPoint connAlive = do mCleanup <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case - RemoteEndPointValid vst@(ValidRemoteEndPointState stream isClosed _ _) -> do + RemoteEndPointValid vst@(ValidRemoteEndPointState stream isClosed) -> do readIORef connAlive >>= \case False -> pure (RemoteEndPointValid vst, Nothing) True -> do writeIORef connAlive False - -- We want to run this cleanup action OUTSIDE of the MVar modification - let cleanup = sendCloseConnection connId stream - pure (RemoteEndPointClosed, Just $ cleanup >> putMVar isClosed ()) + -- Run cleanup OUTSIDE the MVar modification. tryPutMVar keeps this + -- safe against races with the finally in streamToEndpoint that can + -- also signal isClosed on QUIC.Client.run exit. + let cleanup = do + _ <- sendCloseConnection stream + _ <- tryPutMVar isClosed () + pure () + pure (RemoteEndPointClosed, Just cleanup) _ -> pure (RemoteEndPointClosed, Nothing) case mCleanup of diff --git a/src/Network/Transport/QUIC/Internal/Client.hs b/src/Network/Transport/QUIC/Internal/Client.hs --- a/src/Network/Transport/QUIC/Internal/Client.hs +++ b/src/Network/Transport/QUIC/Internal/Client.hs @@ -10,8 +10,8 @@ import Control.Concurrent (forkIOWithUnmask, newEmptyMVar) import Control.Concurrent.Async (withAsync) -import Control.Concurrent.MVar (MVar, putMVar, takeMVar) -import Control.Exception (SomeException, bracket, catch, mask, mask_, throwIO) +import Control.Concurrent.MVar (MVar, putMVar, takeMVar, tryPutMVar) +import Control.Exception (SomeException, bracket, catch, finally, mask, mask_, throwIO) import Data.List.NonEmpty (NonEmpty) import Network.QUIC qualified as QUIC import Network.QUIC.Client qualified as QUIC.Client @@ -28,9 +28,11 @@ EndPointAddress -> -- | Their address EndPointAddress -> - -- | On exception - (SomeException -> IO ()) -> - -- | On a message to forcibly close the connection + -- | Called when the QUIC connection or stream ends without us having initiated the + -- close. Must be idempotent (the caller typically gates on remote endpoint state so + -- that repeated invocations are safe) — this handler is invoked from multiple sites + -- (peer-initiated close signal, QUIC.Client.run exception, thread finally) to cover + -- every termination path. IO () -> IO ( Either @@ -40,7 +42,7 @@ QUIC.Stream ) ) -streamToEndpoint creds validateCreds ourAddress theirAddress onExc onCloseForcibly = +streamToEndpoint creds validateCreds ourAddress theirAddress onConnLoss = case decodeQUICAddr theirAddress of Left errmsg -> pure $ Left (TransportError ConnectNotFound errmsg) Right (QUICAddr hostname servicename _) -> do @@ -75,7 +77,8 @@ (throwIO @SomeException) ) ) - onExc + (\(_ :: SomeException) -> pure ()) + `finally` onConnLoss streamOrError <- takeMVar streamMVar @@ -85,11 +88,24 @@ listenForClose stream doneMVar = receiveMessage stream >>= \case - Right StreamClosed -> do - putMVar doneMVar () - Right (CloseConnection _) -> do - putMVar doneMVar () - Right CloseEndPoint -> do - putMVar doneMVar () - onCloseForcibly + -- Any message from the peer on this stream means we're done listening. + -- Peer-initiated closes (StreamClosed/CloseEndPoint) additionally call + -- onConnLoss; the idempotent gate in the handler dedupes with the finally + -- that also fires on QUIC.Client.run exit. + -- + -- Mask signalling+onConnLoss as an atomic pair: tryPutMVar unblocks + -- runClient's takeMVar, which causes withAsync to cancel this thread. + -- Without mask, the async ThreadKilled can fire partway through + -- onConnLoss, dropping the ErrorEvent. The finally in the parent thread + -- is a backup but cannot recover if surfaceConnectionLost already + -- transitioned the remote state to Closed. + Right StreamClosed -> mask_ $ do + _ <- tryPutMVar doneMVar () + onConnLoss + Right CloseConnection -> + -- Peer closed the logical connection cleanly; no ErrorEvent. + () <$ tryPutMVar doneMVar () + Right CloseEndPoint -> mask_ $ do + _ <- tryPutMVar doneMVar () + onConnLoss other -> throwIO . userError $ "Unexpected incoming message to client: " <> show other diff --git a/src/Network/Transport/QUIC/Internal/Messaging.hs b/src/Network/Transport/QUIC/Internal/Messaging.hs --- a/src/Network/Transport/QUIC/Internal/Messaging.hs +++ b/src/Network/Transport/QUIC/Internal/Messaging.hs @@ -4,6 +4,7 @@ {-# LANGUAGE LambdaCase #-} {-# LANGUAGE NumericUnderscores #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} module Network.Transport.QUIC.Internal.Messaging ( -- * Connections @@ -48,22 +49,20 @@ import Network.Transport (ConnectionId, EndPointAddress) import Network.Transport.Internal (decodeWord32, encodeWord32) import Network.Transport.QUIC.Internal.QUICAddr (QUICAddr (QUICAddr), decodeQUICAddr) -import System.Timeout (timeout) --- | Send a message to a remote endpoint ID +-- | Send a message on the stream. -- -- This function is thread-safe; while the data is sending, asynchronous -- exceptions are masked, to be rethrown after the data is sent. sendMessage :: Stream -> - ClientConnId -> [ByteString] -> IO (Either QUIC.QUICException ()) -sendMessage stream connId messages = +sendMessage stream messages = try ( QUIC.sendStreamMany stream - (encodeMessage connId messages) + (encodeMessage messages) ) -- | Receive a message, including its local destination endpoint ID @@ -87,18 +86,15 @@ -- The encoding is composed of a header, and the payloads. -- The message header is composed of: -- 1. A control byte, to determine how the message should be parsed. --- 2. A 32-bit word that encodes the endpoint ID of the destination endpoint; --- 3. A 32-bit word that encodes the number of frames in the message +-- 2. A 32-bit word that encodes the number of frames in the message -- -- The payload frames are each prepended with the length of the frame. encodeMessage :: - ClientConnId -> [ByteString] -> [ByteString] -encodeMessage connId messages = +encodeMessage messages = BS.concat [ BS.singleton messageControlByte, - encodeWord32 (fromIntegral connId), encodeWord32 (fromIntegral $ length messages) ] : [encodeWord32 (fromIntegral $ BS.length message) <> message | message <- messages] @@ -116,37 +112,37 @@ where go ctrl | ctrl == closeEndPointControlByte = pure $ Right CloseEndPoint - | ctrl == closeConnectionControlByte = Right . CloseConnection . fromIntegral <$> getWord32 + | ctrl == closeConnectionControlByte = pure $ Right CloseConnection | ctrl == messageControlByte = do - connId <- getWord32 numMessages <- getWord32 messages <- replicateM (fromIntegral numMessages) $ do getWord32 >>= get . fromIntegral - pure . Right $ Message (fromIntegral connId) messages + pure . Right $ Message messages | otherwise = pure $ Left $ "Unsupported control byte: " <> show ctrl getWord32 = get 4 <&> decodeWord32 -- | Wrap a method to fetch bytes, to ensure that we always get exactly the --- intended number of bytes. +-- intended number of bytes. Returns early (with the accumulated bytes) if the +-- underlying fetcher signals EOF by returning an empty ByteString; otherwise a +-- fetcher that repeatedly returns empty after a peer FIN would cause this to +-- spin forever. getAllBytes :: -- | Function to fetch at most 'n' bytes (Int -> IO ByteString) -> - -- | Function to fetch exactly 'n' bytes + -- | Function to fetch exactly 'n' bytes (or fewer on EOF) (Int -> IO ByteString) getAllBytes get n = go n mempty where go 0 !acc = pure $ BS.concat acc go m !acc = get m >>= \bytes -> - go - (m - BS.length bytes) - (acc <> [bytes]) + if BS.null bytes + then pure $ BS.concat acc + else go (m - BS.length bytes) (acc <> [bytes]) data MessageReceived - = Message - {-# UNPACK #-} !ClientConnId - {-# UNPACK #-} ![ByteString] - | CloseConnection !ClientConnId + = Message {-# UNPACK #-} ![ByteString] + | CloseConnection | CloseEndPoint | StreamClosed deriving (Show, Eq) @@ -176,11 +172,7 @@ recvAck :: Stream -> IO (Either () ()) recvAck stream = do - -- TODO: make timeout configurable - timeout 500_000 (QUIC.recvStream stream 1) - >>= maybe - (throwIO (AckException "Connection ack not received within acceptable timeframe")) - go + QUIC.recvStream stream 1 >>= go where go response | response == ackMessage = pure $ Right () @@ -221,13 +213,12 @@ closeConnectionControlByte = 255 -- | Send a message to close the connection. -sendCloseConnection :: ClientConnId -> Stream -> IO (Either QUIC.QUICException ()) -sendCloseConnection connId stream = +sendCloseConnection :: Stream -> IO (Either QUIC.QUICException ()) +sendCloseConnection stream = try ( QUIC.sendStream stream - ( BS.concat [BS.singleton closeConnectionControlByte, encodeWord32 (fromIntegral connId)] - ) + (BS.singleton closeConnectionControlByte) ) -- | Send a message to close the connection. @@ -241,8 +232,16 @@ ) -- | Handshake protocol that a client, connecting to a remote endpoint, --- has to perform. --- TODO: encode server part of the handhake +-- has to perform: +-- +-- 1. client -> server: address payload +-- 2. server -> client: ack1 (handshake payload accepted) +-- 3. server -> client: ack2 (ConnectionOpened has been enqueued on server's endpoint) +-- +-- The ack2 step is load-bearing: when this function returns, the server has +-- already written ConnectionOpened to its local queue. Without it, the server's +-- @connect@ would return before the peer's queue has the ConnectionOpened event, +-- which races with subsequent sends on other connections. handshake :: (EndPointAddress, EndPointAddress) -> Stream -> @@ -251,9 +250,6 @@ case decodeQUICAddr theirAddress of Left errmsg -> throwIO $ userError ("Could not decode QUIC address: " <> errmsg) Right (QUICAddr _ _ serverEndPointId) -> do - -- Handshake on connection creation, which simply involves - -- sending our address over, and - -- the endpoint ID of the endpoint we want to communicate with let encodedPayload = BS.toStrict $ Binary.encode (ourAddress, serverEndPointId) payloadLength = encodeWord32 $ fromIntegral (BS.length encodedPayload) @@ -265,10 +261,9 @@ >>= \case Left (_exc :: SomeException) -> pure $ Left () Right _ -> - -- Server acknowledgement that the handshake is complete - -- means that we cannot send messages until the server - -- is ready for them - recvAck stream + recvAck stream >>= \case + Left () -> pure $ Left () + Right () -> recvAck stream -- | Part of the connection ID that is client-allocated. newtype ClientConnId = ClientConnId Word32 diff --git a/src/Network/Transport/QUIC/Internal/QUICTransport.hs b/src/Network/Transport/QUIC/Internal/QUICTransport.hs --- a/src/Network/Transport/QUIC/Internal/QUICTransport.hs +++ b/src/Network/Transport/QUIC/Internal/QUICTransport.hs @@ -60,8 +60,6 @@ ValidRemoteEndPointState (..), remoteStream, remoteStreamIsClosed, - remoteIncoming, - remoteNextConnOutId, Direction (..), -- * Re-exports @@ -70,13 +68,11 @@ where import Control.Concurrent.Async (forConcurrently_) -import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar, putMVar, readMVar) +import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar, readMVar, tryPutMVar) import Control.Concurrent.STM.TQueue (TQueue, writeTQueue) -import Control.Exception (Exception (displayException), SomeException, bracketOnError, try) +import Control.Exception (bracketOnError) import Control.Monad (forM_) import Control.Monad.STM (atomically) -import Data.Binary qualified as Binary -import Data.ByteString qualified as BS import Data.Function ((&)) import Data.List.NonEmpty (NonEmpty) import Data.List.NonEmpty qualified as NE @@ -85,7 +81,6 @@ import Data.Word (Word32) import Lens.Micro.Platform (makeLenses, (%~), (+~), (^.)) import Network.QUIC (Stream) -import Network.QUIC qualified as QUIC import Network.Socket (HostName, ServiceName, Socket) import Network.Socket qualified as N import Network.TLS (Credential) @@ -231,9 +226,7 @@ data ValidRemoteEndPointState = ValidRemoteEndPointState { _remoteStream :: Stream, - _remoteStreamIsClosed :: MVar (), - _remoteIncoming :: !(Maybe ClientConnId), - _remoteNextConnOutId :: !ClientConnId + _remoteStreamIsClosed :: MVar () } makeLenses ''QUICTransport @@ -311,10 +304,13 @@ LocalEndPointStateClosed -> pure (LocalEndPointStateClosed, Nothing) LocalEndPointStateValid st -> pure (LocalEndPointStateClosed, Just st) - forM_ mPreviousState $ \vst -> - forConcurrently_ - (vst ^. incomingConnections <> vst ^. outgoingConnections) - tryCloseRemoteStream + -- Close outgoing remote endpoints before incoming. The peer's handleIncomingMessages + -- reader writes ConnectionClosed in response to our outgoing close; its listenForClose + -- writes ErrorEvent in response to our incoming close. Processing outgoing first gives + -- the peer's event queue the expected ConnectionClosed-before-ErrorEvent ordering. + forM_ mPreviousState $ \vst -> do + forConcurrently_ (vst ^. outgoingConnections) tryCloseRemoteStream + forConcurrently_ (vst ^. incomingConnections) tryCloseRemoteStream atomically $ writeTQueue (localEndPoint ^. localQueue) EndPointClosed where tryCloseRemoteStream :: RemoteEndPoint -> IO () @@ -326,8 +322,9 @@ pure ( RemoteEndPointClosed, Just $ do - sendCloseEndPoint (vst ^. remoteStream) - >> putMVar (vst ^. remoteStreamIsClosed) () + _ <- sendCloseEndPoint (vst ^. remoteStream) + _ <- tryPutMVar (vst ^. remoteStreamIsClosed) () + pure () ) case mCleanup of @@ -344,12 +341,13 @@ mAct <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case RemoteEndPointInit -> pure (RemoteEndPointClosed, Nothing) RemoteEndPointClosed -> pure (RemoteEndPointClosed, Nothing) - RemoteEndPointValid (ValidRemoteEndPointState stream isClosed conns _) -> - let cleanup = - case direction of - Outgoing -> Right <$> forM_ conns (`sendCloseConnection` stream) + RemoteEndPointValid (ValidRemoteEndPointState stream isClosed) -> + let cleanup = do + _ <- case direction of + Outgoing -> sendCloseConnection stream Incoming -> sendCloseEndPoint stream - >> putMVar isClosed () + _ <- tryPutMVar isClosed () + pure () in pure (RemoteEndPointClosed, Just cleanup) case mAct of @@ -399,45 +397,53 @@ Bool -> LocalEndPoint -> EndPointAddress -> - IO (Either (TransportError ConnectErrorCode) (RemoteEndPoint, ClientConnId)) + IO (Either (TransportError ConnectErrorCode) RemoteEndPoint) createConnectionTo creds validateCreds localEndPoint remoteAddress = do createRemoteEndPoint localEndPoint remoteAddress Outgoing >>= \case Left err -> pure $ Left err - Right (remoteEndPoint, _) -> + Right (remoteEndPoint, _) -> do + -- TODO: each call to @connect@ currently opens a dedicated QUIC connection + -- and carries a single logical connection on its stream. Preferred + -- architecture: one QUIC connection per (local endpoint, peer endpoint) + -- pair, with each logical connection carried on its own stream. Streams + -- already give us independent flow control and avoid head-of-line blocking. streamToEndpoint creds validateCreds (localEndPoint ^. localAddress) remoteAddress - (\_ -> closeRemoteEndPoint Outgoing remoteEndPoint) - onConnectionLost + (surfaceConnectionLost remoteEndPoint) >>= \case Left exc -> pure $ Left exc Right (closeStream, stream) -> do - let clientConnId = 0 - validState = + let validState = RemoteEndPointValid $ ValidRemoteEndPointState { _remoteStream = stream, - _remoteStreamIsClosed = closeStream, - _remoteIncoming = Nothing, - _remoteNextConnOutId = clientConnId + 1 + _remoteStreamIsClosed = closeStream } modifyMVar_ (remoteEndPoint ^. remoteEndPointState) (\_ -> pure validState) - - try - ( QUIC.sendStream - stream - ( BS.toStrict $ - Binary.encode clientConnId - ) - ) - >>= \case - Left (exc :: SomeException) -> pure . Left $ TransportError ConnectFailed (displayException exc) - Right () -> pure $ Right (remoteEndPoint, clientConnId) + pure $ Right remoteEndPoint where + -- Idempotent: surfaces EventConnectionLost exactly once, only if the remote + -- endpoint was still Valid when invoked. Called from multiple termination + -- sites (peer-initiated close, QUIC exception, forked-thread finally) so that + -- no close path can leave us silent — the state-transition gate dedupes them. + surfaceConnectionLost remoteEndPoint = do + mAct <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case + RemoteEndPointInit -> pure (RemoteEndPointClosed, Nothing) + RemoteEndPointClosed -> pure (RemoteEndPointClosed, Nothing) + RemoteEndPointValid (ValidRemoteEndPointState stream isClosed) -> + let cleanup = do + _ <- sendCloseConnection stream + _ <- tryPutMVar isClosed () + onConnectionLost + in pure (RemoteEndPointClosed, Just cleanup) + case mAct of + Nothing -> pure () + Just act -> act onConnectionLost = atomically . writeTQueue (localEndPoint ^. localQueue) diff --git a/test/Test/Network/Transport/QUIC.hs b/test/Test/Network/Transport/QUIC.hs --- a/test/Test/Network/Transport/QUIC.hs +++ b/test/Test/Network/Transport/QUIC.hs @@ -10,18 +10,19 @@ import Control.Monad (replicateM_) import Data.ByteString qualified as BS import Data.List.NonEmpty (NonEmpty (..)) -import Network.Transport (EndPoint (..), Event (ConnectionClosed, ConnectionOpened, Received), Reliability (..), Transport (..), close, defaultConnectHints, send) +import Network.Transport (EndPoint (..), Event (ConnectionClosed), Reliability (..), Transport (..), close, defaultConnectHints, send) import Network.Transport.QUIC (QUICTransportConfig (..)) import Network.Transport.QUIC qualified as QUIC import Network.Transport.Tests (echoServer) import Network.Transport.Tests qualified as Tests import Network.Transport.Tests.Auxiliary (forkTry) +import Network.Transport.Tests.Expect (expectConnectionOpened, expectEq, expectReceived, expectRight) import Network.Transport.Util (spawn) import System.FilePath (()) import System.Timeout (timeout) import Test.Tasty (TestName, TestTree, testGroup) -import Test.Tasty.Flaky (flakyTest, limitRetries) -import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=)) +import Test.Tasty.Flaky (flakyTest, limitRetries, constantDelay) +import Test.Tasty.HUnit (Assertion, assertFailure, testCase, (@?=)) tests :: TestTree tests = @@ -47,7 +48,7 @@ ] flaky :: TestTree -> TestTree -flaky = flakyTest (limitRetries 3) +flaky = flakyTest (limitRetries 3 <> constantDelay 1_000) -- | Ensure that a test does not run for too long testCaseWithTimeout :: TestName -> Assertion -> TestTree @@ -91,21 +92,22 @@ let numPings = 10 let bigMessage = BS.replicate 4091 66 -- Using an odd number of bytes (4091) to test message boundaries _ <- forkTry $ do - Right endpoint <- newEndPoint transport + endpoint <- expectRight "newEndPoint" =<< newEndPoint transport ping endpoint server numPings bigMessage putMVar result () takeMVar result where ping endpoint serverAddr numPings message = do - Right conn <- connect endpoint serverAddr ReliableOrdered defaultConnectHints + conn <- expectRight "connect" =<< connect endpoint serverAddr ReliableOrdered defaultConnectHints - ConnectionOpened cid _ _ <- receive endpoint + (cid, _, _) <- expectConnectionOpened =<< receive endpoint replicateM_ numPings $ do _ <- send conn [message] - Received cid' [reply] <- receive endpoint - assertBool mempty $ cid == cid' && reply == message + (cid', payload) <- expectReceived =<< receive endpoint + expectEq "connection id" cid cid' + expectEq "payload" [message] payload close conn diff --git a/test/Test/Network/Transport/QUIC/Internal/Messaging.hs b/test/Test/Network/Transport/QUIC/Internal/Messaging.hs --- a/test/Test/Network/Transport/QUIC/Internal/Messaging.hs +++ b/test/Test/Network/Transport/QUIC/Internal/Messaging.hs @@ -21,20 +21,16 @@ testMessageEncodingAndDecoding :: TestTree testMessageEncodingAndDecoding = testProperty "Encoded messages can be decoded" $ property $ do - -- The connection ID and message length are encoded and decoded the same way, to/from - -- a Word32. - -- To exercise the parsing of Word32s, we need to make sure that the range - -- of data is generated above a Word8 (255), including the connection ID - -- and the number of bytes in the message - endpointId <- fmap fromIntegral <$> forAll $ Gen.word32 Range.constantBounded - + -- The message length is encoded and decoded as a Word32. Generate data above + -- a Word8 (255) to exercise the Word32 parsing of the number of bytes in each + -- message. messages <- forAll (Gen.list (Range.linear 0 3) (Gen.bytes (Range.linear 1 4096))) - let encoded = mconcat $ encodeMessage endpointId messages + let encoded = mconcat $ encodeMessage messages getBytes <- liftIO $ messageDecoder encoded decoded <- liftIO $ decodeMessage getBytes - Right (Message endpointId messages) === decoded + Right (Message messages) === decoded messageDecoder :: ByteString -> IO (Int -> IO ByteString) messageDecoder allBytes = do