Use resourcet for resource managment

servant-server/servant-server.cabal

@@ -69,6 +69,7 @@ library
       , string-conversions >= 0.3  && < 0.5
       , system-filepath    >= 0.4  && < 0.5
       , filepath           >= 1    && < 1.5
+      , resourcet          >= 1.1.6 && <1.2
       , text               >= 1.2  && < 1.3
       , transformers       >= 0.3  && < 0.6
       , transformers-base  >= 0.4.4 && < 0.5
@@ -127,19 +128,20 @@ test-suite spec
     , hspec == 2.*
     , hspec-wai >= 0.8 && <0.9
     , http-types
+    , mtl
     , network >= 2.6
-    , QuickCheck
     , parsec
+    , QuickCheck
+    , resourcet
     , safe
     , servant
     , servant-server
-    , string-conversions
     , should-not-typecheck == 2.1.*
+    , string-conversions
     , temporary
     , text
     , transformers
     , transformers-compat
-    , mtl
     , wai
     , wai-extra
     , warp

servant-server/src/Servant/Server/Internal.hs

@@ -23,7 +23,6 @@ module Servant.Server.Internal
   , module Servant.Server.Internal.ServantErr
   ) where
 
-import           Control.Exception          (finally)
 import           Control.Monad.Trans        (liftIO)
 import qualified Data.ByteString            as B
 import qualified Data.ByteString.Char8      as BC8
@@ -404,9 +403,8 @@ instance HasServer Raw context where
     -- 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
+    r <- runDelayed rawApplication env request
+    go r request respond
 
     where go r request respond = case r of
             Route app   -> app request (respond . Route)

servant-server/src/Servant/Server/Internal/RoutingApplication.hs

