From 81a876c3e3b0795911cb35fd605e7099e3567719 Mon Sep 17 00:00:00 2001 From: Alp Mestanogullari Date: Mon, 24 Oct 2016 19:21:28 +0200 Subject: [PATCH] use an ioref to store clean up actions instead of a field in Delayed, allowing early clean up registration --- servant-server/src/Servant/Server/Internal.hs | 8 ++- .../Server/Internal/RoutingApplication.hs | 72 +++++++++++-------- 2 files changed, 50 insertions(+), 30 deletions(-) diff --git a/servant-server/src/Servant/Server/Internal.hs b/servant-server/src/Servant/Server/Internal.hs index 2e9b3b2d..3ee91af3 100644 --- a/servant-server/src/Servant/Server/Internal.hs +++ b/servant-server/src/Servant/Server/Internal.hs @@ -399,8 +399,12 @@ instance HasServer Raw context where type ServerT Raw m = Application route Proxy _ rawApplication = RawRouter $ \ env request respond -> do - (r, cleanup) <- runDelayed rawApplication env request - go r request respond `finally` cleanup + -- note: a Raw application doesn't register any cleanup + -- but for the sake of consistency, we nonetheless run + -- the cleanup once its done + cleanupRef <- newCleanupRef + r <- runDelayed rawApplication env request cleanupRef + go r request respond `finally` runCleanup cleanupRef where go r request respond = case r of Route app -> app request (respond . Route) diff --git a/servant-server/src/Servant/Server/Internal/RoutingApplication.hs b/servant-server/src/Servant/Server/Internal/RoutingApplication.hs index 107cd798..785bac26 100644 --- a/servant-server/src/Servant/Server/Internal/RoutingApplication.hs +++ b/servant-server/src/Servant/Server/Internal/RoutingApplication.hs @@ -10,10 +10,10 @@ module Servant.Server.Internal.RoutingApplication where import Control.Exception (bracket) -import Control.Monad (ap, liftM) +import Control.Monad (ap, liftM, (>=>)) import Control.Monad.Trans (MonadIO(..)) import Control.Monad.Trans.Except (runExceptT) -import Data.IORef (newIORef, readIORef, writeIORef) +import Data.IORef (newIORef, readIORef, writeIORef, IORef, atomicModifyIORef) import Network.Wai (Application, Request, Response, ResponseReceived) import Prelude () @@ -106,10 +106,6 @@ data Delayed env c where , authD :: DelayedIO auth , bodyD :: DelayedIO body , serverD :: captures -> auth -> body -> Request -> RouteResult c - , cleanupD :: body -> IO () - -- not in DelayedIO because: - -- - most likely should not depend on the request - -- - simpler } -> Delayed env c instance Functor (Delayed env) where @@ -119,12 +115,33 @@ instance Functor (Delayed env) where , .. } -- Note [Existential Record Update] +-- | A mutable cleanup action +newtype CleanupRef = CleanupRef (IORef (IO ())) + +newCleanupRef :: IO CleanupRef +newCleanupRef = CleanupRef <$> newIORef (return ()) + +-- | Add a clean up action to a 'CleanupRef' +addCleanup' :: IO () -> CleanupRef -> IO () +addCleanup' act (CleanupRef ref) = atomicModifyIORef ref (\act' -> (act' >> act, ())) + +addCleanup :: IO () -> DelayedIO () +addCleanup act = DelayedIO $ \_req cleanupRef -> + addCleanup' act cleanupRef >> return (Route ()) + +-- | Run all the clean up actions registered in +-- a 'CleanupRef'. +runCleanup :: CleanupRef -> IO () +runCleanup (CleanupRef ref) = do + act <- readIORef ref + act + -- | Computations used in a 'Delayed' can depend on the -- incoming 'Request', may perform 'IO, and result in a -- 'RouteResult, meaning they can either suceed, fail -- (with the possibility to recover), or fail fatally. -- -newtype DelayedIO a = DelayedIO { runDelayedIO :: Request -> IO (RouteResult a) } +newtype DelayedIO a = DelayedIO { runDelayedIO :: Request -> CleanupRef -> IO (RouteResult a) } instance Functor DelayedIO where fmap = liftM @@ -134,36 +151,36 @@ instance Applicative DelayedIO where (<*>) = ap instance Monad DelayedIO where - return x = DelayedIO (const $ return (Route x)) + return x = DelayedIO (\_req _cleanup -> return (Route x)) DelayedIO m >>= f = - DelayedIO $ \ req -> do - r <- m req + DelayedIO $ \ req cl -> do + r <- m req cl case r of Fail e -> return $ Fail e FailFatal e -> return $ FailFatal e - Route a -> runDelayedIO (f a) req + Route a -> runDelayedIO (f a) req cl instance MonadIO DelayedIO where - liftIO m = DelayedIO (const $ Route <$> m) + liftIO m = DelayedIO (\_req _cl -> Route <$> m) -- | A 'Delayed' without any stored checks. emptyDelayed :: RouteResult a -> Delayed env a emptyDelayed result = - Delayed (const r) r r r (\ _ _ _ _ -> result) (const $ return ()) + Delayed (const r) r r r (\ _ _ _ _ -> result) where r = return () -- | Fail with the option to recover. delayedFail :: ServantErr -> DelayedIO a -delayedFail err = DelayedIO (const $ return $ Fail err) +delayedFail err = DelayedIO (\_req _cleanup -> return $ Fail err) -- | Fail fatally, i.e., without any option to recover. delayedFailFatal :: ServantErr -> DelayedIO a -delayedFailFatal err = DelayedIO (const $ return $ FailFatal err) +delayedFailFatal err = DelayedIO (\_req _cleanup -> return $ FailFatal err) -- | Gain access to the incoming request. withRequest :: (Request -> DelayedIO a) -> DelayedIO a -withRequest f = DelayedIO (\ req -> runDelayedIO (f req) req) +withRequest f = DelayedIO (\ req cl -> runDelayedIO (f req) req cl) -- | Add a capture to the end of the capture block. addCapture :: Delayed env (a -> b) @@ -205,7 +222,6 @@ addBodyCheck Delayed{..} new = Delayed { bodyD = (,) <$> bodyD <*> new , serverD = \ c a (z, v) req -> ($ v) <$> serverD c a z req - , cleanupD = cleanupD . fst -- not sure that's right , .. } -- Note [Existential Record Update] @@ -248,19 +264,18 @@ passToServer Delayed{..} x = runDelayed :: Delayed env a -> env -> Request - -> IO (RouteResult a, IO ()) -runDelayed Delayed{..} env req = do - cleanupRef <- newIORef (return ()) - routeRes <- runDelayedIO + -> CleanupRef + -> IO (RouteResult a) +runDelayed Delayed{..} env req cleanupRef = + runDelayedIO (do c <- capturesD env methodD a <- authD b <- bodyD - liftIO (writeIORef cleanupRef $ cleanupD b) - DelayedIO $ \ r -> return (serverD c a b r) + DelayedIO $ \ r _cleanup -> return (serverD c a b r) ) req - fmap (routeRes,) $ readIORef cleanupRef + cleanupRef -- | Runs a delayed server and the resulting action. -- Takes a continuation that lets us send a response. @@ -272,10 +287,11 @@ runAction :: Delayed env (Handler a) -> (RouteResult Response -> IO r) -> (a -> RouteResult Response) -> IO r -runAction action env req respond k = - bracket (runDelayed action env req) - snd - (\(res, _) -> go res >>= respond) +runAction action env req respond k = do + cleanupRef <- newCleanupRef + bracket (runDelayed action env req cleanupRef) + (const $ runCleanup cleanupRef) + (go >=> respond) where go (Fail e) = return $ Fail e