Close failed connections without reconnecting during exception cleanup

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 ()

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

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

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)

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