fb-util-0.1.0.1: Util/Async.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

-- | This module provides some higher level concurrency facilities
module Util.Async
  ( window_,
    windowUnorderedReduce,
    runHereOrConcurrently
  ) where

import Control.Concurrent.Async
import Control.Concurrent.Chan
import Control.Monad (forM_, replicateM_)

-- | This function is similar to folly::window. It takes a list of input items
--   and launch a pool of `maxConcurrency` of workers to process them
window_ :: Int -> (a -> IO ()) -> [a] -> IO ()
window_ maxConcurrency process tasks = do
  queue <- newChan
  forM_ tasks $ writeChan queue . Just
  replicateM_ maxConcurrency $ writeChan queue Nothing
  replicateConcurrently_ maxConcurrency $ runWorker queue process
  where
    runWorker :: Chan (Maybe a) -> (a -> IO ()) -> IO ()
    runWorker queue process' = do
      token <- readChan queue
      case token of
        Just task -> do
          process' task
          runWorker queue process'
        Nothing -> return ()


-- | A variant of window function where we collect the result without
--   preserving the order
windowUnorderedReduce
  :: Monoid b => Int -> (a -> IO b) -> [a] -> IO [b]
windowUnorderedReduce maxConcurrency process tasks = do
  queue <- newChan
  forM_ tasks $ writeChan queue . Just
  replicateM_ maxConcurrency $ writeChan queue Nothing
  replicateConcurrently maxConcurrency
    $ runWorker queue mempty
  where
    runWorker queue acc = do
      token <- readChan queue
      case token of
        Just task -> do
          result <- process task
          runWorker queue $ mappend result acc
        Nothing -> return acc

-- | A wrapper for 'mapConcurrently' so that a singleton list of
-- actions runs in this thread, instead of asynchronously
runHereOrConcurrently :: [IO a] -> IO [a]
runHereOrConcurrently [] = return []
runHereOrConcurrently [x] = (:[]) <$> x
runHereOrConcurrently xs = mapConcurrently id xs