make cleanup in Delayed more resistant to exceptions

Author: alpmestan

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

@@ -23,6 +23,7 @@ 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
@@ -400,11 +401,13 @@ instance HasServer Raw context where
   type ServerT Raw m = Application
 
   route Proxy _ rawApplication = RawRouter $ \ env request respond -> do
-    (r, _) <- runDelayed rawApplication env request
-    case r of
-      Route app   -> app request (respond . Route)
-      Fail a      -> respond $ Fail a
-      FailFatal e -> respond $ FailFatal e
+    (r, cleanup) <- runDelayed rawApplication env request
+    go r request respond `finally` cleanup
+
+    where go r request respond = case r of
+            Route app   -> app request (respond . Route)
+            Fail a      -> respond $ Fail a
+            FailFatal e -> respond $ FailFatal e
 
 -- | If you use 'ReqBody' in one of the endpoints for your API,
 -- this automatically requires your server-side handler to be a function

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

@@ -9,6 +9,7 @@
 {-# LANGUAGE TupleSections              #-}
 module Servant.Server.Internal.RoutingApplication where
 
+import           Control.Exception                  (bracket)
 import           Control.Monad                      (ap, liftM)
 import           Control.Monad.Trans                (MonadIO(..))
 import           Control.Monad.Trans.Except         (runExceptT)
@@ -272,11 +273,11 @@ runAction :: Delayed env (Handler a)
           -> (RouteResult Response -> IO r)
           -> (a -> RouteResult Response)
           -> IO r
-runAction action env req respond k = do
-  (routeResult, cleanup) <- runDelayed action env req
-  resp <- go routeResult
-  cleanup
-  respond resp
+runAction action env req respond k =
+  bracket (runDelayed action env req)
+          snd
+          (\(res, _) -> go res >>= respond)
+
   where
     go (Fail e)      = return $ Fail e
     go (FailFatal e) = return $ FailFatal e