{-# 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
{ 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
, reaperDelay :: {-# UNPACK #-} !Int
-- ^ Number of microseconds to delay between calls of 'reaperAction'.
--
-- Default: 30 seconds.
--
-- @since 0.1.0
, reaperCons :: item -> workload -> workload
-- ^ Add an item onto a workload.
--
-- Default: list consing.
--
-- @since 0.1.0
, reaperNull :: workload -> Bool
-- ^ Check if a workload is empty, in which case the worker thread
-- will shut down.
--
-- Default: 'null'.
--
-- @since 0.1.0
, reaperEmpty :: workload
-- ^ An empty workload.
--
-- Default: empty list.
--
-- @since 0.1.0
, 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 =
ReaperSettings
{ reaperAction = \wl -> pure (wl ++)
, reaperDelay = 30000000
, reaperCons = (:)
, reaperNull = null
, reaperEmpty = []
, reaperThreadName = "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 settings@ReaperSettings {..} = do
stateRef <- newIORef NoReaper
tidRef <- newIORef Nothing
return
Reaper
{ reaperAdd = add settings stateRef tidRef
, reaperRead = readRef stateRef
, reaperModify = modifyRef stateRef
, reaperStop = stop stateRef
, reaperKill = kill tidRef
}
where
readRef stateRef = do
mx <- readIORef stateRef
case mx of
NoReaper -> return reaperEmpty
Workload wl -> return wl
modifyRef stateRef modifier = atomicModifyIORef' stateRef $ \case
NoReaper ->
(NoReaper, reaperEmpty)
Workload wl ->
let !wl' = modifier wl
in (Workload wl', wl')
stop stateRef = atomicModifyIORef' stateRef $ \case
NoReaper -> (NoReaper, reaperEmpty)
Workload x -> (Workload reaperEmpty, x)
kill tidRef = do
mtid <- readIORef tidRef
forM_ mtid killThread
add ::
(MonadUnliftIO m) =>
ReaperSettings m workload item ->
IORef (State workload) ->
IORef (Maybe ThreadId) ->
item ->
m ()
add settings@ReaperSettings {..} stateRef tidRef item =
mask_ $ join $ atomicModifyIORef' stateRef cons
where
cons NoReaper =
let wl = reaperCons item reaperEmpty
in (Workload wl, spawn settings stateRef tidRef)
cons (Workload wl) =
let wl' = reaperCons item wl
in (Workload wl', return ())
spawn ::
(MonadUnliftIO m) =>
ReaperSettings m workload item ->
IORef (State workload) ->
IORef (Maybe ThreadId) ->
m ()
spawn settings stateRef tidRef = do
tid <- forkIO $ reaper settings stateRef tidRef
liftIO . labelThread tid $ reaperThreadName settings
writeIORef tidRef $ Just tid
reaper ::
(MonadUnliftIO m) =>
ReaperSettings m workload item ->
IORef (State workload) ->
IORef (Maybe ThreadId) ->
m ()
reaper settings@ReaperSettings {..} stateRef tidRef = do
threadDelay reaperDelay
-- Getting the current jobs. Push an empty job to the reference.
wl <- atomicModifyIORef' stateRef swapWithEmpty
-- Do the jobs. A function to merge the left jobs and
-- new jobs is returned.
!merge <- reaperAction wl
-- Merging the left jobs and new jobs.
-- If there is no jobs, this thread finishes.
cont <- atomicModifyIORef' stateRef (check merge)
if cont
then reaper settings stateRef tidRef
else writeIORef tidRef Nothing
where
swapWithEmpty NoReaper = error "Control.Reaper.reaper: unexpected NoReaper (1)"
swapWithEmpty (Workload wl) = (Workload reaperEmpty, wl)
check _ NoReaper = error "Control.Reaper.reaper: unexpected NoReaper (2)"
check merge (Workload wl)
-- If there is no job, reaper is terminated.
| reaperNull wl' = (NoReaper, False)
-- If there are jobs, carry them out.
| otherwise = (Workload wl', True)
where
wl' = merge 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 f =
go id
where
go !front [] = pure front
go !front (x : xs) = do
my <- f x
let front' =
case my of
Nothing -> front
Just y -> front . (y :)
go front' 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)
@
-}