{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

{- | This is an almost direct copy of [Control.Reaper](https://hackage.haskell.org/package/auto-update/docs/Control-Reaper.html)
 from the /auto-update/ package. The salient difference is that this module allows us to define cleanup threads in arbitrary
 monads using 'MonadUnliftIO'.

 This module provides the ability to create reapers: dedicated cleanup
 threads. These threads will automatically spawn and die based on the
 presence of a workload to process on. Example uses include:

 * Killing long-running jobs
 * Closing unused connections in a connection pool
 * Pruning a cache of old items (see example below)

 For real-world usage, search the <https://github.com/yesodweb/wai WAI family of packages>
 for imports of "Control.Reaper".
-}
module UnliftIO.Reaper
  ( -- * Example: Regularly cleaning a cache
    -- $example1

    -- * Settings
    ReaperSettings
  , defaultReaperSettings

    -- * Accessors
  , reaperAction
  , reaperDelay
  , reaperCons
  , reaperNull
  , reaperEmpty
  , reaperThreadName

    -- * Type
  , Reaper
  , reaperAdd
  , reaperRead
  , reaperModify
  , reaperStop
  , reaperKill

    -- * Creation
  , mkReaper

    -- * Helper
  , mkListAction
  )
where

import Control.Monad (forM_, join)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Functor.Identity
import GHC.Conc.Sync (labelThread)
import UnliftIO (MonadUnliftIO)
import UnliftIO.Concurrent (ThreadId, forkIO, killThread, threadDelay)
import UnliftIO.Exception (mask_)
import UnliftIO.IORef (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef)
import UnliftIO.Reaper.Internal

{- | Settings for creating a reaper. This type has two parameters:
 @workload@ gives the entire workload, whereas @item@ gives an
 individual piece of the queue. A common approach is to have @workload@
 be a list of @item@s. This is encouraged by 'defaultReaperSettings' and
 'mkListAction'.

 @since 0.1.0
-}
data ReaperSettings m workload item = ReaperSettings
  { forall (m :: * -> *) workload item.
ReaperSettings m workload item
-> workload -> m (workload -> workload)
reaperAction :: workload -> m (workload -> workload)
  -- ^ The action to perform on a workload. The result of this is a
  -- \"workload modifying\" function. In the common case of using lists,
  -- the result should be a difference list that prepends the remaining
  -- workload to the temporary workload. The temporary workload here
  -- refers to items added to the workload while the reaper action is
  -- running. For help with setting up such an action, see 'mkListAction'.
  --
  -- Default: do nothing with the workload, and then prepend it to the
  -- temporary workload. This is incredibly useless; you should
  -- definitely override this default.
  --
  -- @since 0.1.0
  , forall (m :: * -> *) workload item.
ReaperSettings m workload item -> Int
reaperDelay :: {-# UNPACK #-} !Int
  -- ^ Number of microseconds to delay between calls of 'reaperAction'.
  --
  -- Default: 30 seconds.
  --
  -- @since 0.1.0
  , forall (m :: * -> *) workload item.
ReaperSettings m workload item -> item -> workload -> workload
reaperCons :: item -> workload -> workload
  -- ^ Add an item onto a workload.
  --
  -- Default: list consing.
  --
  -- @since 0.1.0
  , forall (m :: * -> *) workload item.
ReaperSettings m workload item -> workload -> Bool
reaperNull :: workload -> Bool
  -- ^ Check if a workload is empty, in which case the worker thread
  -- will shut down.
  --
  -- Default: 'null'.
  --
  -- @since 0.1.0
  , forall (m :: * -> *) workload item.
ReaperSettings m workload item -> workload
reaperEmpty :: workload
  -- ^ An empty workload.
  --
  -- Default: empty list.
  --
  -- @since 0.1.0
  , forall (m :: * -> *) workload item.
ReaperSettings m workload item -> String
reaperThreadName :: String
  -- ^ Label of the thread spawned by the reaper.
  --
  -- Default: @"Reaper"@.
  --
  -- @since 0.1.0
  }

{- | Default @ReaperSettings@ value, biased towards having a list of work
 items.

 @since 0.1.0
-}
defaultReaperSettings :: ReaperSettings Identity [item] item
defaultReaperSettings :: forall item. ReaperSettings Identity [item] item
defaultReaperSettings =
  ReaperSettings
    { reaperAction :: [item] -> Identity ([item] -> [item])
reaperAction = \[item]
wl -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ([item]
wl forall a. [a] -> [a] -> [a]
++)
    , reaperDelay :: Int
reaperDelay = Int
30000000
    , reaperCons :: item -> [item] -> [item]
reaperCons = (:)
    , reaperNull :: [item] -> Bool
reaperNull = forall (t :: * -> *) a. Foldable t => t a -> Bool
null
    , reaperEmpty :: [item]
reaperEmpty = []
    , reaperThreadName :: String
reaperThreadName = String
"Reaper"
    }

-- | State of reaper.
data State workload
  = -- | No reaper thread
    NoReaper
  | -- | The current jobs
    Workload !workload

{- | Create a reaper addition function. This function can be used to add
 new items to the workload. Spawning of reaper threads will be handled
 for you automatically.

 @since 0.1.0
-}
mkReaper :: MonadUnliftIO m => ReaperSettings m workload item -> m (Reaper m workload item)
mkReaper :: forall (m :: * -> *) workload item.
MonadUnliftIO m =>
ReaperSettings m workload item -> m (Reaper m workload item)
mkReaper settings :: ReaperSettings m workload item
settings@ReaperSettings {workload
Int
String
workload -> m (workload -> workload)
workload -> Bool
item -> workload -> workload
reaperThreadName :: String
reaperEmpty :: workload
reaperNull :: workload -> Bool
reaperCons :: item -> workload -> workload
reaperDelay :: Int
reaperAction :: workload -> m (workload -> workload)
reaperThreadName :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> String
reaperEmpty :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> workload
reaperNull :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> workload -> Bool
reaperCons :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> item -> workload -> workload
reaperDelay :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> Int
reaperAction :: forall (m :: * -> *) workload item.
ReaperSettings m workload item
-> workload -> m (workload -> workload)
..} = do
  IORef (State workload)
stateRef <- forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef forall workload. State workload
NoReaper
  IORef (Maybe ThreadId)
tidRef <- forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef forall a. Maybe a
Nothing
  forall (m :: * -> *) a. Monad m => a -> m a
return
    Reaper
      { reaperAdd :: item -> m ()
reaperAdd = forall (m :: * -> *) workload item.
MonadUnliftIO m =>
ReaperSettings m workload item
-> IORef (State workload) -> IORef (Maybe ThreadId) -> item -> m ()
add ReaperSettings m workload item
settings IORef (State workload)
stateRef IORef (Maybe ThreadId)
tidRef
      , reaperRead :: m workload
reaperRead = forall {m :: * -> *}.
MonadIO m =>
IORef (State workload) -> m workload
readRef IORef (State workload)
stateRef
      , reaperModify :: (workload -> workload) -> m workload
reaperModify = forall {m :: * -> *}.
MonadIO m =>
IORef (State workload) -> (workload -> workload) -> m workload
modifyRef IORef (State workload)
stateRef
      , reaperStop :: m workload
reaperStop = forall {m :: * -> *}.
MonadIO m =>
IORef (State workload) -> m workload
stop IORef (State workload)
stateRef
      , reaperKill :: m ()
reaperKill = forall {m :: * -> *} {t :: * -> *}.
(MonadIO m, Foldable t) =>
IORef (t ThreadId) -> m ()
kill IORef (Maybe ThreadId)
tidRef
      }
  where
    readRef :: IORef (State workload) -> m workload
readRef IORef (State workload)
stateRef = do
      State workload
mx <- forall (m :: * -> *) a. MonadIO m => IORef a -> m a
readIORef IORef (State workload)
stateRef
      case State workload
mx of
        State workload
NoReaper -> forall (m :: * -> *) a. Monad m => a -> m a
return workload
reaperEmpty
        Workload workload
wl -> forall (m :: * -> *) a. Monad m => a -> m a
return workload
wl
    modifyRef :: IORef (State workload) -> (workload -> workload) -> m workload
modifyRef IORef (State workload)
stateRef workload -> workload
modifier = forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' IORef (State workload)
stateRef forall a b. (a -> b) -> a -> b
$ \case
      State workload
NoReaper ->
        (forall workload. State workload
NoReaper, workload
reaperEmpty)
      Workload workload
wl ->
        let !wl' :: workload
wl' = workload -> workload
modifier workload
wl
        in  (forall workload. workload -> State workload
Workload workload
wl', workload
wl')
    stop :: IORef (State workload) -> m workload
stop IORef (State workload)
stateRef = forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' IORef (State workload)
stateRef forall a b. (a -> b) -> a -> b
$ \case
      State workload
NoReaper -> (forall workload. State workload
NoReaper, workload
reaperEmpty)
      Workload workload
x -> (forall workload. workload -> State workload
Workload workload
reaperEmpty, workload
x)
    kill :: IORef (t ThreadId) -> m ()
kill IORef (t ThreadId)
tidRef = do
      t ThreadId
mtid <- forall (m :: * -> *) a. MonadIO m => IORef a -> m a
readIORef IORef (t ThreadId)
tidRef
      forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t ThreadId
mtid forall (m :: * -> *). MonadIO m => ThreadId -> m ()
killThread

add ::
  MonadUnliftIO m =>
  ReaperSettings m workload item ->
  IORef (State workload) ->
  IORef (Maybe ThreadId) ->
  item ->
  m ()
add :: forall (m :: * -> *) workload item.
MonadUnliftIO m =>
ReaperSettings m workload item
-> IORef (State workload) -> IORef (Maybe ThreadId) -> item -> m ()
add settings :: ReaperSettings m workload item
settings@ReaperSettings {workload
Int
String
workload -> m (workload -> workload)
workload -> Bool
item -> workload -> workload
reaperThreadName :: String
reaperEmpty :: workload
reaperNull :: workload -> Bool
reaperCons :: item -> workload -> workload
reaperDelay :: Int
reaperAction :: workload -> m (workload -> workload)
reaperThreadName :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> String
reaperEmpty :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> workload
reaperNull :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> workload -> Bool
reaperCons :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> item -> workload -> workload
reaperDelay :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> Int
reaperAction :: forall (m :: * -> *) workload item.
ReaperSettings m workload item
-> workload -> m (workload -> workload)
..} IORef (State workload)
stateRef IORef (Maybe ThreadId)
tidRef item
item =
  forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
mask_ forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' IORef (State workload)
stateRef State workload -> (State workload, m ())
cons
  where
    cons :: State workload -> (State workload, m ())
cons State workload
NoReaper =
      let wl :: workload
wl = item -> workload -> workload
reaperCons item
item workload
reaperEmpty
      in  (forall workload. workload -> State workload
Workload workload
wl, forall (m :: * -> *) workload item.
MonadUnliftIO m =>
ReaperSettings m workload item
-> IORef (State workload) -> IORef (Maybe ThreadId) -> m ()
spawn ReaperSettings m workload item
settings IORef (State workload)
stateRef IORef (Maybe ThreadId)
tidRef)
    cons (Workload workload
wl) =
      let wl' :: workload
wl' = item -> workload -> workload
reaperCons item
item workload
wl
      in  (forall workload. workload -> State workload
Workload workload
wl', forall (m :: * -> *) a. Monad m => a -> m a
return ())

spawn ::
  MonadUnliftIO m =>
  ReaperSettings m workload item ->
  IORef (State workload) ->
  IORef (Maybe ThreadId) ->
  m ()
spawn :: forall (m :: * -> *) workload item.
MonadUnliftIO m =>
ReaperSettings m workload item
-> IORef (State workload) -> IORef (Maybe ThreadId) -> m ()
spawn ReaperSettings m workload item
settings IORef (State workload)
stateRef IORef (Maybe ThreadId)
tidRef = do
  ThreadId
tid <- forall (m :: * -> *). MonadUnliftIO m => m () -> m ThreadId
forkIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) workload item.
MonadUnliftIO m =>
ReaperSettings m workload item
-> IORef (State workload) -> IORef (Maybe ThreadId) -> m ()
reaper ReaperSettings m workload item
settings IORef (State workload)
stateRef IORef (Maybe ThreadId)
tidRef
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. ThreadId -> String -> IO ()
labelThread ThreadId
tid forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) workload item.
ReaperSettings m workload item -> String
reaperThreadName ReaperSettings m workload item
settings
  forall (m :: * -> *) a. MonadIO m => IORef a -> a -> m ()
writeIORef IORef (Maybe ThreadId)
tidRef forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just ThreadId
tid

reaper ::
  MonadUnliftIO m =>
  ReaperSettings m workload item ->
  IORef (State workload) ->
  IORef (Maybe ThreadId) ->
  m ()
reaper :: forall (m :: * -> *) workload item.
MonadUnliftIO m =>
ReaperSettings m workload item
-> IORef (State workload) -> IORef (Maybe ThreadId) -> m ()
reaper settings :: ReaperSettings m workload item
settings@ReaperSettings {workload
Int
String
workload -> m (workload -> workload)
workload -> Bool
item -> workload -> workload
reaperThreadName :: String
reaperEmpty :: workload
reaperNull :: workload -> Bool
reaperCons :: item -> workload -> workload
reaperDelay :: Int
reaperAction :: workload -> m (workload -> workload)
reaperThreadName :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> String
reaperEmpty :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> workload
reaperNull :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> workload -> Bool
reaperCons :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> item -> workload -> workload
reaperDelay :: forall (m :: * -> *) workload item.
ReaperSettings m workload item -> Int
reaperAction :: forall (m :: * -> *) workload item.
ReaperSettings m workload item
-> workload -> m (workload -> workload)
..} IORef (State workload)
stateRef IORef (Maybe ThreadId)
tidRef = do
  forall (m :: * -> *). MonadIO m => Int -> m ()
threadDelay Int
reaperDelay
  -- Getting the current jobs. Push an empty job to the reference.
  workload
wl <- forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' IORef (State workload)
stateRef forall {b}. State b -> (State workload, b)
swapWithEmpty
  -- Do the jobs. A function to merge the left jobs and
  -- new jobs is returned.
  !workload -> workload
merge <- workload -> m (workload -> workload)
reaperAction workload
wl
  -- Merging the left jobs and new jobs.
  -- If there is no jobs, this thread finishes.
  Bool
cont <- forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' IORef (State workload)
stateRef (forall {workload}.
(workload -> workload) -> State workload -> (State workload, Bool)
check workload -> workload
merge)
  if Bool
cont
    then forall (m :: * -> *) workload item.
MonadUnliftIO m =>
ReaperSettings m workload item
-> IORef (State workload) -> IORef (Maybe ThreadId) -> m ()
reaper ReaperSettings m workload item
settings IORef (State workload)
stateRef IORef (Maybe ThreadId)
tidRef
    else forall (m :: * -> *) a. MonadIO m => IORef a -> a -> m ()
writeIORef IORef (Maybe ThreadId)
tidRef forall a. Maybe a
Nothing
  where
    swapWithEmpty :: State b -> (State workload, b)
swapWithEmpty State b
NoReaper = forall a. HasCallStack => String -> a
error String
"Control.Reaper.reaper: unexpected NoReaper (1)"
    swapWithEmpty (Workload b
wl) = (forall workload. workload -> State workload
Workload workload
reaperEmpty, b
wl)

    check :: (workload -> workload) -> State workload -> (State workload, Bool)
check workload -> workload
_ State workload
NoReaper = forall a. HasCallStack => String -> a
error String
"Control.Reaper.reaper: unexpected NoReaper (2)"
    check workload -> workload
merge (Workload workload
wl)
      -- If there is no job, reaper is terminated.
      | workload -> Bool
reaperNull workload
wl' = (forall workload. State workload
NoReaper, Bool
False)
      -- If there are jobs, carry them out.
      | Bool
otherwise = (forall workload. workload -> State workload
Workload workload
wl', Bool
True)
      where
        wl' :: workload
wl' = workload -> workload
merge workload
wl

{- | A helper function for creating 'reaperAction' functions. You would
 provide this function with a function to process a single work item and
 return either a new work item, or @Nothing@ if the work item is
 expired.

 @since 0.1.0
-}
mkListAction ::
  Monad m =>
  (item -> m (Maybe item')) ->
  [item] ->
  m ([item'] -> [item'])
mkListAction :: forall (m :: * -> *) item item'.
Monad m =>
(item -> m (Maybe item')) -> [item] -> m ([item'] -> [item'])
mkListAction item -> m (Maybe item')
f =
  forall {c}. ([item'] -> c) -> [item] -> m ([item'] -> c)
go forall a. a -> a
id
  where
    go :: ([item'] -> c) -> [item] -> m ([item'] -> c)
go ![item'] -> c
front [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure [item'] -> c
front
    go ![item'] -> c
front (item
x : [item]
xs) = do
      Maybe item'
my <- item -> m (Maybe item')
f item
x
      let front' :: [item'] -> c
front' =
            case Maybe item'
my of
              Maybe item'
Nothing -> [item'] -> c
front
              Just item'
y -> [item'] -> c
front forall b c a. (b -> c) -> (a -> b) -> a -> c
. (item'
y forall a. a -> [a] -> [a]
:)
      ([item'] -> c) -> [item] -> m ([item'] -> c)
go [item'] -> c
front' [item]
xs

{- $example1
 In this example code, we use a 'Data.Map.Strict.Map' to cache fibonacci numbers, and a 'Reaper' to prune the cache.

 NOTE: When using this module as a cache you should keep in mind that while
 the reaper thread is active running your "reaperAction", the cache will
 appear empty to concurrently running threads.  Any newly created cache
 entries will be on the temporary worklist, and will merged back into the the
 main cache only once the "reaperAction" completes (together with the portion
 of the extant worklist that the @cleaner@ callback decided to retain).

 If you're looking for a cache that supports concurrent purging of stale
 items, but without exposing a transient empty cache during cleanup, this is
 not the cache implementation you need.  This module was primarily designed
 for cleaning up /stuck/ processes, or idle threads in a thread pool.  The cache
 use-case was not a primary design focus.

 The @main@ function first creates a 'Reaper', with fields to initialize the
 cache ('reaperEmpty'), add items to it ('reaperCons'), and prune it ('reaperAction').
 The reaper will run every two seconds ('reaperDelay'), but will stop running while
 'reaperNull' is true.

 @main@ then loops infinitely ('Control.Monad.forever'). Each second it calculates the fibonacci number
 for a value between 30 and 34, first trying the cache ('reaperRead' and 'Data.Map.Strict.lookup'),
 then falling back to manually calculating it (@fib@)
 and updating the cache with the result ('reaperAdd')

 @clean@ simply removes items cached for more than 10 seconds.
 This function is where you would perform IO-related cleanup,
 like killing threads or closing connections, if that was the purpose of your reaper.

 @
 module Main where

 import "Data.Time" (UTCTime, getCurrentTime, diffUTCTime)
 import "Control.Reaper"
 import "Control.Concurrent" (threadDelay)
 import "Data.Map.Strict" (Map)
 import qualified "Data.Map.Strict" as Map
 import "Control.Monad" (forever)
 import "System.Random" (getStdRandom, randomR)

 fib :: 'Int' -> 'Int'
 fib 0 = 0
 fib 1 = 1
 fib n = fib (n-1) + fib (n-2)

 type Cache = 'Data.Map.Strict.Map' 'Int' ('Int', 'Data.Time.Clock.UTCTime')

 main :: IO ()
 main = do
   reaper <- 'mkReaper' 'defaultReaperSettings'
     { 'reaperEmpty' = Map.'Data.Map.Strict.empty'
     , 'reaperCons' = \\(k, v, time) workload -> Map.'Data.Map.Strict.insert' k (v, time) workload
     , 'reaperAction' = clean
     , 'reaperDelay' = 1000000 * 2 -- Clean every 2 seconds
     , 'reaperNull' = Map.'Data.Map.Strict.null'
     }
   forever $ do
     fibArg <- 'System.Random.getStdRandom' ('System.Random.randomR' (30,34))
     cache <- 'reaperRead' reaper
     let cachedResult = Map.'Data.Map.Strict.lookup' fibArg cache
     case cachedResult of
       'Just' (fibResult, _createdAt) -> 'putStrLn' $ "Found in cache: `fib " ++ 'show' fibArg ++ "` " ++ 'show' fibResult
       'Nothing' -> do
         let fibResult = fib fibArg
         'putStrLn' $ "Calculating `fib " ++ 'show' fibArg ++ "` " ++ 'show' fibResult
         time <- 'Data.Time.Clock.getCurrentTime'
         ('reaperAdd' reaper) (fibArg, fibResult, time)
     'threadDelay' 1000000 -- 1 second

 -- Remove items > 10 seconds old
 clean :: Cache -> IO (Cache -> Cache)
 clean oldMap = do
   currentTime <- 'Data.Time.Clock.getCurrentTime'
   let pruned = Map.'Data.Map.Strict.filter' (\\(_, createdAt) -> currentTime \`diffUTCTime\` createdAt < 10.0) oldMap
   return (\\newData -> Map.'Data.Map.Strict.union' pruned newData)
 @
-}