diff --git a/.travis.yml b/.travis.yml --- a/.travis.yml +++ b/.travis.yml @@ -1,8 +1,8 @@ # This Travis job script has been generated by a script via # -# runghc make_travis_yml_2.hs 'cabal.project' +# runghc make_travis_yml_2.hs 'timer-wheel.cabal' # -# For more information, see https://github.com/hvr/multi-ghc-travis +# For more information, see https://github.com/haskell-CI/haskell-ci # language: c sudo: false @@ -28,12 +28,18 @@ matrix: include: - - compiler: "ghc-8.0.2" - addons: {apt: {packages: [ghc-ppa-tools,cabal-install-head,ghc-8.0.2], sources: [hvr-ghc]}} + - compiler: "ghc-8.6.3" + # env: TEST=--disable-tests BENCH=--disable-benchmarks + addons: {apt: {packages: [ghc-ppa-tools,cabal-install-2.4,ghc-8.6.3], sources: [hvr-ghc]}} + - compiler: "ghc-8.4.4" + # env: TEST=--disable-tests BENCH=--disable-benchmarks + addons: {apt: {packages: [ghc-ppa-tools,cabal-install-2.4,ghc-8.4.4], sources: [hvr-ghc]}} - compiler: "ghc-8.2.2" - addons: {apt: {packages: [ghc-ppa-tools,cabal-install-head,ghc-8.2.2], sources: [hvr-ghc]}} - - compiler: "ghc-8.4.3" - addons: {apt: {packages: [ghc-ppa-tools,cabal-install-head,ghc-8.4.3], sources: [hvr-ghc]}} + # env: TEST=--disable-tests BENCH=--disable-benchmarks + addons: {apt: {packages: [ghc-ppa-tools,cabal-install-2.4,ghc-8.2.2], sources: [hvr-ghc]}} + - compiler: "ghc-8.0.2" + # env: TEST=--disable-tests BENCH=--disable-benchmarks + addons: {apt: {packages: [ghc-ppa-tools,cabal-install-2.4,ghc-8.0.2], sources: [hvr-ghc]}} before_install: - HC=${CC} @@ -51,20 +57,25 @@ - BENCH=${BENCH---enable-benchmarks} - TEST=${TEST---enable-tests} - HADDOCK=${HADDOCK-true} - - INSTALLED=${INSTALLED-true} + - UNCONSTRAINED=${UNCONSTRAINED-true} + - NOINSTALLEDCONSTRAINTS=${NOINSTALLEDCONSTRAINTS-false} - GHCHEAD=${GHCHEAD-false} - travis_retry cabal update -v - "sed -i.bak 's/^jobs:/-- jobs:/' ${HOME}/.cabal/config" - rm -fv cabal.project cabal.project.local - grep -Ev -- '^\s*--' ${HOME}/.cabal/config | grep -Ev '^\s*$' - "printf 'packages: \".\"\\n' > cabal.project" - - cat cabal.project + - "printf 'write-ghc-environment-files: always\\n' >> cabal.project" + - touch cabal.project.local + - "if ! $NOINSTALLEDCONSTRAINTS; then for pkg in $($HCPKG list --simple-output); do echo $pkg | grep -vw -- timer-wheel | sed 's/^/constraints: /' | sed 's/-[^-]*$/ installed/' >> cabal.project.local; done; fi" + - cat cabal.project || true + - cat cabal.project.local || true - if [ -f "./configure.ac" ]; then (cd "." && autoreconf -i); fi - rm -f cabal.project.freeze - - cabal new-build -w ${HC} ${TEST} ${BENCH} --project-file="cabal.project" --dep -j2 - - cabal new-build -w ${HC} --disable-tests --disable-benchmarks --project-file="cabal.project" --dep -j2 + - cabal new-build -w ${HC} ${TEST} ${BENCH} --project-file="cabal.project" --dep -j2 all + - cabal new-build -w ${HC} --disable-tests --disable-benchmarks --project-file="cabal.project" --dep -j2 all - rm -rf .ghc.environment.* "."/dist - DISTDIR=$(mktemp -d /tmp/dist-test.XXXX) @@ -72,23 +83,31 @@ # any command which exits with a non-zero exit code causes the build to fail. script: # test that source-distributions can be generated - - (cd "." && cabal sdist) - - mv "."/dist/timer-wheel-*.tar.gz ${DISTDIR}/ + - cabal new-sdist all + - mv dist-newstyle/sdist/*.tar.gz ${DISTDIR}/ - cd ${DISTDIR} || false - find . -maxdepth 1 -name '*.tar.gz' -exec tar -xvf '{}' \; - "printf 'packages: timer-wheel-*/*.cabal\\n' > cabal.project" - - cat cabal.project + - "printf 'write-ghc-environment-files: always\\n' >> cabal.project" + - touch cabal.project.local + - "if ! $NOINSTALLEDCONSTRAINTS; then for pkg in $($HCPKG list --simple-output); do echo $pkg | grep -vw -- timer-wheel | sed 's/^/constraints: /' | sed 's/-[^-]*$/ installed/' >> cabal.project.local; done; fi" + - cat cabal.project || true + - cat cabal.project.local || true # this builds all libraries and executables (without tests/benchmarks) - cabal new-build -w ${HC} --disable-tests --disable-benchmarks all - # Build with installed constraints for packages in global-db - - if $INSTALLED; then echo cabal new-build -w ${HC} --disable-tests --disable-benchmarks $(${HCPKG} list --global --simple-output --names-only | sed 's/\([a-zA-Z0-9-]\{1,\}\) */--constraint="\1 installed" /g') all | sh; else echo "Not building with installed constraints"; fi - # build & run tests, build benchmarks - cabal new-build -w ${HC} ${TEST} ${BENCH} all + - if [ "x$TEST" = "x--enable-tests" ]; then cabal new-test -w ${HC} ${TEST} ${BENCH} all; fi # cabal check - (cd timer-wheel-* && cabal check) -# REGENDATA ["cabal.project"] + # haddock + - if $HADDOCK; then cabal new-haddock -w ${HC} ${TEST} ${BENCH} all; else echo "Skipping haddock generation";fi + + # Build without installed constraints for packages in global-db + - if $UNCONSTRAINED; then rm -f cabal.project.local; echo cabal new-build -w ${HC} --disable-tests --disable-benchmarks all; else echo "Not building without installed constraints"; fi + +# REGENDATA ["timer-wheel.cabal"] # EOF diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,23 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/) and this project adheres to the [Haskell Package Versioning Policy](https://pvp.haskell.org/). +## Unreleased +### Added +- `destroy` function, for reaping the background thread +- `recurring_` function + +### Changed +- If the timer wheel reaper thread crashes, it will propagate the exception to +the thread that spawned it +- `new` may now throw `InvalidTimerWheelConfig` +- The cancel action returned by `register` is now memoized, which fixes a bug +involving trying to cancel a `recurring` timer twice. The second call used to +spin forever and peg a CPU +- Use `Config` type for creating a timer wheel +- Change argument order around +- Rename `new` to `create` +- Make recurring timers more accurate + ## [0.1.0] - 2018-07-18 ### Added diff --git a/Setup.hs b/Setup.hs deleted file mode 100644 --- a/Setup.hs +++ /dev/null @@ -1,2 +0,0 @@ -import Distribution.Simple -main = defaultMain diff --git a/src/Data/TimerWheel.hs b/src/Data/TimerWheel.hs --- a/src/Data/TimerWheel.hs +++ b/src/Data/TimerWheel.hs @@ -1,52 +1,43 @@ -{-# language CPP #-} -{-# language LambdaCase #-} -{-# language MagicHash #-} -{-# language NamedFieldPuns #-} -{-# language RecursiveDo #-} -{-# language ScopedTypeVariables #-} -{-# language UnboxedTuples #-} -{-# language ViewPatterns #-} - -{-# options_ghc -funbox-strict-fields #-} +{-# OPTIONS_GHC -funbox-strict-fields #-} module Data.TimerWheel ( -- * Timer wheel TimerWheel - , new + , create + , destroy + , Config(..) , register , register_ , recurring + , recurring_ + , InvalidTimerWheelConfig(..) + , TimerWheelDied(..) ) where -import Entries (Entries) import Supply (Supply) +import Wheel (Wheel) -import qualified Entries as Entries import qualified Supply - -import Control.Concurrent -import Control.Exception -import Control.Monad -import Data.Fixed (E6, Fixed(MkFixed)) -import Data.Foldable -import Data.IORef -import Data.Primitive.MutVar -import Data.Primitive.UnliftedArray -import Data.Word (Word64) -import GHC.Base (IO(IO), mkWeak#) -#if MIN_VERSION_base(4,11,0) -import GHC.Clock (getMonotonicTimeNSec) -#else -import System.Clock (Clock(Monotonic), getTime, toNanoSecs) -#endif -import GHC.Prim (RealWorld) -import GHC.Weak (Weak(Weak), deRefWeak) +import qualified Wheel -import qualified GHC.Event as GHC +import Control.Concurrent (ThreadId, forkIOWithUnmask, killThread, myThreadId, + throwTo) +import Control.Exception (AsyncException(ThreadKilled), + Exception(fromException, toException), SomeException, + asyncExceptionFromException, + asyncExceptionToException, catch, throwIO) +import Control.Monad (join, void, when) +import Data.Coerce (coerce) +import Data.Fixed (E6, Fixed(MkFixed)) +import Data.IORef (IORef, newIORef, readIORef, writeIORef) +import Data.Word (Word64) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) -- | A 'TimerWheel' is a vector-of-collections-of timers to fire. It is --- configured with a /spoke count/ and /resolution/. A timeout thread is spawned --- to step through the timer wheel and fire expired timers at regular intervals. +-- configured with a /spoke count/ and /resolution/. Timers may be scheduled +-- arbitrarily far in the future. A timeout thread is spawned to step through +-- the timer wheel and fire expired timers at regular intervals. -- -- * The /spoke count/ determines the size of the timer vector. -- @@ -69,7 +60,7 @@ -- each spoke, __more accurate__ timers, and will require __more wakeups__ -- by the timeout thread. -- --- * The timeout thread has three important properties: +-- * The timeout thread has some important properties: -- -- * There is only one, and it fires expired timers synchronously. If your -- timer actions execute quicky, 'register' them directly. Otherwise, @@ -77,13 +68,10 @@ -- performed on a job queue. -- -- * Synchronous exceptions thrown by enqueued @IO@ actions will bring the --- thread down, and no more timeouts will ever fire. If you want to catch --- exceptions and log them, for example, you will have to bake this into --- the registered actions yourself. --- --- * The life of the timeout thread is scoped to the life of the timer --- wheel. When the timer wheel is garbage collected, the timeout thread --- will automatically stop doing work, and die gracefully. +-- thread down, which will cause it to asynchronously throw a +-- 'TimerWheelDied' exception to the thread that 'create'd it. If you want +-- to catch exceptions and log them, for example, you will have to bake +-- this into the registered actions yourself. -- -- Below is a depiction of a timer wheel with @6@ timers inserted across @8@ -- spokes, and a resolution of @0.1s@. @@ -95,201 +83,179 @@ -- +-----+-----+-----+-----+-----+-----+-----+-----+ -- @ data TimerWheel = TimerWheel - { wheelResolution :: !Word64 - -- ^ The length of time that each entry corresponds to, in nanoseconds. - , wheelSupply :: !Supply + { wheelSupply :: !Supply -- ^ A supply of unique ints. - , wheelEntries :: !(UnliftedArray (MutVar RealWorld Entries)) + , wheelWheel :: !Wheel -- ^ The array of collections of timers. + , wheelThread :: !ThreadId } --- | @new n s@ creates a 'TimerWheel' with __@n@__ spokes and a resolution of --- __@s@__ seconds. -new :: Int -> Fixed E6 -> IO TimerWheel -new slots (MkFixed (fromInteger -> resolution)) = do - wheel :: UnliftedArray (MutVar RealWorld Entries) <- do - wheel :: MutableUnliftedArray RealWorld (MutVar RealWorld Entries) <- - unsafeNewUnliftedArray slots - for_ [0..slots-1] $ \i -> - writeUnliftedArray wheel i =<< newMutVar Entries.empty - freezeUnliftedArray wheel 0 slots - - supply :: Supply <- - Supply.new +data Config + = Config + { spokes :: !Natural -- ^ Spoke count. + , resolution :: !(Fixed E6) -- ^ Resolution, in seconds. + } deriving (Generic, Show) - weakWheel :: Weak (UnliftedArray (MutVar RealWorld Entries)) <- - case wheel of - UnliftedArray wheel# -> - IO $ \s -> - case mkWeak# wheel# wheel (\t -> (# t, () #)) s of - (# s', w #) -> - (# s', Weak w #) +-- | The timeout thread died. +newtype TimerWheelDied + = TimerWheelDied SomeException + deriving (Show) - (void . forkIO) - (reaper resolution (sizeofUnliftedArray wheel) weakWheel) +instance Exception TimerWheelDied where + toException = asyncExceptionToException + fromException = asyncExceptionFromException - pure TimerWheel - { wheelResolution = fromIntegral resolution * 1000 - , wheelSupply = supply - , wheelEntries = wheel - } +-- | The timer wheel config was invalid. +-- +-- * @spokes@ must be positive, and less than @maxBound@ of @Int@. +-- * @resolution@ must be positive. +data InvalidTimerWheelConfig + = InvalidTimerWheelConfig !Config + deriving (Show) -reaper - :: Int - -> Int - -> Weak (UnliftedArray (MutVar RealWorld Entries)) - -> IO () -reaper resolution len weakWheel = do - manager :: GHC.TimerManager <- - GHC.getSystemTimerManager - loop manager 0 - where - -- Reaper loop: we haven't yet run the entries at index 'i'. - loop :: GHC.TimerManager -> Int -> IO () - loop manager i = do - -- Sleep until the roughly the next bucket. - waitVar :: MVar () <- - newEmptyMVar +instance Exception InvalidTimerWheelConfig - mask_ $ do - key :: GHC.TimeoutKey <- - GHC.registerTimeout manager resolution (putMVar waitVar ()) - takeMVar waitVar `onException` GHC.unregisterTimeout manager key +-- | Create a timer wheel. +-- +-- /Throws./ If the config is invalid, throws 'InvalidTimerWheelConfig'. +-- +-- /Throws./ If the timeout thread dies, asynchronously throws 'TimerWheelDied' +-- to the thread that called 'create'. +create :: Config -> IO TimerWheel +create config@(Config { spokes, resolution }) = do + when (invalidConfig config) + (throwIO (InvalidTimerWheelConfig config)) - now :: Word64 <- - getMonotonicTime + wheel :: Wheel <- + Wheel.create + (fromIntegral spokes) + (fromIntegral (coerce resolution :: Integer)) - -- De-ref the wheel, and if it's gone, we go too. - deRefWeak weakWheel >>= \case - Nothing -> - pure () - Just wheel -> do - -- Figure out which bucket we're in. Usually this will be 'i+1', but - -- maybe we were scheduled a bit early and ended up in 'i', or maybe - -- running the entries in bucket 'i-1' took a long time and we slept all - -- the way to 'i+2'. In any case, we should run the entries in buckets - -- up-to-but-not-including this one, beginning with bucket 'i'. + supply :: Supply <- + Supply.new - let - j :: Int - j = - fromIntegral (now `div` (fromIntegral resolution * 1000)) - `mod` len + thread <- myThreadId - let - is :: [Int] - is = - if j >= i - then - [i .. j-1] - else - [i .. len - 1] ++ [0 .. j-1] + reaperThread <- + forkIOWithUnmask $ \unmask -> + unmask (Wheel.reap wheel) + `catch` \e -> + case fromException e of + Just ThreadKilled -> pure () + _ -> throwTo thread (TimerWheelDied e) - -- To actually run the entries in a bucket, partition them into expired - -- (count == 0) and alive (count > 0). Run the expired entries and - -- decrement the alive entries' counts by 1. + pure TimerWheel + { wheelSupply = supply + , wheelWheel = wheel + , wheelThread = reaperThread + } - for_ is $ \k -> do - let - entriesRef :: MutVar RealWorld Entries - entriesRef = - indexUnliftedArray wheel k +-- | Tear down a timer wheel by killing the timeout thread. +destroy :: TimerWheel -> IO () +destroy wheel = + killThread (wheelThread wheel) - join - (atomicModifyMutVar' entriesRef - (\entries -> - if Entries.null entries - then - (entries, pure ()) - else - case Entries.squam entries of - (expired, alive) -> - (alive, sequence_ expired))) - loop manager j +invalidConfig :: Config -> Bool +invalidConfig Config { spokes, resolution } = + or + [ spokes == 0 + , spokes > fromIntegral (maxBound :: Int) + , resolution <= 0 + ] --- | @register n m w@ registers an action __@m@__ in timer wheel __@w@__ to fire --- after __@n@__ seconds. +-- | @register wheel action delay@ registers an action __@action@__ in timer +-- wheel __@wheel@__ to fire after __@delay@__ seconds. -- -- Returns an action that, when called, attempts to cancel the timer, and -- returns whether or not it was successful (@False@ means the timer has already -- fired). -register :: Fixed E6 -> IO () -> TimerWheel -> IO (IO Bool) -register (MkFixed ((*1000) . fromIntegral -> delay)) action wheel = do - newEntryId :: Int <- - Supply.next (wheelSupply wheel) - - entriesVar :: MutVar RealWorld Entries <- - entriesIn delay wheel +-- +-- Subsequent calls to the cancel action have no effect, and continue to return +-- whatever the first result was. +register :: + TimerWheel -- ^ + -> IO () -- ^ Action + -> Fixed E6 -- ^ Delay, in seconds + -> IO (IO Bool) +register wheel action (MkFixed (fromIntegral -> delay)) = + _register wheel action delay - atomicModifyMutVar' entriesVar - (\entries -> - (Entries.insert newEntryId entryCount action entries, ())) +-- | Like 'register', but for when you don't intend to cancel the timer. +register_ :: + TimerWheel -- ^ + -> IO () -- ^ Action + -> Fixed E6 -- ^ Delay, in seconds + -> IO () +register_ wheel action delay = + void (register wheel action delay) - pure $ do - atomicModifyMutVar' entriesVar - (\entries -> - case Entries.delete newEntryId entries of - Nothing -> - (entries, False) - Just entries' -> - (entries', True)) - where - entryCount :: Word64 - entryCount = - delay `div` - (fromIntegral (sizeofUnliftedArray (wheelEntries wheel)) - * wheelResolution wheel) +_register :: TimerWheel -> IO () -> Word64 -> IO (IO Bool) +_register wheel action delay = do + key <- Supply.next (wheelSupply wheel) + Wheel.insert (wheelWheel wheel) key action delay --- | Like 'register', but for when you don't intend to cancel the timer. -register_ :: Fixed E6 -> IO () -> TimerWheel -> IO () -register_ delay action wheel = - void (register delay action wheel) +_register_ :: TimerWheel -> IO () -> Word64 -> IO () +_register_ wheel action delay = + void (_register wheel action delay) --- | @recurring n m w@ registers an action __@m@__ in timer wheel __@w@__ to --- fire every __@n@__ seconds. +-- | @recurring wheel action delay@ registers an action __@action@__ in timer +-- wheel __@wheel@__ to fire every __@delay@__ seconds. -- -- Returns an action that, when called, cancels the recurring timer. -recurring :: Fixed E6 -> IO () -> TimerWheel -> IO (IO ()) -recurring delay action wheel = mdo +recurring :: + TimerWheel + -> IO () -- ^ Action + -> Fixed E6 -- ^ Delay, in seconds + -> IO (IO ()) +recurring wheel action (MkFixed (fromIntegral -> delay)) = mdo + let + doAction :: IO () + doAction = do + -- Re-register one bucket early, to account for the fact that timers are + -- expired at the *end* of a bucket. + -- + -- +---+---+---+---+ + -- { A | | | } + -- +---+---+---+---+ + -- | + -- The reaper thread fires 'A' approximately here, so if it's meant + -- to be repeated every two buckets, and we just re-register it at + -- this time, three buckets will pass before it's run again. So, we + -- act as if it's still "one bucket ago" at the moment we re-register + -- it. + writeIORef cancelRef =<< + _register + wheel + doAction + (delay - Wheel.resolution (wheelWheel wheel)) + action + cancel :: IO Bool <- - register delay (action' cancelRef) wheel + _register wheel doAction delay + cancelRef :: IORef (IO Bool) <- newIORef cancel + pure (untilTrue (join (readIORef cancelRef))) - where - action' :: IORef (IO Bool) -> IO () - action' cancelRef = do - action - cancel :: IO Bool <- - register delay (action' cancelRef) wheel - writeIORef cancelRef cancel --- | @entriesIn delay wheel@ returns the bucket in @wheel@ that corresponds to --- @delay@ nanoseconds from now. -entriesIn :: Word64 -> TimerWheel -> IO (MutVar RealWorld Entries) -entriesIn delay TimerWheel{wheelResolution, wheelEntries} = do - now :: Word64 <- - getMonotonicTime - pure (index ((now+delay) `div` wheelResolution)) - where - index :: Word64 -> MutVar RealWorld Entries - index i = - indexUnliftedArray wheelEntries - (fromIntegral i `rem` sizeofUnliftedArray wheelEntries) +-- | Like 'recurring', but for when you don't intend to cancel the timer. +recurring_ :: + TimerWheel + -> IO () -- ^ Action + -> Fixed E6 -- ^ Delay, in seconds + -> IO () +recurring_ wheel action (MkFixed (fromIntegral -> delay)) = + _register_ wheel doAction delay --- | Repeat an IO action until it returns 'True'. + where + doAction :: IO () + doAction = do + _register_ wheel doAction (delay - Wheel.resolution (wheelWheel wheel)) + action + +-- Repeat an IO action until it returns 'True'. untilTrue :: IO Bool -> IO () untilTrue action = action >>= \case - True -> - pure () - False -> - untilTrue action - -getMonotonicTime :: IO Word64 -getMonotonicTime = -#if MIN_VERSION_base(4,11,0) - getMonotonicTimeNSec -#else - fromIntegral . toNanoSecs <$> getTime Monotonic -#endif + True -> pure () + False -> untilTrue action diff --git a/src/Entries.hs b/src/Entries.hs --- a/src/Entries.hs +++ b/src/Entries.hs @@ -8,7 +8,7 @@ , size , insert , delete - , squam + , partition ) where import Data.Coerce @@ -55,8 +55,8 @@ (\(_, _, ys) -> ys) <$> IntPSQ.deleteView i xs -- | Extract expired timers. -squam :: Entries -> ([IO ()], Entries) -squam (Entries entries) = +partition :: Entries -> ([IO ()], Entries) +partition (Entries entries) = case IntPSQ.atMostView 0 entries of (expired, alive) -> (map f expired, Entries (IntPSQ.unsafeMapMonotonic g alive)) @@ -68,4 +68,4 @@ g :: Int -> Word64 -> IO () -> (Word64, IO ()) g _ n m = (n-1, m) -{-# INLINABLE squam #-} +{-# INLINABLE partition #-} diff --git a/src/Wheel.hs b/src/Wheel.hs new file mode 100644 --- /dev/null +++ b/src/Wheel.hs @@ -0,0 +1,179 @@ +{-# LANGUAGE CPP #-} + +module Wheel + ( Wheel(resolution) + , create + , lenMicros + , insert + , reap + ) where + +import Entries (Entries) + +import qualified Entries as Entries + +import Control.Concurrent (threadDelay) +import Control.Concurrent.MVar +import Control.Monad (join, when) +import Data.Foldable (for_) +import Data.Primitive.MutVar (MutVar, atomicModifyMutVar', newMutVar) +import Data.Primitive.UnliftedArray (MutableUnliftedArray, UnliftedArray, + freezeUnliftedArray, indexUnliftedArray, + sizeofUnliftedArray, + unsafeNewUnliftedArray, writeUnliftedArray) +import Data.Word (Word64) +import GHC.Prim (RealWorld) +import System.IO.Unsafe (unsafeInterleaveIO) + +#if MIN_VERSION_base(4,11,0) +import GHC.Clock (getMonotonicTimeNSec) +#else +import System.Clock (Clock(Monotonic), getTime, toNanoSecs) +#endif + + + +type IORef + = MutVar RealWorld + +data Wheel + = Wheel + { buckets :: !(UnliftedArray (IORef Entries)) + , resolution :: !Word64 -- micros + } + +create :: + Int + -> Word64 + -> IO Wheel +create spokes resolution = do + mbuckets :: MutableUnliftedArray RealWorld (IORef Entries) <- + unsafeNewUnliftedArray spokes + + for_ [0 .. spokes-1] $ \i -> + writeUnliftedArray mbuckets i =<< newMutVar Entries.empty + + buckets :: UnliftedArray (IORef Entries) <- + freezeUnliftedArray mbuckets 0 spokes + + pure Wheel + { buckets = buckets + , resolution = resolution + } + +numSpokes :: Wheel -> Int +numSpokes wheel = + sizeofUnliftedArray (buckets wheel) + +lenMicros :: Wheel -> Word64 +lenMicros wheel = + fromIntegral (numSpokes wheel) * resolution wheel + +bucket :: Wheel -> Word64 -> IORef Entries +bucket wheel time = + indexUnliftedArray + (buckets wheel) + (index wheel time) + +index :: Wheel -> Word64 -> Int +index wheel@Wheel{resolution} time = + fromIntegral (time `div` resolution) `rem` numSpokes wheel + +insert :: + Wheel + -> Int + -> IO () + -> Word64 + -> IO (IO Bool) +insert wheel key action delay = do + now :: Word64 <- + getMonotonicMicros + + let + time :: Word64 + time = + now + delay + + let + count :: Word64 + count = + delay `div` lenMicros wheel + + let + bucketVar :: MutVar RealWorld Entries + bucketVar = + bucket wheel time + + atomicModifyMutVar' bucketVar + (\entries -> + (Entries.insert key count action entries, ())) + + canceledVar :: MVar (Maybe Bool) <- + unsafeInterleaveIO (newMVar Nothing) + + pure $ do + modifyMVar canceledVar $ \result -> do + canceled <- + maybe + (atomicModifyMutVar' bucketVar + (\entries -> + maybe + (entries, False) + (, True) + (Entries.delete key entries))) + pure + result + pure (Just canceled, canceled) + + +reap :: Wheel -> IO () +reap wheel@Wheel{buckets, resolution} = do + now :: Word64 <- + getMonotonicMicros + + let + -- How far in time are we into the very first bucket? Sleep until it's over. + elapsedBucketMicros :: Word64 + elapsedBucketMicros = + now `rem` resolution + + let + remainingBucketMicros :: Word64 + remainingBucketMicros = + resolution - elapsedBucketMicros + + threadDelay (fromIntegral remainingBucketMicros) + + loop + (now + remainingBucketMicros + resolution) + (index wheel now) + + where + loop :: Word64 -> Int -> IO () + loop nextTime i = do + join + (atomicModifyMutVar' (indexUnliftedArray buckets i) + (\entries -> + if Entries.null entries + then + (entries, pure ()) + else + case Entries.partition entries of + (expired, alive) -> + (alive, sequence_ expired))) + + afterTime :: Word64 <- + getMonotonicMicros + + when (afterTime < nextTime) $ + (threadDelay (fromIntegral (nextTime - afterTime))) + + loop (nextTime + resolution) ((i+1) `rem` numSpokes wheel) + +getMonotonicMicros :: IO Word64 +getMonotonicMicros = +#if MIN_VERSION_base(4,11,0) + (`div` 1000) <$> getMonotonicTimeNSec +#else + ((`div` 1000) . fromIntegral . toNanoSecs <$> getTime Monotonic) +#endif diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,90 @@ +{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-} + +import Control.Concurrent +import Control.Exception +import Control.Monad +import Data.Fixed (E6, Fixed(..)) +import Data.IORef +import Data.TimerWheel +import System.Random + +main :: IO () +main = do + do + putStrLn "Timer wheel runs scheduled actions" + let n = 10000 + ref <- newIORef (0::Int) + wheel <- create Config { spokes = 4, resolution = 0.25 } + putStrLn ("> Inserting " ++ show n ++ " timers") + replicateM_ n $ do + delay <- randomRIO (0, 4*1000*1000) + register_ wheel (modifyIORef' ref (+1)) (MkFixed delay) + putStrLn "> Sleeping for 5s" + sleep 5 + readIORef ref `is` n + destroy wheel + + do + putStrLn "Timers can be canceled" + let n = 10000 + ref <- newIORef (0::Int) + wheel <- create Config { spokes = 4, resolution = 0.25 } + putStrLn ("> Inserting " ++ show n ++ " timers") + cancels <- + replicateM n $ do + delay <- randomRIO (0, 4*1000*1000) + register wheel (modifyIORef' ref (+1)) (MkFixed delay) + putStrLn "> Sleeping for 2s" + sleep 2 + putStrLn "> Canceling all timers" + successes <- sequence cancels + readIORef ref `is` (n - length (filter id successes)) + destroy wheel + + do + putStrLn "Re-calling a successful cancel works" + wheel <- create Config { spokes = 4, resolution = 0.25 } + cancel <- register wheel (pure ()) 1 + cancel `is` True + cancel `is` True + destroy wheel + + do + putStrLn "Re-calling a failed cancel works" + wheel <- create Config { spokes = 4, resolution = 0.25 } + cancel <- register wheel (pure ()) 0.5 + sleep 1 + cancel `is` False + cancel `is` False + destroy wheel + + do + putStrLn "Recurring timers work" + ref <- newIORef (0::Int) + wheel <- create Config { spokes = 4, resolution = 0.05 } + recurring_ wheel (modifyIORef' ref (+1)) 0.15 + sleep 1 + readIORef ref `is` (6::Int) + destroy wheel + + do + putStrLn "Recurring timers can be canceled" + ref <- newIORef (0::Int) + wheel <- create Config { spokes = 4, resolution = 0.05 } + cancel <- recurring wheel (modifyIORef' ref (+1)) 0.15 + sleep 1 + cancel + sleep 1 + readIORef ref `is` (6::Int) + destroy wheel + +class Assert a b where + is :: a -> b -> IO () + +instance (Eq a, Show a) => Assert (IO a) a where + is mx y = do + x <- mx + unless (x == y) (throwIO (userError (show x ++ " /= " ++ show y))) + +sleep :: Fixed E6 -> IO () +sleep (MkFixed micro) = threadDelay (fromIntegral micro) diff --git a/timer-wheel.cabal b/timer-wheel.cabal --- a/timer-wheel.cabal +++ b/timer-wheel.cabal @@ -1,7 +1,7 @@ cabal-version: 2.0 name: timer-wheel -version: 0.1.0 +version: 0.2.0 category: Data description: This library provides a timer wheel data structure for @@ -27,7 +27,11 @@ license: BSD3 license-file: LICENSE build-type: Simple -tested-with: GHC == 8.0.2, GHC == 8.2.2, GHC == 8.4.3 +tested-with: + GHC == 8.0.2 + , GHC == 8.2.2 + , GHC == 8.4.4 + , GHC == 8.6.3 extra-source-files: .travis.yml @@ -39,18 +43,32 @@ library build-depends: - atomic-primops ^>= 0.8 - , base ^>= 4.9 || ^>= 4.10 || ^>= 4.11 - , ghc-prim ^>= 0.5 - , primitive ^>= 0.6 - , psqueues ^>= 0.2.7 + atomic-primops ^>= 0.8, + base ^>= 4.9 || ^>= 4.10 || ^>= 4.11 || ^>= 4.12, + ghc-prim ^>= 0.5, + primitive ^>= 0.6, + psqueues ^>= 0.2.7 if !impl(ghc >= 8.4) build-depends: clock ^>= 0.7 + + default-extensions: + DeriveGeneric + GeneralizedNewtypeDeriving + LambdaCase + NamedFieldPuns + RecursiveDo + ScopedTypeVariables + TupleSections + TypeApplications + ViewPatterns + default-language: Haskell2010 + exposed-modules: Data.TimerWheel + ghc-options: -fprint-expanded-synonyms -fprint-explicit-foralls @@ -69,8 +87,34 @@ ghc-options: -Wpartial-fields -Wmissing-export-lists + hs-source-dirs: src + other-modules: Entries Supply + Wheel + +test-suite tests + build-depends: + base, + random, + timer-wheel + + default-language: + Haskell2010 + + ghc-options: + -threaded + -with-rtsopts=-N2 + -Wall + + hs-source-dirs: + test + + main-is: + Main.hs + + type: + exitcode-stdio-1.0