@@ -10,10 +10,10 @@
1010module Servant.Server.Internal.RoutingApplication where
1111
1212import Control.Exception (bracket)
13−import Control.Monad (ap, liftM)
13+import Control.Monad (ap, liftM, (>=>))
1414import Control.Monad.Trans (MonadIO(..))
1515import Control.Monad.Trans.Except (runExceptT)
16−import Data.IORef (newIORef, readIORef, writeIORef)
16+import Data.IORef (newIORef, readIORef, writeIORef, IORef, atomicModifyIORef)
1717import Network.Wai (Application, Request,
1818 Response, ResponseReceived)
1919import Prelude ()
@@ -107,10 +107,6 @@ data Delayed env c where
107107 , authD :: DelayedIO auth
108108 , bodyD :: DelayedIO body
109109 , serverD :: captures -> auth -> body -> Request -> RouteResult c
110− , cleanupD :: body -> IO ()
111− -- not in DelayedIO because:
112− -- - most likely should not depend on the request
113− -- - simpler
114110 } -> Delayed env c
115111
116112instance Functor (Delayed env) where
@@ -120,12 +116,33 @@ instance Functor (Delayed env) where
120116 , ..
121117 } -- Note [Existential Record Update]
122118
119+-- | A mutable cleanup action
120+newtype CleanupRef = CleanupRef (IORef (IO ()))
121+
122+newCleanupRef :: IO CleanupRef
123+newCleanupRef = CleanupRef <$> newIORef (return ())
124+
125+-- | Add a clean up action to a 'CleanupRef'
126+addCleanup' :: IO () -> CleanupRef -> IO ()
127+addCleanup' act (CleanupRef ref) = atomicModifyIORef ref (\act' -> (act' >> act, ()))
128+
129+addCleanup :: IO () -> DelayedIO ()
130+addCleanup act = DelayedIO $ \_req cleanupRef ->
131+ addCleanup' act cleanupRef >> return (Route ())
132+
133+-- | Run all the clean up actions registered in
134+-- a 'CleanupRef'.
135+runCleanup :: CleanupRef -> IO ()
136+runCleanup (CleanupRef ref) = do
137+ act <- readIORef ref
138+ act
139+
123140-- | Computations used in a 'Delayed' can depend on the
124141-- incoming 'Request', may perform 'IO, and result in a
125142-- 'RouteResult, meaning they can either suceed, fail
126143-- (with the possibility to recover), or fail fatally.
127144--
128−newtype DelayedIO a = DelayedIO { runDelayedIO :: Request -> IO (RouteResult a) }
145+newtype DelayedIO a = DelayedIO { runDelayedIO :: Request -> CleanupRef -> IO (RouteResult a) }
129146
130147instance Functor DelayedIO where
131148 fmap = liftM
@@ -135,36 +152,36 @@ instance Applicative DelayedIO where
135152 (<*>) = ap
136153
137154instance Monad DelayedIO where
138− return x = DelayedIO (const $ return (Route x))
155+ return x = DelayedIO (\_req _cleanup -> return (Route x))
139156 DelayedIO m >>= f =
140− DelayedIO $ \ req -> do
141− r <- m req
157+ DelayedIO $ \ req cl -> do
158+ r <- m req cl
142159 case r of
143160 Fail e -> return $ Fail e
144161 FailFatal e -> return $ FailFatal e
145− Route a -> runDelayedIO (f a) req
162+ Route a -> runDelayedIO (f a) req cl
146163
147164instance MonadIO DelayedIO where
148− liftIO m = DelayedIO (const $ Route <$> m)
165+ liftIO m = DelayedIO (\_req _cl -> Route <$> m)
149166
150167-- | A 'Delayed' without any stored checks.
151168emptyDelayed :: RouteResult a -> Delayed env a
152169emptyDelayed result =
153− Delayed (const r) r r r (\ _ _ _ _ -> result) (const $ return ())
170+ Delayed (const r) r r r (\ _ _ _ _ -> result)
154171 where
155172 r = return ()
156173
157174-- | Fail with the option to recover.
158175delayedFail :: ServantErr -> DelayedIO a
159−delayedFail err = DelayedIO (const $ return $ Fail err)
176+delayedFail err = DelayedIO (\_req _cleanup -> return $ Fail err)
160177
161178-- | Fail fatally, i.e., without any option to recover.
162179delayedFailFatal :: ServantErr -> DelayedIO a
163−delayedFailFatal err = DelayedIO (const $ return $ FailFatal err)
180+delayedFailFatal err = DelayedIO (\_req _cleanup -> return $ FailFatal err)
164181
165182-- | Gain access to the incoming request.
166183withRequest :: (Request -> DelayedIO a) -> DelayedIO a
167−withRequest f = DelayedIO (\ req -> runDelayedIO (f req) req)
184+withRequest f = DelayedIO (\ req cl -> runDelayedIO (f req) req cl)
168185
169186-- | Add a capture to the end of the capture block.
170187addCapture :: Delayed env (a -> b)
@@ -206,7 +223,6 @@ addBodyCheck Delayed{..} new =
206223 Delayed
207224 { bodyD = (,) <$> bodyD <*> new
208225 , serverD = \ c a (z, v) req -> ($ v) <$> serverD c a z req
209− , cleanupD = cleanupD . fst -- not sure that's right
210226 , ..
211227 } -- Note [Existential Record Update]
212228
@@ -249,19 +265,18 @@ passToServer Delayed{..} x =
249265runDelayed :: Delayed env a
250266 -> env
251267 -> Request
252− -> IO (RouteResult a, IO ())
253−runDelayed Delayed{..} env req = do
254− cleanupRef <- newIORef (return ())
255− routeRes <- runDelayedIO
268+ -> CleanupRef
269+ -> IO (RouteResult a)
270+runDelayed Delayed{..} env req cleanupRef =
271+ runDelayedIO
256272 (do c <- capturesD env
257273 methodD
258274 a <- authD
259275 b <- bodyD
260− liftIO (writeIORef cleanupRef $ cleanupD b)
261− DelayedIO $ \ r -> return (serverD c a b r)
276+ DelayedIO $ \ r _cleanup -> return (serverD c a b r)
262277 )
263278 req
264− fmap (routeRes,) $ readIORef cleanupRef
279+ cleanupRef
265280
266281-- | Runs a delayed server and the resulting action.
267282-- Takes a continuation that lets us send a response.
@@ -273,10 +288,11 @@ runAction :: Delayed env (Handler a)
273288 -> (RouteResult Response -> IO r)
274289 -> (a -> RouteResult Response)
275290 -> IO r
276−runAction action env req respond k =
277− bracket (runDelayed action env req)
278− snd
279− (\(res, _) -> go res >>= respond)
291+runAction action env req respond k = do
292+ cleanupRef <- newCleanupRef
293+ bracket (runDelayed action env req cleanupRef)
294+ (const $ runCleanup cleanupRef)
295+ (go >=> respond)
280296
281297 where
282298 go (Fail e) = return $ Fail e