{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE CPP                 #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Util.Control.Exception
  ( -- * Catching all exceptions safely
    catchAll
  , handleAll
  , tryAll
    -- * Exception predicates
  , isSyncException
  , isAsyncException
    -- * Other utilities
  , throwLeftIO
  , throwLeftExceptionIO
  , tryBracket
  , tryFinally
  , onSomeException
  , afterwards
  , swallow
  , logExceptions
  ) where

#if __GLASGOW_HASKELL__ == 804
import Control.Exception ( SomeAsyncException(..) )
#endif
import Control.Exception.Lifted
import Control.Monad
import Control.Monad.Trans.Control

import GHC.Stack (HasCallStack, withFrozenCallStack)

import Util.Log

-- | Catch all exceptions *except* asynchronous exceptions
-- (technically, children of 'SomeAsyncException').  Catching
-- asynchronous exceptions is almost never what you want to do: it can
-- result in ignoring 'ThreadKilled' which can lead to deadlock (see
-- <https://our.internmc.facebook.com/intern/diff/D4745709/
-- D4745709>).
--
-- Use this instead of the raw 'catch' when catching 'SomeException'.
--
catchAll :: MonadBaseControl IO m => m a -> (SomeException -> m a) -> m a
catchAll action handler =
  action `catch` \ex ->
    case fromException ex of
      Just (_ :: SomeAsyncException) -> throwIO ex
      Nothing -> handler ex

-- | The "try" version of 'catchAll'
tryAll :: MonadBaseControl IO m => m a -> m (Either SomeException a)
tryAll action = (Right <$> action) `catchAll` (return . Left)

-- | Flipped version of 'catchAll'
handleAll :: MonadBaseControl IO m => (SomeException -> m a) -> m a -> m a
handleAll = flip catchAll

throwLeftIO :: Exception e => Either e a -> IO a
throwLeftIO = throwLeftExceptionIO id

throwLeftExceptionIO :: Exception e => (a -> e) -> Either a b -> IO b
throwLeftExceptionIO mkEx e = either (throwIO . mkEx) pure e

-- | Detect 'SomeAsyncException' wrapped exceptions versus all others
isSyncException :: Exception e => e -> Bool
isSyncException e = case fromException (toException e) of
  Just (SomeAsyncException _) -> False
  Nothing -> True

-- | Detect 'SomeAsyncException' wrapped exceptions versus all others
isAsyncException :: Exception e => e -> Bool
isAsyncException = not . isSyncException

-- | A variant of 'bracket' where the release action also gets to see whether
-- the inner action succeeded or threw an exception.
tryBracket
  :: IO a                                   -- ^ run first
  -> (a -> Either SomeException b -> IO ()) -- ^ run finally
  -> (a -> IO b)                            -- ^ run in between
  -> IO b
tryBracket before after inner =
  mask $ \restore -> do
    a <- before
    r <- restore (inner a) `catch` \ex -> do
      after a (Left ex)
      throwIO ex
    _ <- after a (Right r)
    return r

-- | A variant of 'finally' where the final action also gets to see whether
-- the first action succeeded or threw an exception.
tryFinally
  :: IO a                              -- ^ run first
  -> (Either SomeException a -> IO ()) -- ^ run finally
  -> IO a
tryFinally inner after =
  mask $ \restore -> do
    r <- restore inner `catch` \ex -> do
      after (Left ex)
      throwIO ex
    _ <- after (Right r)
    return r

-- | Execute an action and invoke a function if it throws any exception. The
-- exception is then rethrown. Any exceptions from the function are ignored
-- (but logged).
onSomeException :: HasCallStack => IO a -> (SomeException -> IO ()) -> IO a
onSomeException io f = io `catch` \exc -> do
  withFrozenCallStack $ swallow $ f exc
  throwIO exc

-- | Execute an action and do something with its result even if it throws a
-- synchronous exception. Any exceptions from the function are ignored
-- (but logged).
afterwards :: HasCallStack => IO a -> (Either SomeException a -> IO ()) -> IO a
afterwards io f = do
  r <- tryAll io
  withFrozenCallStack $ swallow $ f r
  case r of
    Right result -> return result
    Left exc -> throwIO exc

-- | Execute an action and drop its result or any synchronous
-- exception it throws.  Exceptions are logged.
swallow :: HasCallStack => IO a -> IO ()
swallow io = void io `catchAll` \exc -> do
  withFrozenCallStack $ logError $ "swallowing exception: " ++ show exc

-- | Log and rethrow all synchronous exceptions arising from an
-- IO computation.
logExceptions :: HasCallStack => (String -> String) -> IO a -> IO a
logExceptions f io = withFrozenCallStack $
  io `onSomeException` (logError . f . show)