@@ -6,10 +6,13 @@
66{-# LANGUAGE KindSignatures #-}
77{-# LANGUAGE RecordWildCards #-}
88{-# LANGUAGE StandaloneDeriving #-}
9+{-# LANGUAGE TupleSections #-}
910module Servant.Server.Internal.RoutingApplication where
1011
12+import Control.Exception (finally)
1113import Control.Monad (ap, liftM)
1214import Control.Monad.Trans (MonadIO(..))
15+import Data.IORef (newIORef, readIORef, IORef, atomicModifyIORef)
1316import Network.Wai (Application, Request,
1417 Response, ResponseReceived)
1518import Prelude ()
@@ -112,12 +115,33 @@ instance Functor (Delayed env) where
112115 , ..
113116 } -- Note [Existential Record Update]
114117
118+-- | A mutable cleanup action
119+newtype CleanupRef = CleanupRef (IORef (IO ()))
120+
121+newCleanupRef :: IO CleanupRef
122+newCleanupRef = CleanupRef <$> newIORef (return ())
123+
124+-- | Add a clean up action to a 'CleanupRef'
125+addCleanup' :: IO () -> CleanupRef -> IO ()
126+addCleanup' act (CleanupRef ref) = atomicModifyIORef ref (\act' -> (act' >> act, ()))
127+
128+addCleanup :: IO () -> DelayedIO ()
129+addCleanup act = DelayedIO $ \_req cleanupRef ->
130+ addCleanup' act cleanupRef >> return (Route ())
131+
132+-- | Run all the clean up actions registered in
133+-- a 'CleanupRef'.
134+runCleanup :: CleanupRef -> IO ()
135+runCleanup (CleanupRef ref) = do
136+ act <- readIORef ref
137+ act
138+
115139-- | Computations used in a 'Delayed' can depend on the
116140-- incoming 'Request', may perform 'IO, and result in a
117141-- 'RouteResult, meaning they can either suceed, fail
118142-- (with the possibility to recover), or fail fatally.
119143--
120−newtype DelayedIO a = DelayedIO { runDelayedIO :: Request -> IO (RouteResult a) }
144+newtype DelayedIO a = DelayedIO { runDelayedIO :: Request -> CleanupRef -> IO (RouteResult a) }
121145
122146instance Functor DelayedIO where
123147 fmap = liftM
@@ -127,17 +151,17 @@ instance Applicative DelayedIO where
127151 (<*>) = ap
128152
129153instance Monad DelayedIO where
130− return x = DelayedIO (const $ return (Route x))
154+ return x = DelayedIO (\_req _cleanup -> return (Route x))
131155 DelayedIO m >>= f =
132− DelayedIO $ \ req -> do
133− r <- m req
156+ DelayedIO $ \ req cl -> do
157+ r <- m req cl
134158 case r of
135159 Fail e -> return $ Fail e
136160 FailFatal e -> return $ FailFatal e
137− Route a -> runDelayedIO (f a) req
161+ Route a -> runDelayedIO (f a) req cl
138162
139163instance MonadIO DelayedIO where
140− liftIO m = DelayedIO (const $ Route <$> m)
164+ liftIO m = DelayedIO (\_req _cl -> Route <$> m)
141165
142166-- | A 'Delayed' without any stored checks.
143167emptyDelayed :: RouteResult a -> Delayed env a
@@ -148,15 +172,15 @@ emptyDelayed result =
148172
149173-- | Fail with the option to recover.
150174delayedFail :: ServantErr -> DelayedIO a
151−delayedFail err = DelayedIO (const $ return $ Fail err)
175+delayedFail err = DelayedIO (\_req _cleanup -> return $ Fail err)
152176
153177-- | Fail fatally, i.e., without any option to recover.
154178delayedFailFatal :: ServantErr -> DelayedIO a
155−delayedFailFatal err = DelayedIO (const $ return $ FailFatal err)
179+delayedFailFatal err = DelayedIO (\_req _cleanup -> return $ FailFatal err)
156180
157181-- | Gain access to the incoming request.
158182withRequest :: (Request -> DelayedIO a) -> DelayedIO a
159−withRequest f = DelayedIO (\ req -> runDelayedIO (f req) req)
183+withRequest f = DelayedIO (\ req cl -> runDelayedIO (f req) req cl)
160184
161185-- | Add a capture to the end of the capture block.
162186addCapture :: Delayed env (a -> b)
@@ -196,8 +220,8 @@ addBodyCheck :: Delayed env (a -> b)
196220 -> Delayed env b
197221addBodyCheck Delayed{..} new =
198222 Delayed
199− { bodyD = (,) <$> bodyD <*> new
200− , serverD = \ c a (z, v) req -> ($ v) <$> serverD c a z req
223+ { bodyD = (,) <$> bodyD <*> new
224+ , serverD = \ c a (z, v) req -> ($ v) <$> serverD c a z req
201225 , ..
202226 } -- Note [Existential Record Update]
203227
@@ -240,13 +264,18 @@ passToServer Delayed{..} x =
240264runDelayed :: Delayed env a
241265 -> env
242266 -> Request
267+ -> CleanupRef
243268 -> IO (RouteResult a)
244−runDelayed Delayed{..} env = runDelayedIO $ do
245− c <- capturesD env
246− methodD
247− a <- authD
248− b <- bodyD
249− DelayedIO (\ req -> return $ serverD c a b req)
269+runDelayed Delayed{..} env req cleanupRef =
270+ runDelayedIO
271+ (do c <- capturesD env
272+ methodD
273+ a <- authD
274+ b <- bodyD
275+ DelayedIO $ \ r _cleanup -> return (serverD c a b r)
276+ )
277+ req
278+ cleanupRef
250279
251280-- | Runs a delayed server and the resulting action.
252281-- Takes a continuation that lets us send a response.
@@ -258,8 +287,11 @@ runAction :: Delayed env (Handler a)
258287 -> (RouteResult Response -> IO r)
259288 -> (a -> RouteResult Response)
260289 -> IO r
261−runAction action env req respond k =
262− runDelayed action env req >>= go >>= respond
290+runAction action env req respond k = do
291+ cleanupRef <- newCleanupRef
292+ (runDelayed action env req cleanupRef >>= go >>= respond)
293+ `finally` runCleanup cleanupRef
294+
263295 where
264296 go (Fail e) = return $ Fail e
265297 go (FailFatal e) = return $ FailFatal e