diff --git a/documentation/docs/testing/t004-errors/index.lhs b/documentation/docs/testing/t004-errors/index.lhs index 90353f8f..2cb05705 100644 --- a/documentation/docs/testing/t004-errors/index.lhs +++ b/documentation/docs/testing/t004-errors/index.lhs @@ -26,25 +26,23 @@ module Main (main) where import ClickHaskell -- GHC included -import Control.Exception (try) +import Control.Exception (bracket, try) import GHC.Generics (Generic) main :: IO () main = do - connection <- openConnection defaultConnectionArgs - connOld <- openConnection (overrideMaxRevision 1 defaultConnectionArgs) - - runErrorsTest connection - runErrorsTest connOld + runErrorsTest defaultConnectionArgs + runErrorsTest (overrideMaxRevision 1 defaultConnectionArgs) putStrLn "Ok" -runErrorsTest :: Connection -> IO () -runErrorsTest connection = do - validateUnmatchedColumn connection - validateUnmatchedType connection - validateUnmatchedColumnsCount connection +runErrorsTest :: ConnectionArgs -> IO () +runErrorsTest args = do + -- A failed operation invalidates its handle; retry with a fresh connection. + bracket (openConnection args) closeConnection validateUnmatchedColumn + bracket (openConnection args) closeConnection validateUnmatchedType + bracket (openConnection args) closeConnection validateUnmatchedColumnsCount validateUnmatchedColumn :: Connection -> IO () diff --git a/library/ClickHaskell.cabal b/library/ClickHaskell.cabal index b7d82280..567b8978 100644 --- a/library/ClickHaskell.cabal +++ b/library/ClickHaskell.cabal @@ -30,10 +30,7 @@ Flag dev Default: False Manual: True -Library - Exposed-Modules: - ClickHaskell - ClickHaskell.Protocol +common clickhaskell-settings Other-Modules: ClickHaskell.Connection ClickHaskell.Protocol.Client @@ -120,3 +117,20 @@ Library ScopedTypeVariables StandaloneDeriving UndecidableInstances + + +Library + import: clickhaskell-settings + Exposed-Modules: + ClickHaskell + ClickHaskell.Protocol + +Test-Suite connection-lifecycle-regressions + import: clickhaskell-settings + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: test/connection-lifecycle + other-modules: + ClickHaskell + ClickHaskell.Protocol + ghc-options: -threaded diff --git a/library/ClickHaskell.hs b/library/ClickHaskell.hs index 552a9230..26713ad1 100644 --- a/library/ClickHaskell.hs +++ b/library/ClickHaskell.hs @@ -13,7 +13,7 @@ module ClickHaskell {- * Connection -} ConnectionArgs, defaultConnectionArgs , setHost, setPort, setUser, setDatabase, setPassword - , Connection(..), openConnection + , Connection(..), openConnection, closeConnection {- ** Hacking -} , overrideInitConnection, BufferArgs(..) , overrideHostname @@ -91,7 +91,7 @@ import ClickHaskell.Protocol.Client (addParameter, IsParameterType) -- GHC included import Control.Concurrent (newMVar, putMVar, takeMVar) -import Control.Exception (Exception, mask, onException, throw, throwIO) +import Control.Exception (Exception, finally, mask, onException, throw, throwIO) import Control.Monad (when) import Data.Binary.Get import Data.ByteString.Builder @@ -110,15 +110,15 @@ import Data.WideWord (Int128 (..), Word128 (..), Int256(..), Word256(..)) -- * Connection openConnection :: HasCallStack => ConnectionArgs -> IO Connection -openConnection creds@MkConnectionArgs{mHostname, mOsUser} = do +openConnection creds@MkConnectionArgs{mHostname, mOsUser} = mask $ \restore -> do hostname <- maybe (lookupEnv "HOSTNAME") (pure . Just) mHostname osUser <- maybe (lookupEnv "USER") (pure . Just) mOsUser - connectionState <- + connectionState <- restore $ createConnectionState auth . (maybe id overrideHostname hostname) . (maybe id overrideOsUser osUser) $ creds - MkConnection <$> newMVar connectionState + MkConnection <$> newMVar (Just connectionState) -- * Statements and commands @@ -317,18 +317,29 @@ readBuffer MkBuffer{readBuff, writeBuff} parser = runBufferReader (runGetIncreme (Fail _leftover _consumed msg) -> throwIO (InternalError $ DeserializationError msg) +-- Failed operations invalidate the handle; its owner chooses when to reopen. +-- Never perform network reconnection while unwinding an exception. withConnection :: HasCallStack => Connection -> (ConnectionState -> IO a) -> IO a -withConnection (MkConnection connStateMVar) f = +withConnection (MkConnection cell) f = mask $ \restore -> do - connState <- takeMVar connStateMVar - b <- onException - (restore (f connState)) - (do - newConnState <- recreateConnectionState auth connState - putMVar connStateMVar newConnState - ) - putMVar connStateMVar connState - return b + cached <- takeMVar cell + case cached of + Nothing -> putMVar cell Nothing >> throwIO ConnectionClosed + Just state -> do + result <- restore (f state) `onException` + (destroyBuff (buffer state) `finally` putMVar cell Nothing) + putMVar cell (Just state) + pure result + +-- | Close the physical connection. Repeated calls do nothing. +-- Waits for any current operation; stop concurrent IO before disposing a pool. +-- After close, or after a failed operation, further operations throw +-- 'ConnectionClosed'. Open a fresh connection to retry: cleanup never reconnects. +closeConnection :: Connection -> IO () +closeConnection (MkConnection cell) = + mask $ \_ -> do + cached <- takeMVar cell + maybe (pure ()) (destroyBuff . buffer) cached `finally` putMVar cell Nothing auth :: Buffer -> ConnectionArgs -> IO ConnectionState auth buffer creds@MkConnectionArgs{db, user, pass, mOsUser, mHostname, maxRevision} = do diff --git a/library/ClickHaskell/Connection.hs b/library/ClickHaskell/Connection.hs index 787ec13b..29bcb229 100644 --- a/library/ClickHaskell/Connection.hs +++ b/library/ClickHaskell/Connection.hs @@ -5,7 +5,7 @@ import ClickHaskell.Primitive -- GHC included import Control.Concurrent (MVar) -import Control.Exception (SomeException, bracketOnError, catch, finally, throwIO) +import Control.Exception (SomeException, bracketOnError, catch, finally, mask, onException, throwIO) import Data.Binary.Builder (Builder, toLazyByteString) import Data.ByteString as BS (ByteString, null) import Data.IORef (newIORef, readIORef, writeIORef) @@ -25,7 +25,9 @@ import System.Timeout (timeout) Errors occured on connection operations -} data ConnectionError - = NoAdressResolved + = ConnectionClosed + -- ^ The connection was explicitly closed or invalidated by a failed operation. + | NoAdressResolved -- ^ Occurs when 'getAddrInfo' returns an empty result | EstablishTimeout -- ^ Occurs on 'socket' connection timeout @@ -47,7 +49,7 @@ writeToConnection :: ConnectionState -> (ProtocolRevision -> Builder) -> IO () writeToConnection MkConnectionState{revision, buffer} serializer = (writeConn buffer) (serializer revision) -data Connection where MkConnection :: (MVar ConnectionState) -> Connection +data Connection where MkConnection :: (MVar (Maybe ConnectionState)) -> Connection data ConnectionState = MkConnectionState { buffer :: Buffer @@ -64,8 +66,12 @@ createConnectionState -> IO ConnectionState createConnectionState postInitAction creds@MkConnectionArgs{..} = do let port = fromMaybe defPort mPort - buffer <- mkBuffer =<< mkBufferArgs host =<< initSocket =<< resolveHostName host port - postInitAction buffer creds + address <- resolveHostName host port + mask $ \restore -> do + sock <- initSocket address + args <- restore (mkBufferArgs host sock) `onException` close sock + buffer <- mkBuffer args `onException` closeSock args + restore (postInitAction buffer creds) `onException` destroyBuff buffer recreateConnectionState :: (Buffer -> ConnectionArgs -> IO ConnectionState) diff --git a/library/test/connection-lifecycle/Main.hs b/library/test/connection-lifecycle/Main.hs new file mode 100644 index 00000000..ed9d67ac --- /dev/null +++ b/library/test/connection-lifecycle/Main.hs @@ -0,0 +1,125 @@ +module Main (main) where + +import ClickHaskell +import ClickHaskell.Connection (ConnectionArgs(..)) +import Control.Concurrent (newEmptyMVar, takeMVar) +import Control.Exception (IOException, SomeException, bracket, throwIO, try) +import Control.Monad (replicateM_, unless) +import qualified Data.ByteString as BS +import Data.IORef +import Network.Socket +import System.Timeout (timeout) + +main :: IO () +main = do + result <- timeout 5000000 runTests + case result of + Just () -> pure () + Nothing -> fail "regression suite timed out" + +runTests :: IO () +runTests = do + withTransport $ \args closed readAction _ -> do + writeIORef readAction (pure hello) + conn <- openConnection args + closeConnection conn + closeConnection conn + assertCount "idempotent close" closed 1 + expectClosed conn + withTransport $ \args closed readAction _ -> do + writeIORef readAction (pure hello) + conn <- openConnection args + writeIORef readAction (throwIO (userError "read failed")) + failed <- try @IOException (ping conn) + unless (either (const True) (const False) failed) (fail "operation did not fail") + assertCount "failed operation closes transport" closed 1 + expectClosed conn + closeConnection conn + assertCount "failed handle does not close twice" closed 1 + replicateM_ 4 $ withTransport $ \args closed readAction _ -> do + writeIORef readAction (pure hello) + conn <- openConnection args + blocked <- newEmptyMVar + writeIORef readAction (takeMVar blocked) + result <- timeout 100000 (ping conn) + unless (result == Nothing) (fail "stalled operation did not time out") + assertCount "cancelled operation closes transport" closed 1 + expectClosed conn + withTransport $ \args closed readAction _ -> do + writeIORef readAction (throwIO (userError "authentication failed")) + expectOpenFailure args + assertCount "failed authentication closes transport" closed 1 + withTransport $ \args closed readAction _ -> do + blocked <- newEmptyMVar + writeIORef readAction (takeMVar blocked) + result <- timeout 100000 (openConnection args) + case result of + Nothing -> pure () + Just conn -> closeConnection conn >> fail "authentication did not time out" + assertCount "cancelled authentication closes transport" closed 1 + withTransport $ \args _ _ sockets -> do + let broken = args { mkBufferArgs = \_ _ -> throwIO (userError "transport init failed") } + expectOpenFailure broken + allocated <- readIORef sockets + case allocated of + [sock] -> do + result <- try @IOException (getSocketName sock) + unless (either (const True) (const False) result) (fail "initialization leaked its socket") + _ -> fail "unexpected reconnection during initialization" + -- Even a custom transport whose close throws must leave a closed handle. + withTransport $ \args closed readAction _ -> do + writeIORef readAction (pure hello) + let broken = args { mkBufferArgs = \host sock -> do + bufferArgs <- mkBufferArgs args host sock + pure bufferArgs { closeSock = closeSock bufferArgs >> throwIO (userError "close failed") } } + conn <- openConnection broken + _ <- try @IOException (closeConnection conn) + expectClosed conn + closeConnection conn + assertCount "throwing close invalidates handle" closed 1 + putStrLn "Connection lifecycle regressions passed" + +-- Revision 1 HelloResponse: no optional handshake fields. +hello :: BS.ByteString +hello = BS.pack [0, 4, 116, 101, 115, 116, 1, 1, 1] + +expectClosed :: Connection -> IO () +expectClosed conn = do + result <- timeout 100000 (try @ConnectionError (ping conn)) + case result of + Just (Left ConnectionClosed) -> pure () + _ -> fail "failed/closed handle blocked or was reused" + +expectOpenFailure :: ConnectionArgs -> IO () +expectOpenFailure args = do + result <- try @SomeException (openConnection args) + case result of + Left _ -> pure () + Right conn -> closeConnection conn >> fail "connection unexpectedly opened" + +assertCount :: String -> IORef Int -> Int -> IO () +assertCount label ref expected = do + actual <- readIORef ref + unless (actual == expected) (fail (label ++ ": expected " ++ show expected ++ ", got " ++ show actual)) + +-- The real socket is owned by the driver. Reads/writes use a scripted transport, +-- so these failure and cancellation tests do not need a ClickHouse server. +withTransport :: (ConnectionArgs -> IORef Int -> IORef (IO BS.ByteString) -> IORef [Socket] -> IO a) -> IO a +withTransport action = bracket (newIORef []) cleanup $ \sockets -> do + closed <- newIORef 0 + readAction <- newIORef (pure hello) + let args = (overrideMaxRevision 1 defaultConnectionArgs) + { resolveHostName = \_ _ -> pure defaultHints + , initSocket = \_ -> do + sock <- socket AF_INET Stream defaultProtocol + modifyIORef' sockets (sock :) + pure sock + , mkBufferArgs = \_ sock -> pure MkBufferArgs + { writeSock = const (pure ()) + , readSock = readIORef readAction >>= id + , closeSock = modifyIORef' closed (+1) >> close sock + } + } + action args closed readAction sockets + where + cleanup sockets = readIORef sockets >>= mapM_ close