concurrent-shoutout

module Main where

import System.TimeIt
import Data.ByteString.Lazy qualified as LBS
import Data.ByteString qualified as BS
import Data.ByteString.Lazy (ByteString)
import Crypto.Random (getRandomBytes)
import Control.Monad
import Data.Functor
import Data.Function
import UnliftIO
import Data.List (unfoldr)
import Data.Foldable
import Data.Traversable
import System.FilePath
import Data.Map qualified as Map
import Data.Map (Map)
import Data.Hashable
import Data.Vector qualified as V
import Data.Vector ((!))
import Safe
import Control.Concurrent.STM (flushTQueue)

keySize :: Int
keySize = 32

keysFile :: FilePath
keysFile = "keys"

readKeys :: FilePath -> IO [ByteString]
readKeys fn = do
  lbss <- LBS.readFile fn
  pure $ unfoldr emit lbss
  where
    emit ss =
      let (a,b) = LBS.splitAt (fromIntegral keySize) ss
      in  if LBS.null a then
            Nothing
          else
            Just (a,b)

justReadKeys :: FilePath -> IO ()
justReadKeys fn = do
  keys <- readKeys fn <&> length
  print keys

keysToTVMapSeq :: FilePath -> IO ()
keysToTVMapSeq fn = do
  tvm <- newTVarIO (mempty :: Map ByteString ByteString)
  keys <- readKeys fn
  for_ keys $ \k -> do
    atomically $ modifyTVar' tvm (Map.insert k k)


keysToTVMapCo :: FilePath -> IO ()
keysToTVMapCo fn = do
  tvm <- newTVarIO (mempty :: Map ByteString ByteString)
  keys <- readKeys fn
  forConcurrently_ keys $ \k -> do
    atomically $ modifyTVar' tvm (Map.insert k k)


keysToTVMapBuc :: Int -> FilePath -> IO ()
keysToTVMapBuc n fn = do

  tqs <- replicateM n newTQueueIO
             <&> V.fromList

  keys <- readKeys fn

  tq <- newTQueueIO
  prod <- async do
            for_ keys $ \k -> do

              let z = maybe 0 fst (LBS.uncons k)
              let i = fromIntegral z `mod` n
              let tq = tqs ! i

              atomically $ writeTQueue tq (Just k)

            for_ tqs $ \q -> atomically $ writeTQueue q Nothing

  ps <- for tqs $ \q -> async do
          myBuc <- newTVarIO (mempty :: Map ByteString ByteString)
          fix \next -> do
            r <- atomically $ readTQueue q
            case r of
              Nothing -> pure ()
              Just v  -> do
                atomically $ modifyTVar' myBuc (Map.insert v v)
                next

          pure ()

  mapM_ wait (prod : V.toList ps)


main :: IO ()
main = do
  withTempDirectory "." "shootout" $ \dir -> do

    let fn = dir </> keysFile
    timeItNamed "keys-to-file" do

      withBinaryFile fn AppendMode $ \ha -> do
        for_ [1 .. 100_000 ] $ \_ -> do
          bs <- getRandomBytes keySize <&> LBS.fromStrict
          LBS.hPutStr ha bs

    withFile fn ReadMode hFileSize >>= print

    timeItNamed "justReadKeys" (justReadKeys fn)
    timeItNamed "keysToTVMapSeq" (keysToTVMapSeq fn)
    timeItNamed "keysToTVMapCo"  (keysToTVMapCo fn)
    timeItNamed "keysToTVMapBuc 2" (keysToTVMapBuc 2 fn)
    timeItNamed "keysToTVMapBuc 4" (keysToTVMapBuc 4 fn)
    timeItNamed "keysToTVMapBuc 8" (keysToTVMapBuc 8 fn)
    timeItNamed "keysToTVMapBuc 16" (keysToTVMapBuc 16 fn)