throttle.hs

import Control.Concurrent           (forkIO, killThread,
                                     threadDelay, newEmptyMVar,
                                     tryTakeMVar, putMVar)
import Control.Monad                (void, forever)
import Control.Concurrent.STM       (atomically)
import Control.Monad.Trans          (MonadIO(..))
import Control.Monad.Trans.Resource (MonadResource, allocate)
import Data.Conduit                 (GInfConduit, awaitE)

throttle :: (MonadResource m) => Int -> GInfConduit a m a
throttle micro = do
  inputThrottle <- liftIO $ newEmptyMVar
  yieldInitialInput
  startThrottleThread inputThrottle
  throttleRestOfInput inputThrottle
  where
    yieldInitialInput = do
      inputE <- awaitE
      case inputE of
        Right input -> yield input
        Left _ -> return ()
    throttleRestOfInput inputThrottle =
      awaitForever $ \input -> do
        result <- liftIO $ tryTakeMVar inputThrottle
        case result of
          Just _ -> yield $ input
          Nothing -> return ()
    startThrottleThread inputThrottle =
      void $ allocate (forkIO $ throttleThread inputThrottle)
                      (killThread)
    throttleThread inputThrottle = forever $ do
      threadDelay micro
      void $ putMVar inputThrottle True