Acquisition timeout

Author: robx

CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased (PostgREST fork)
 
+Added support for timing out resource acquisition ([PR #](https://github.com/PostgREST/hasql-pool/pull/3)).
 Added support for flushing the pool without destroying it ([PR #2](https://github.com/PostgREST/hasql-pool/pull/2)).
 
 # 0.7.2

hasql-pool.cabal

@@ -66,9 +66,11 @@ test-suite test
     Arrows, BangPatterns, ConstraintKinds, DataKinds, DefaultSignatures, DeriveDataTypeable, DeriveFoldable, DeriveFunctor, DeriveGeneric, DeriveTraversable, EmptyDataDecls, FlexibleContexts, FlexibleInstances, FunctionalDependencies, GADTs, GeneralizedNewtypeDeriving, LambdaCase, LiberalTypeSynonyms, MagicHash, MultiParamTypeClasses, MultiWayIf, NoImplicitPrelude, NoMonomorphismRestriction, OverloadedStrings, PatternGuards, ParallelListComp, QuasiQuotes, RankNTypes, RecordWildCards, ScopedTypeVariables, StandaloneDeriving, TemplateHaskell, TupleSections, TypeFamilies, TypeOperators, UnboxedTuples
   default-language:
     Haskell2010
+  ghc-options: -threaded
   build-depends:
     hasql,
     hasql-pool,
+    async,
     hspec >=2.6 && <3,
     rerebase >=1.15 && <2,
     stm >=2.5 && <3

library/Hasql/Pool.hs

@@ -22,6 +22,8 @@ import qualified Hasql.Session as Session
 data Pool = Pool
   { -- | Connection settings.
     poolFetchConnectionSettings :: IO Connection.Settings,
+    -- | Acquisition timeout, in microseconds.
+    poolAcquisitionTimeout :: Maybe Int,
     -- | Avail connections.
     poolConnectionQueue :: TQueue Connection,
     -- | Remaining capacity.
@@ -40,9 +42,9 @@ data Pool = Pool
 --
 -- No connections actually get established by this function. It is delegated
 -- to 'use'.
-acquire :: Int -> Connection.Settings -> IO Pool
-acquire poolSize connectionSettings =
-  acquireDynamically poolSize (pure connectionSettings)
+acquire :: Int -> Maybe Int -> Connection.Settings -> IO Pool
+acquire poolSize timeout connectionSettings =
+  acquireDynamically poolSize timeout (pure connectionSettings)
 
 -- | Given the pool-size and connection settings constructor action
 -- create a connection-pool.
@@ -52,9 +54,9 @@ acquire poolSize connectionSettings =
 --
 -- In difference to 'acquire' new settings get fetched each time a connection
 -- is created. This may be useful for some security models.
-acquireDynamically :: Int -> IO Connection.Settings -> IO Pool
-acquireDynamically poolSize fetchConnectionSettings = do
-  Pool fetchConnectionSettings
+acquireDynamically :: Int -> Maybe Int -> IO Connection.Settings -> IO Pool
+acquireDynamically poolSize timeout fetchConnectionSettings = do
+  Pool fetchConnectionSettings timeout
     <$> newTQueueIO
     <*> newTVarIO poolSize
     <*> (newTVarIO =<< newTVarIO True)
@@ -95,7 +97,13 @@ flush Pool {..} =
 -- and a slot gets freed up for a new connection to be established the next
 -- time one is needed. The error still gets returned from this function.
 use :: Pool -> Session.Session a -> IO (Either UsageError a)
-use Pool {..} sess =
+use Pool {..} sess = do
+  timeout <- case poolAcquisitionTimeout of
+    Just delta -> do
+      delay <- registerDelay delta
+      return $ readTVar delay
+    Nothing ->
+      return $ return False
   join . atomically $ do
     aliveVar <- readTVar poolAlive
     alive <- readTVar aliveVar
@@ -109,6 +117,11 @@ use Pool {..} sess =
                 then do
                   writeTVar poolCapacity $! pred capVal
                   return $ onNewConn aliveVar
+                else retry,
+            do
+              timedOut <- timeout
+              if timedOut
+                then return . return . Left $ AcquisitionTimeout
                 else retry
           ]
       else return . return . Left $ PoolIsReleasedUsageError
@@ -152,6 +165,8 @@ data UsageError
     SessionUsageError Session.QueryError
   | -- | Attempt to use a pool, which has already been called 'release' upon.
     PoolIsReleasedUsageError
+  | -- | Timeout acquiring a connection.
+    AcquisitionTimeout
   deriving (Show, Eq)
 
 instance Exception UsageError

test/Main.hs

@@ -10,70 +10,85 @@ import Test.Hspec
 import Prelude
 import qualified System.Environment
 import qualified Data.ByteString.Char8 as B8
+import Control.Concurrent.Async (race)
 
 main = do
   connectionSettings <- getConnectionSettings
   hspec $ describe "" $ do
     it "Releases a spot in the pool when there is a query error" $ do
-      pool <- acquire 1 connectionSettings
+      pool <- acquire 1 Nothing connectionSettings
       use pool badQuerySession `shouldNotReturn` (Right ())
       use pool selectOneSession `shouldReturn` (Right 1)
     it "Simulation of connection error works" $ do
-      pool <- acquire 3 connectionSettings
+      pool <- acquire 3 Nothing connectionSettings
       res <- use pool $ closeConnSession >> selectOneSession
       shouldSatisfy res $ \case
         Left (SessionUsageError (Session.QueryError _ _ (Session.ClientError _))) -> True
         _ -> False
     it "Connection errors cause eviction of connection" $ do
-      pool <- acquire 3 connectionSettings
+      pool <- acquire 3 Nothing connectionSettings
       res <- use pool $ closeConnSession >> selectOneSession
       res <- use pool $ closeConnSession >> selectOneSession
       res <- use pool $ closeConnSession >> selectOneSession
       res <- use pool $ selectOneSession
       shouldSatisfy res $ isRight
     it "Connection gets returned to the pool after normal use" $ do
-      pool <- acquire 3 connectionSettings
+      pool <- acquire 3 Nothing connectionSettings
       res <- use pool $ selectOneSession
       res <- use pool $ selectOneSession
       res <- use pool $ selectOneSession
       res <- use pool $ selectOneSession
       res <- use pool $ selectOneSession
       shouldSatisfy res $ isRight
     it "Connection gets returned to the pool after non-connection error" $ do
-      pool <- acquire 3 connectionSettings
+      pool <- acquire 3 Nothing connectionSettings
       res <- use pool $ badQuerySession
       res <- use pool $ badQuerySession
       res <- use pool $ badQuerySession
       res <- use pool $ badQuerySession
       res <- use pool $ selectOneSession
       shouldSatisfy res $ isRight
     it "Getting and setting session variables works" $ do
-      pool <- acquire 1 connectionSettings
+      pool <- acquire 1 Nothing connectionSettings
       res <- use pool $ getSettingSession "testing.foo"
       res `shouldBe` Right Nothing
       res <- use pool $ do
         setSettingSession "testing.foo" "hello world"
         getSettingSession "testing.foo"
       res `shouldBe` Right (Just "hello world")
     it "Session variables stay set when a connection gets reused" $ do
-      pool <- acquire 1 connectionSettings
+      pool <- acquire 1 Nothing connectionSettings
       res <- use pool $ setSettingSession "testing.foo" "hello world"
       res `shouldBe` Right ()
       res2 <- use pool $ getSettingSession "testing.foo"
       res2 `shouldBe` Right (Just "hello world")
     it "Flushing the pool resets session variables" $ do
-      pool <- acquire 1 connectionSettings
+      pool <- acquire 1 Nothing connectionSettings
       res <- use pool $ setSettingSession "testing.foo" "hello world"
       res `shouldBe` Right ()
       flush pool
       res <- use pool $ getSettingSession "testing.foo"
       res `shouldBe` Right Nothing
     it "Flushing a released pool leaves it dead" $ do
-      pool <- acquire 1 connectionSettings
+      pool <- acquire 1 Nothing connectionSettings
       release pool
       flush pool
       res <- use pool $ selectOneSession
       res `shouldBe` Left PoolIsReleasedUsageError
+    it "Times out connection acquisition" $ do
+      pool <- acquire 1 (Just 1000) connectionSettings -- 1ms timeout
+      sleeping <- newEmptyMVar
+      t0 <- getCurrentTime
+      res <- race
+        (use pool $ liftIO $ do
+           putMVar sleeping ()
+           threadDelay 1000000) -- 1s
+        (do
+           takeMVar sleeping
+           use pool $ selectOneSession)
+      t1 <- getCurrentTime
+      res `shouldBe` Right (Left AcquisitionTimeout)
+      diffUTCTime t1 t0 `shouldSatisfy` (< 0.5) -- 0.5s
 
 getConnectionSettings :: IO Connection.Settings
 getConnectionSettings = B8.unwords . catMaybes <$> sequence