@@ -1,24 +1,32 @@
 {-# LANGUAGE CPP                        #-}
 {-# LANGUAGE DeriveFunctor              #-}
-{-# LANGUAGE OverloadedStrings          #-}
-{-# LANGUAGE TypeOperators              #-}
+{-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GADTs                      #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE KindSignatures             #-}
+{-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE OverloadedStrings          #-}
 {-# LANGUAGE RecordWildCards            #-}
 {-# LANGUAGE StandaloneDeriving         #-}
 {-# LANGUAGE TupleSections              #-}
+{-# LANGUAGE TypeFamilies               #-}
+{-# LANGUAGE TypeOperators              #-}
+{-# LANGUAGE UndecidableInstances       #-}
 module Servant.Server.Internal.RoutingApplication where
 
-import           Control.Exception                  (finally)
 import           Control.Monad                      (ap, liftM)
-import           Control.Monad.Trans                (MonadIO(..))
-import           Data.IORef                         (newIORef, readIORef, IORef, atomicModifyIORef)
-import           Network.Wai                        (Application, Request,
-                                                     Response, ResponseReceived)
+import           Control.Monad.Base                 (MonadBase (..))
+import           Control.Monad.Catch                (MonadThrow (..))
+import           Control.Monad.Reader               (MonadReader (..), ReaderT, runReaderT)
+import           Control.Monad.Trans                (MonadIO (..), MonadTrans (..))
+import           Control.Monad.Trans.Control        (ComposeSt, MonadBaseControl (..), MonadTransControl (..),
+                                                     defaultLiftBaseWith, defaultRestoreM)
+import           Control.Monad.Trans.Resource       (MonadResource (..), ResourceT, runResourceT)
+import           Network.Wai                        (Application, Request, Response, ResponseReceived)
 import           Prelude                            ()
 import           Prelude.Compat
-import           Servant.Server.Internal.ServantErr
 import           Servant.Server.Internal.Handler
+import           Servant.Server.Internal.ServantErr
 
 type RoutingApplication =
      Request -- ^ the request, the field 'pathInfo' may be modified by url routing
@@ -32,6 +40,58 @@ data RouteResult a =
   | Route !a
   deriving (Eq, Show, Read, Functor)
 
+instance Applicative RouteResult where
+    pure = return
+    (<*>) = ap
+
+instance Monad RouteResult where
+    return = Route
+    Route a     >>= f = f a
+    Fail e      >>= _ = Fail e
+    FailFatal e >>= _ = FailFatal e
+
+newtype RouteResultT m a = RouteResultT { runRouteResultT :: m (RouteResult a) }
+  deriving (Functor)
+
+-- As we write these instances, we get instances for `DelayedIO` with GND.
+instance MonadTrans RouteResultT where
+    lift = RouteResultT . liftM Route
+
+instance (Functor m, Monad m) => Applicative (RouteResultT m) where
+    pure  = return
+    (<*>) = ap
+
+instance Monad m => Monad (RouteResultT m) where
+    return = RouteResultT . return . Route
+    m >>= k = RouteResultT $ do
+        a <- runRouteResultT m
+        case a of
+            Fail e      -> return $ Fail e
+            FailFatal e -> return $ FailFatal e
+            Route b     -> runRouteResultT (k b)
+
+instance MonadIO m => MonadIO (RouteResultT m) where
+    liftIO = lift . liftIO
+
+instance MonadBase b m => MonadBase b (RouteResultT m) where
+    liftBase = lift . liftBase
+
+instance MonadBaseControl b m => MonadBaseControl b (RouteResultT m) where
+    type StM (RouteResultT m) a = ComposeSt RouteResultT m a
+    liftBaseWith = defaultLiftBaseWith
+    restoreM     = defaultRestoreM
+
+instance MonadTransControl RouteResultT where
+    type StT RouteResultT a = RouteResult a
+    liftWith f = RouteResultT $ liftM return $ f $ runRouteResultT
+    restoreT = RouteResultT
+
+instance MonadThrow m => MonadThrow (RouteResultT m) where
+    throwM = lift . throwM
+
+-- instance MonadCatch m => MonadCatch (RouteResultT m) where
+-- instance MonadError ServantErr (RouteResultT m) where
+
 toApplication :: RoutingApplication -> Application
 toApplication ra request respond = ra request routingRespond
  where
@@ -115,53 +175,30 @@ 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 -> CleanupRef -> IO (RouteResult a) }
+newtype DelayedIO a = DelayedIO { runDelayedIO' :: ReaderT Request (ResourceT (RouteResultT IO)) a }
+  deriving
+    ( Functor, Applicative, Monad
+    , MonadIO, MonadReader Request
+    , MonadBase IO
+    , MonadThrow
+    , MonadResource
+    )
 
-instance Functor DelayedIO where
-  fmap = liftM
+returnRouteResult :: RouteResult a -> DelayedIO a
+returnRouteResult x = DelayedIO $ lift . lift $ RouteResultT . return $ x
 
-instance Applicative DelayedIO where
-  pure = return
-  (<*>) = ap
+instance MonadBaseControl IO DelayedIO where
+    type StM DelayedIO a = StM (ReaderT Request (ResourceT (RouteResultT IO))) a
+    liftBaseWith f = DelayedIO $ liftBaseWith $ \g -> f (g . runDelayedIO')
+    restoreM       = DelayedIO . restoreM
 
-instance Monad DelayedIO where
-  return x = DelayedIO (\_req _cleanup -> return (Route x))
-  DelayedIO m >>= f =
-    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 cl
-
-instance MonadIO DelayedIO where
-  liftIO m = DelayedIO (\_req _cl -> Route <$> m)
+runDelayedIO :: DelayedIO a -> Request -> IO (RouteResult a)
+runDelayedIO m req = runRouteResultT $ runResourceT $ runReaderT (runDelayedIO' m) req
 
 -- | A 'Delayed' without any stored checks.
 emptyDelayed :: RouteResult a -> Delayed env a
@@ -172,15 +209,17 @@ emptyDelayed result =
 
 -- | Fail with the option to recover.
 delayedFail :: ServantErr -> DelayedIO a
-delayedFail err = DelayedIO (\_req _cleanup -> return $ Fail err)
+delayedFail err = returnRouteResult $ Fail err
 
 -- | Fail fatally, i.e., without any option to recover.
 delayedFailFatal :: ServantErr -> DelayedIO a
-delayedFailFatal err = DelayedIO (\_req _cleanup -> return $ FailFatal err)
+delayedFailFatal err = returnRouteResult $ FailFatal err
 
 -- | Gain access to the incoming request.
 withRequest :: (Request -> DelayedIO a) -> DelayedIO a
-withRequest f = DelayedIO (\ req cl -> runDelayedIO (f req) req cl)
+withRequest f = do
+    req <- ask
+    f req
 
 -- | Add a capture to the end of the capture block.
 addCapture :: Delayed env (a -> b)
@@ -264,18 +303,17 @@ passToServer Delayed{..} x =
 runDelayed :: Delayed env a
            -> env
            -> Request
-           -> CleanupRef
            -> IO (RouteResult a)
-runDelayed Delayed{..} env req cleanupRef =
+runDelayed Delayed{..} env req  =
   runDelayedIO
     (do c <- capturesD env
         methodD
         a <- authD
         b <- bodyD
-        DelayedIO $ \ r _cleanup -> return (serverD c a b r)
+        r <- ask
+        returnRouteResult (serverD c a b r)
     )
     req
-    cleanupRef
 
 -- | Runs a delayed server and the resulting action.
 -- Takes a continuation that lets us send a response.
@@ -288,10 +326,7 @@ runAction :: Delayed env (Handler a)
           -> (a -> RouteResult Response)
           -> IO r
 runAction action env req respond k = do
-  cleanupRef <- newCleanupRef
-  (runDelayed action env req cleanupRef >>= go >>= respond)
-    `finally` runCleanup cleanupRef
-
+    runDelayed action env req >>= go >>= respond
   where
     go (Fail e)      = return $ Fail e
     go (FailFatal e) = return $ FailFatal e

servant-server/test/Servant/Server/Internal/RoutingApplicationSpec.hs

@@ -4,6 +4,7 @@ import Prelude ()
 import Prelude.Compat
 
 import Control.Exception hiding (Handler)
+import Control.Monad.Trans.Resource (register)
 import Control.Monad.IO.Class
 import Data.Maybe (isJust)
 import Data.IORef
@@ -13,23 +14,20 @@ import Test.Hspec
 
 import System.IO.Unsafe (unsafePerformIO)
 
-ok :: IO (RouteResult ())
-ok = return (Route ())
-
 -- Let's not write to the filesystem
 delayedTestRef :: IORef (Maybe String)
 delayedTestRef = unsafePerformIO $ newIORef Nothing
 
 delayed :: DelayedIO () -> RouteResult (Handler ()) -> Delayed () (Handler ())
 delayed body srv = Delayed
-  { capturesD = \() -> DelayedIO $ \_req _cl -> ok
-  , methodD = DelayedIO $ \_req_ _cl -> ok
-  , authD = DelayedIO $ \_req _cl -> ok
-  , bodyD = do
+  { capturesD = \_ -> return ()
+  , methodD   = return ()
+  , authD     = return ()
+  , bodyD     = do
       liftIO (writeIORef delayedTestRef (Just "hia") >> putStrLn "garbage created")
-      addCleanup (writeIORef delayedTestRef Nothing >> putStrLn "garbage collected")
+      _ <- register (writeIORef delayedTestRef Nothing >> putStrLn "garbage collected")
       body
-  , serverD = \() () _body _req -> srv
+  , serverD   = \() () _body _req -> srv
   }
 
 simpleRun :: Delayed () (Handler ())