testsuite: Add test for #13916
[ghc.git] / testsuite / tests / concurrent / should_run / T13916_Bracket.hs
1 {-# LANGUAGE RankNTypes #-}
2 {-# LANGUAGE RecordWildCards #-}
3 {- |
4 Module : Bracket
5 Description : Handling multiple environments with bracket-like apis
6 Maintainer : robertkennedy@clearwateranalytics.com
7 Stability : stable
8
9 This module is meant for ie Sql or mongo connections, where you may wish for some number of easy to grab
10 environments. In particular, this assumes your connection has some initialization/release functions
11
12 This module creates bugs with any optimizations enabled. The bugs do not occur if the program is in the same
13 module.
14 -}
15 module T13916_Bracket (
16 -- * Data Types
17 Spawner(..), Limit(..), Cache,
18 -- * Usage
19 withEnvCache, withEnv
20 ) where
21
22 import Control.Concurrent.STM
23 import Control.Concurrent.STM.TSem
24 import Control.Exception hiding (handle)
25 import Control.Monad
26 import Data.Vector (Vector)
27 import qualified Data.Vector as Vector
28
29 -- * Data Types
30 -- | Tells the program how many environments it is allowed to spawn.
31 -- A `Lax` limit will spawn extra connections if the `Cache` is empty,
32 -- while a `Hard` limit will not spawn any more than the given number of connections simultaneously.
33 --
34 -- @since 0.3.7
35 data Limit = Hard {getLimit :: {-# unpack #-} !Int}
36
37 data Spawner env = Spawner
38 { maker :: IO env
39 , killer :: env -> IO ()
40 , isDead :: env -> IO Bool
41 }
42
43 type VCache env = Vector (TMVar env)
44 data Cache env = Unlimited { spawner :: Spawner env
45 , vcache :: !(VCache env)
46 }
47 | Limited { spawner :: Spawner env
48 , vcache :: !(VCache env)
49 , envsem :: TSem
50 }
51
52 -- ** Initialization
53 withEnvCache :: Limit -> Spawner env -> (Cache env -> IO a) -> IO a
54 withEnvCache limit spawner = bracket starter releaseCache
55 where starter = case limit of
56 Hard n -> Limited spawner <$> initializeEmptyCache n <*> atomically (newTSem n)
57
58 -- ** Using a single value
59 withEnv :: Cache env -> (env -> IO a) -> IO a
60 withEnv cache = case cache of
61 Unlimited{..} -> withEnvUnlimited spawner vcache
62 Limited{..} -> withEnvLimited spawner vcache envsem
63
64 -- *** Unlimited
65 -- | Takes an env and returns it on completion of the function.
66 -- If all envs are already taken or closed, this will spin up a new env.
67 -- When the function finishes, this will attempt to put the env into the cache. If it cannot,
68 -- it will kill the env. Note this can lead to many concurrent connections.
69 --
70 -- @since 0.3.5
71 withEnvUnlimited :: Spawner env -> VCache env -> (env -> IO a) -> IO a
72 withEnvUnlimited Spawner{..} cache = bracket taker putter
73 where
74 taker = do
75 mpipe <- atomically $ tryTakeEnv cache
76 case mpipe of
77 Nothing -> maker
78 Just env -> isDead env >>= \b -> if not b then return env else killer env >> maker
79
80 putter env = do
81 accepted <- atomically $ tryPutEnv cache env
82 unless accepted $ killer env
83
84 -- *** Limited
85 -- | Takes an env and returns it on completion of the function.
86 -- If all envs are already taken, this will wait. This should have a constant number of environments
87 --
88 -- @since 0.3.6
89 withEnvLimited :: Spawner env -> VCache env -> TSem -> (env -> IO a) -> IO a
90 withEnvLimited spawner vcache envsem = bracket taker putter
91 where
92 taker = limitMakeEnv spawner vcache envsem
93 putter env = atomically $ putEnv vcache env
94
95 limitMakeEnv :: Spawner env -> VCache env -> TSem -> IO env
96 limitMakeEnv Spawner{..} vcache envsem = go
97 where
98 go = do
99 eenvpermission <- atomically $ ( Left <$> takeEnv vcache )
100 `orElse` ( Right <$> waitTSem envsem )
101 case eenvpermission of
102 Right () -> maker
103 Left env -> do
104 -- Given our env, we check if it's dead. If it's not, we are done and return it.
105 -- If it is dead, we release it, signal that a new env can be created, and then recurse
106 isdead <- isDead env
107 if not isdead then return env
108 else do
109 killer env
110 atomically $ signalTSem envsem
111 go
112
113 -- * Low level
114 initializeEmptyCache :: Int -> IO (VCache env)
115 initializeEmptyCache n | n < 1 = return mempty
116 | otherwise = Vector.replicateM n newEmptyTMVarIO
117
118 takeEnv :: VCache env -> STM env
119 takeEnv = Vector.foldl folding retry
120 where folding m stmenv = m `orElse` takeTMVar stmenv
121
122 tryTakeEnv :: VCache env -> STM (Maybe env)
123 tryTakeEnv cache = (Just <$> takeEnv cache) `orElse` pure Nothing
124
125 putEnv :: VCache env -> env -> STM ()
126 putEnv cache env = Vector.foldl folding retry cache
127 where folding m stmenv = m `orElse` putTMVar stmenv env
128
129 tryPutEnv :: VCache env -> env -> STM Bool
130 tryPutEnv cache env = (putEnv cache env *> return True) `orElse` pure False
131
132 releaseCache :: Cache env -> IO ()
133 releaseCache cache = Vector.mapM_ qkRelease (vcache cache)
134 where qkRelease tenv = atomically (tryTakeTMVar tenv)
135 >>= maybe (return ()) (killer $ spawner cache)