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

@@ -401,8 +401,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
+    -- note: a Raw application doesn't register any cleanup
-    go r request respond `finally` 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)

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                            ()
@@ -107,10 +107,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
@@ -120,12 +116,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
@@ -135,36 +152,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
+    DelayedIO $ \ req cl -> do
-      r <- m req
+      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)
@@ -206,7 +223,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]
 
@@ -249,19 +265,18 @@ passToServer Delayed{..} x =
 runDelayed :: Delayed env a
            -> env
            -> Request
-           -> IO (RouteResult a, IO ())
+           -> CleanupRef
-runDelayed Delayed{..} env req = do
+           -> IO (RouteResult a)
-  cleanupRef <- newIORef (return ())
+runDelayed Delayed{..} env req cleanupRef =
-  routeRes <- runDelayedIO
+  runDelayedIO
     (do c <- capturesD env
         methodD
         a <- authD
         b <- bodyD
-        liftIO (writeIORef cleanupRef $ cleanupD b)
+        DelayedIO $ \ r _cleanup -> return (serverD c a b r)
-        DelayedIO $ \ r -> 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.
@@ -273,10 +288,11 @@ runAction :: Delayed env (Handler a)
           -> (RouteResult Response -> IO r)
           -> (a -> RouteResult Response)
           -> IO r
-runAction action env req respond k =
+runAction action env req respond k = do
-  bracket (runDelayed action env req)
+  cleanupRef <- newCleanupRef
-          snd
+  bracket (runDelayed action env req cleanupRef)
-          (\(res, _) -> go res >>= respond)
+          (const $ runCleanup cleanupRef)
+          (go >=> respond)
 
   where
     go (Fail e)      = return $ Fail e