{-
  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.
-}


module Glean.Util.Warden
  ( Warden
  , withWarden
  , create
  , shutdown
  , spawn
  , spawn_
  , spawnMask
  , spawnDaemon
  , adopt
  ) where

import Control.Concurrent (forkIO)
import Control.Concurrent.Async (Async)
import qualified Control.Concurrent.Async as Async
import Control.Concurrent.MVar
import Control.Exception
import Control.Monad
import Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
import System.IO (fixIO)

import Util.Log

-- | A 'Warden' is an owner of 'Async's which cancels them on 'shutdown'.
--
-- 'Nothing' in the MVar means the 'Warden' has been shut down.
newtype Warden = Warden (MVar (Maybe (HashSet (Async ()))))

-- | Run the action with a new 'Warden', shutting it down when the action
-- exits.
withWarden :: (Warden -> IO a) -> IO a
withWarden = bracket create shutdown

-- | Create a new 'Warden'.
create :: IO Warden
create = Warden <$> newMVar (Just mempty)

-- | Shutdown a 'Warden', cancelling all owned threads. Subsequent calls to
-- 'spawn' and 'shutdown' will be no-ops.
shutdown :: Warden -> IO ()
shutdown (Warden v) = do
  r <- swapMVar v Nothing
  mapM_ (Async.mapConcurrently_ Async.cancel) r

forget :: Warden -> Async a -> IO ()
forget (Warden v) async = modifyMVar_ v $ \x -> case x of
  Just xs -> return $! Just $! HashSet.delete (void async) xs
  Nothing -> return Nothing

-- | Spawn a thread with masked exceptions and pass an unmask function to the
-- action.
spawnMask :: Warden -> ((forall b. IO b -> IO b) -> IO a) -> IO (Async a)
spawnMask (Warden v) action = modifyMVar v $ \r -> case r of
  Just asyncs -> do
    -- Create a new thread which removes itself from the 'HashSet' when it
    -- exits.
    this <- fixIO $ \this -> mask_ $ Async.asyncWithUnmask $ \unmask ->
      action unmask `finally` forget (Warden v) this
    return (Just $ HashSet.insert (void this) asyncs, this)
  Nothing -> (,) Nothing <$> Async.async (throwIO Async.AsyncCancelled)

-- | Spawn a new thread owned by the 'Warden'.
spawn :: Warden -> IO a -> IO (Async a)
spawn warden action = spawnMask warden $ \unmask -> unmask action

-- | Spawn a new thread owned by the 'Warden'.
spawn_ :: Warden -> IO () -> IO ()
spawn_ w = void . spawn w

-- | Run the given IO action repeatedly in a separate thread, catching
-- and logging any exceptions (Sync or Async) that it throws. Ends
-- at 'Warden' shutdown.
spawnDaemon :: Warden -> String -> IO () -> IO ()
spawnDaemon w name io =
  void $ spawnMask w $ \restore ->
    forever $
      handle (\ex -> case ex of
        _ | Just Async.AsyncCancelled{} <- fromException ex -> throwIO ex
          | otherwise -> logError (name <> ": " <> show ex)) $
        restore io

-- | Adopt an existing thread which will be cancelled when 'Warden' shut down.
-- If the 'Warden' has already been shut down, the thread will be cancelled
-- immediately.
adopt :: Warden -> Async a -> IO ()
adopt (Warden v) async = modifyMVar_ v $ \r -> case r of
  Just asyncs -> do
    void $ forkIO $ do
      void $ Async.waitCatch async
      forget (Warden v) async
    return $ Just $ HashSet.insert (void async) asyncs
  Nothing -> do
    Async.cancel async
    return Nothing