diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Moritz Schulte (c) 2017 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Moritz Schulte nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,22 @@ +# throttle-io-stream [![Hackage version](https://img.shields.io/hackage/v/throttle-io-stream.svg?label=Hackage)](https://hackage.haskell.org/package/throttle-io-stream) [![Stackage version](https://www.stackage.org/package/throttle-io-stream/badge/lts?label=Stackage)](https://www.stackage.org/package/throttle-io-stream) [![Build Status](https://travis-ci.org/mtesseract/throttle-io-stream.svg?branch=master)](https://travis-ci.org/mtesseract/throttle-io-stream) + +### About + +This packages provides throttling functionality for arbitrary IO +producers and consumers. The core function exported is the following: + +```haskell +throttle :: ThrottleConf a -- ^ Throttling configuration + -> IO (Maybe a) -- ^ Input callback + -> (Maybe a -> IO ()) -- ^ Output callback + -> IO (Async ()) -- ^ Returns an async handler for this + -- throttling process +``` + +This will spawn asynchronous operations that + +1. consume data using the provided input callback and write it into an + internal buffering queue and + +1. produce data from the buffering queue using the provided consumer + callback. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/src/Control/Concurrent/Throttle.hs b/src/Control/Concurrent/Throttle.hs new file mode 100644 --- /dev/null +++ b/src/Control/Concurrent/Throttle.hs @@ -0,0 +1,272 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} + +-- | Throttle + +module Control.Concurrent.Throttle + ( Measure + , ThrottleConf + , newThrottleConf + , throttleConfSetMeasure + , throttleConfThrottleConsumer + , throttleConfThrottleProducer + , throttleConfSetInterval + , throttleConfSetMaxThroughput + , throttleConfSetBufferSize + , throttleConfSetEmaAlpha + , throttle + ) where + +import Control.Concurrent +import Control.Concurrent.Async +import Control.Concurrent.STM +import Control.Concurrent.STM.TBMQueue +import Control.Concurrent.Throttle.Ema +import Control.Monad +import Control.Monad.IO.Class +import System.Clock + +-- | Type of a measure function for items of the specified type. The +-- measure function is used by the user to specify a notion of size +-- used for throughput computations. +type Measure a = a -> Double + +-- | Defines the throttling mode. Consuming and producing can be +-- throttled independently. +data ThrottleMode = ThrottleMode + { throttleConsumer :: Bool + , throttleProducer :: Bool + } + +-- | Configuration for throttling. +data ThrottleConf a = ThrottleConf + { throttleConfMeasure :: Measure a -- ^ Measure function for values of type a + , throttleConfMode :: ThrottleMode -- ^ Throtting mode + , throttleConfInterval :: Double -- ^ Interval in milliseconds + , throttleConfMaxThroughput :: Double -- ^ Maximum throughput allowed + , throttleConfBufferSize :: Int -- ^ Size for buffer queue + , throttleConfEmaAlpha :: Double -- ^ Exponential weight for computing current item size + } + +-- | Type used for internal statistics collection. It is used for +-- approximating the "current" size of items processed. +data Stats = Stats + { statsEmaItemSizeIn :: Ema -- ^ Exponentially weighted moving + -- average for item size of read + -- items + , statsEmaItemSizeOut :: Ema -- ^ Exponentially weighted moving + -- average for item size of written + -- items + } deriving (Show) + +-- | Produce a new 'ThrottleConf'. +newThrottleConf :: ThrottleConf a +newThrottleConf = defaultThrottleConf + +-- | Default 'ThrottleConf'. +defaultThrottleConf :: ThrottleConf a +defaultThrottleConf = ThrottleConf + { throttleConfMeasure = const 1 + , throttleConfMode = ThrottleMode { throttleConsumer = False + , throttleProducer = False } + , throttleConfInterval = 1000 + , throttleConfMaxThroughput = 100 + , throttleConfBufferSize = 1024 + , throttleConfEmaAlpha = defaultEmaAlpha } + +-- | Set measure function in configuration. +throttleConfSetMeasure :: Measure a -> ThrottleConf a -> ThrottleConf a +throttleConfSetMeasure measure conf = conf { throttleConfMeasure = measure } + +throttleConfThrottleProducer :: ThrottleConf a -> ThrottleConf a +throttleConfThrottleProducer conf @ ThrottleConf { .. } = + conf { throttleConfMode = throttleConfMode { throttleProducer = True } } + +throttleConfThrottleConsumer :: ThrottleConf a -> ThrottleConf a +throttleConfThrottleConsumer conf @ ThrottleConf { .. } = + conf { throttleConfMode = throttleConfMode { throttleConsumer = True } } + +-- | Set interval in configuration. +throttleConfSetInterval :: Double -> ThrottleConf a -> ThrottleConf a +throttleConfSetInterval interval conf = conf { throttleConfInterval = interval } + +-- | Set max throughput in configuration. +throttleConfSetMaxThroughput :: Double -> ThrottleConf a -> ThrottleConf a +throttleConfSetMaxThroughput throughput conf = + conf { throttleConfMaxThroughput = throughput } + +-- | Set buffer size in configuration. +throttleConfSetBufferSize :: Int -> ThrottleConf a -> ThrottleConf a +throttleConfSetBufferSize n conf = conf { throttleConfBufferSize = n } + +-- | Set exponential weight factor used for computing current item +-- size. +throttleConfSetEmaAlpha :: Double -> ThrottleConf a -> ThrottleConf a +throttleConfSetEmaAlpha alpha conf = conf { throttleConfEmaAlpha = alpha } + +-- | Default exponential weight factor for computing current item +-- size. +defaultEmaAlpha :: Double +defaultEmaAlpha = 0.5 + +-- | Asynchonously read items with the given input callback and write +-- them throttled with the given output callback. +throttle :: ThrottleConf a -- ^ Throttling configuration + -> IO (Maybe a) -- ^ Input callback + -> (Maybe a -> IO ()) -- ^ Output callback + -> IO (Async ()) -- ^ Returns an async handler for this + -- throttling process +throttle (conf @ ThrottleConf { .. }) readItem writeItem = do + queueBuffer <- atomically $ newTBMQueue throttleConfBufferSize + statsTVar <- atomically $ newTVar Stats { statsEmaItemSizeIn = newEma throttleConfEmaAlpha 0 + , statsEmaItemSizeOut = newEma throttleConfEmaAlpha 0 } + async $ + withAsync (consumer conf statsTVar queueBuffer readItem ) $ \ consumerThread -> + withAsync (producer conf statsTVar queueBuffer writeItem) $ \ producerThread -> do + link consumerThread + link producerThread + void $ waitBoth consumerThread producerThread + +-- | Unthrottled consumer. Fills the provided buffer with new items as +-- fast as possible. +consumerUnthrottled :: TBMQueue a -> IO (Maybe a) -> IO () +consumerUnthrottled buffer readItem = go + where go = + readItem >>= \ case + Just a -> atomically (writeTBMQueue buffer a) >> go + Nothing -> atomically (closeTBMQueue buffer) + +-- | Throttled Consumer. +consumerThrottled :: ThrottleConf a + -> TVar Stats + -> TBMQueue a + -> IO (Maybe a) + -> IO () +consumerThrottled (conf @ ThrottleConf { .. }) statsTVar buffer readItem = go + where go = do + (maybeStats, consumeDuration) <- timeAction $ + readItem >>= \ case + Just a -> atomically $ do + writeTBMQueue buffer a + modifyTVar statsTVar (updateStatsIn conf a) + Just <$> readTVar statsTVar + Nothing -> atomically $ do + closeTBMQueue buffer + return Nothing + case maybeStats of + Just stats -> do + liftIO . threadDelay $ throttleDelayIn stats conf consumeDuration + go + Nothing -> return () + +-- | Producer thread, dispatches to throttled or non-throttled case. +producer :: ThrottleConf a + -> TVar Stats + -> TBMQueue a + -> (Maybe a -> IO ()) + -> IO () +producer (conf @ ThrottleConf { .. }) stats = + if throttleProducer throttleConfMode + then producerThrottled conf stats + else producerUnthrottled + +-- | Consumer, dispatches to throttled or non-throttled case. +consumer :: ThrottleConf a + -> TVar Stats + -> TBMQueue a + -> IO (Maybe a) + -> IO () +consumer (conf @ ThrottleConf { .. }) stats = + if throttleConsumer throttleConfMode + then consumerThrottled conf stats + else consumerUnthrottled + +-- | Throttled Producer. Reads items from the provided buffer and +-- writes them throttled. When the queue is empty, this function +-- returns. +producerThrottled :: ThrottleConf a + -> TVar Stats + -> TBMQueue a + -> (Maybe a -> IO ()) + -> IO () +producerThrottled (conf @ ThrottleConf { .. }) statsTVar buffer writeItem = go + where go = do + (maybeStats, produceDuration) <- timeAction $ + atomically (readTBMQueue buffer) >>= \ case + Just a -> do + writeItem (Just a) + atomically $ do + modifyTVar statsTVar (updateStatsOut conf a) + Just <$> readTVar statsTVar + Nothing -> do + writeItem Nothing + return Nothing + case maybeStats of + Just stats -> do + liftIO . threadDelay $ throttleDelayOut stats conf produceDuration + go + Nothing -> return () + +-- | Unthrottled producer. +producerUnthrottled :: TBMQueue a + -> (Maybe a -> IO ()) + -> IO () +producerUnthrottled buffer writeItem = go + where go = + atomically (readTBMQueue buffer) >>= \ case + Just a -> writeItem (Just a) >> go + Nothing -> writeItem Nothing + +-- | Update provided statistics. +updateStatsOut :: ThrottleConf a -> a -> Stats -> Stats +updateStatsOut ThrottleConf { .. } a (stats @ Stats { .. }) = + stats { statsEmaItemSizeOut = emaUpdate aSize statsEmaItemSizeOut } + where aSize = throttleConfMeasure a + +-- | Update provided statistics. +updateStatsIn :: ThrottleConf a -> a -> Stats -> Stats +updateStatsIn ThrottleConf { .. } a (stats @ Stats { .. }) = + stats { statsEmaItemSizeIn = emaUpdate aSize statsEmaItemSizeIn } + where aSize = throttleConfMeasure a + +-- | Measure execution of an IO action. Returns a pair consisting of +-- the result of the IO action and the duration in milliseconds. +timeAction :: IO a -> IO (a, Double) +timeAction io = do + t0 <- getTime Monotonic + a <- io + t1 <- getTime Monotonic + return (a, t1 `timeSpecDiff` t0) + + where timeSpecDiff :: TimeSpec -> TimeSpec -> Double + timeSpecDiff ts1 ts0 = + fromIntegral (sec ts1 - sec ts0) * 10^3 + fromIntegral (nsec ts1 - nsec ts0) / 10^6 + +-- | Given a throttle configuration and the current averaged item +-- size, compute the desired delay between two writes. +computeDelay :: ThrottleConf a -> Double -> Double +computeDelay _ 0 = 0 +computeDelay ThrottleConf { .. } itemSize = + throttleConfInterval / (throttleConfMaxThroughput / itemSize) + +-- | Delay execution. +throttleDelayIn :: Stats -- ^ Current throughput statistics + -> ThrottleConf a -- ^ Throttle configuration + -> Double -- ^ Duration of last write + -> Int -- ^ Resulting delay in microseconds +throttleDelayIn Stats { .. } = throttleDelay (emaCurrent statsEmaItemSizeIn) + +-- | Delay execution. +throttleDelayOut :: Stats -- ^ Current throughput statistics + -> ThrottleConf a -- ^ Throttle configuration + -> Double -- ^ Duration of last write + -> Int -- ^ Resulting delay in microseconds +throttleDelayOut Stats { .. } = throttleDelay (emaCurrent statsEmaItemSizeOut) + +throttleDelay :: Double -- ^ Current item size + -> ThrottleConf a -- ^ Throttle configuration + -> Double -- ^ Duration of last write + -> Int -- ^ Resulting delay in microseconds +throttleDelay itemSize (conf @ ThrottleConf { .. }) duration = + round . subtract duration . (10^3 *) $ computeDelay conf itemSize diff --git a/src/Control/Concurrent/Throttle/Ema.hs b/src/Control/Concurrent/Throttle/Ema.hs new file mode 100644 --- /dev/null +++ b/src/Control/Concurrent/Throttle/Ema.hs @@ -0,0 +1,25 @@ +-- | Exponentially weighted moving average + +{-# LANGUAGE RecordWildCards #-} + +module Control.Concurrent.Throttle.Ema + ( Ema + , emaCurrent + , emaWeight + , newEma + , emaUpdate + ) where + +data Ema = Ema { emaWeight :: !Double + , emaCurrent :: !Double + } deriving (Show) + +newEma :: Double -> Double -> Ema +newEma α initialValue = Ema { emaWeight = α, emaCurrent = initialValue } + +emaModifyCurrent :: (Ema -> Double) -> Ema -> Ema +emaModifyCurrent f ema = ema { emaCurrent = f ema } + +emaUpdate :: Double -> Ema -> Ema +emaUpdate newSample = emaModifyCurrent $ + \ Ema { .. } -> emaWeight * newSample + (1 - emaWeight) * emaCurrent diff --git a/tests/Simple.hs b/tests/Simple.hs new file mode 100644 --- /dev/null +++ b/tests/Simple.hs @@ -0,0 +1,71 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Main where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async +import Control.Concurrent.STM +import Control.Concurrent.Throttle +import Data.ByteString (ByteString) +import qualified Data.ByteString as ByteString +import qualified Data.ByteString.Char8 as ByteString.Char8 +import Data.Function ((&)) +import Data.Text.Encoding +import Say +import Test.Framework (defaultMain, testGroup) +import Test.Framework.Providers.HUnit (testCase) + +newtype Elem = Elem { unElem :: ByteString } deriving (Show) + +measure :: Elem -> Double +measure = fromIntegral . ByteString.length . unElem + +produceNext :: TBQueue Elem -> IO (Maybe Elem) +produceNext queue = do + e <- atomically (readTBQueue queue) + if measure e > 2 + then return Nothing + else return $ Just e + +consumeNext :: Maybe Elem -> IO () +consumeNext (Just e) = say . decodeUtf8 . unElem $ e +consumeNext Nothing = return () + +initialDelay :: Int +initialDelay = 5 * 10^6 + +modifyDelay :: Int -> Int +modifyDelay x = round (0.6 * fromIntegral x) + +producer :: TBQueue Elem -> IO () +producer queue = go 0 initialDelay + where go n delay = do + atomically $ writeTBQueue queue (Elem (ByteString.Char8.pack (show n))) + threadDelay delay + go (n + 1) (modifyDelay delay) + +simpleTest :: IO () +simpleTest = do + let conf = newThrottleConf + & throttleConfSetMeasure measure + & throttleConfThrottleProducer + & throttleConfThrottleConsumer + & throttleConfSetInterval 1000 -- Interval is one second + & throttleConfSetMaxThroughput 100 -- 100 Bytes per interval + & throttleConfSetBufferSize 32 + + queue <- atomically $ newTBQueue 1024 + handle <- throttle conf (produceNext queue) consumeNext + _ <- async (producer queue) + wait handle + + return () + +main :: IO () +main = do + putStrLn "" + defaultMain tests + + where tests = [ testGroup "Simple" + [ testCase "Simple" simpleTest ] + ] diff --git a/throttle-io-stream.cabal b/throttle-io-stream.cabal new file mode 100644 --- /dev/null +++ b/throttle-io-stream.cabal @@ -0,0 +1,65 @@ +-- This file has been generated from package.yaml by hpack version 0.17.0. +-- +-- see: https://github.com/sol/hpack + +name: throttle-io-stream +version: 0.2.0.0 +synopsis: Throttler between arbitrary IO producer and consumer functions +description: This packages provides functionality for throttling IO using general IO callbacks. The throttling depends on a provided configuration. The supported throttling modes are producer throttling, consumer throttling or producer & consumer throttling. +category: Concurrency +homepage: https://github.com/mtesseract/io-throttle#readme +bug-reports: https://github.com/mtesseract/io-throttle/issues +author: Moritz Schulte +maintainer: mtesseract@silverratio.net +copyright: (c) 2017 Moritz Schulte +license: BSD3 +license-file: LICENSE +build-type: Simple +cabal-version: >= 1.10 + +extra-source-files: + README.md + +source-repository head + type: git + location: https://github.com/mtesseract/io-throttle + +library + hs-source-dirs: + src + ghc-options: -Wall -fno-warn-type-defaults + build-depends: + base >=4.7 && <5 + , stm + , async + , stm-chans + , clock + exposed-modules: + Control.Concurrent.Throttle + other-modules: + Control.Concurrent.Throttle.Ema + Paths_throttle_io_stream + default-language: Haskell2010 + +test-suite throttle-io-stream-test + type: exitcode-stdio-1.0 + main-is: Simple.hs + hs-source-dirs: + tests + ghc-options: -Wall -fno-warn-type-defaults -threaded -rtsopts -with-rtsopts=-N -Wall -fno-warn-type-defaults + build-depends: + base >=4.7 && <5 + , stm + , async + , stm-chans + , clock + , base + , throttle-io-stream + , HUnit + , test-framework + , test-framework-hunit + , say + , text + , bytestring + , async + default-language: Haskell2010