diff --git a/benchmark/Main.hs b/benchmark/Main.hs --- a/benchmark/Main.hs +++ b/benchmark/Main.hs @@ -11,7 +11,13 @@ main = defaultMain $ [ - bench "extractLinesTransform" $ whnfIO $ produceAndConsume + bench "extractLinesConcurrently" $ whnfIO $ produceAndConsume + (Produce.fileBytes "data/2.tsv") + (right' (Consume.transform + (Transform.extractLinesConcurrently numCapabilities) + (Consume.count))) + , + bench "extractLines" $ whnfIO $ produceAndConsume (Produce.fileBytes "data/2.tsv") (right' (Consume.transform (Transform.extractLines) diff --git a/library/Potoki/Core/Transform.hs b/library/Potoki/Core/Transform.hs --- a/library/Potoki/Core/Transform.hs +++ b/library/Potoki/Core/Transform.hs @@ -20,7 +20,10 @@ executeIO, mapInIO, -- * ByteString - module Potoki.Core.Transform.ByteString, + builderChunks, + extractLines, + extractLinesWithoutTrail, + extractLinesConcurrently, -- * State R.runState, R.execState, @@ -35,9 +38,11 @@ N.bufferize, N.bufferizeFlushing, N.concurrently, + N.concurrentlyInOrder, N.unsafeConcurrently, N.async, N.concurrentlyWithBatching, + N.concurrentlyInOrderWithBatching, -- * File IO deleteFile, appendBytesToFile, diff --git a/library/Potoki/Core/Transform/ByteString.hs b/library/Potoki/Core/Transform/ByteString.hs --- a/library/Potoki/Core/Transform/ByteString.hs +++ b/library/Potoki/Core/Transform/ByteString.hs @@ -4,6 +4,7 @@ import Potoki.Core.Prelude hiding (filter) import Potoki.Core.Transform.Basic import Potoki.Core.Transform.Instances () +import Potoki.Core.Transform.Concurrency import Potoki.Core.Types import qualified Potoki.Core.Fetch as A import qualified Potoki.Core.Produce as H @@ -91,3 +92,48 @@ in do writeIORef pokingRef mempty return (Just [bytes]) + +extractLinesConcurrently :: Int -> Transform ByteString ByteString +extractLinesConcurrently concurrency = + concurrentlyInOrderWithBatching 100 concurrency (arr (B.split 10)) >>> mergeLineChunks + +mergeLineChunks :: Transform [ByteString] ByteString +mergeLineChunks = Transform $ \ (Fetch fetchChunkList) -> liftIO $ do + cachedChunkListRef <- newIORef (Just []) + incompleteChunkRef <- newIORef Nothing + return $ Fetch $ let + loop = do + cachedChunkListIfAny <- readIORef cachedChunkListRef + case cachedChunkListIfAny of + Just cachedChunkList -> case cachedChunkList of + (!a) : b : tail -> do + writeIORef cachedChunkListRef (Just (b : tail)) + return (Just a) + (!a) : _ -> do + writeIORef cachedChunkListRef (Just []) + writeIORef incompleteChunkRef (Just a) + loop + _ -> do + newChunkListIfAny <- fetchChunkList + incompleteChunkIfAny <- readIORef incompleteChunkRef + case incompleteChunkIfAny of + Just incompleteChunk -> case newChunkListIfAny of + Just newChunkList -> case newChunkList of + a : b : tail -> do + completeChunk <- evaluate (incompleteChunk <> a) + writeIORef cachedChunkListRef (Just (b : tail)) + writeIORef incompleteChunkRef Nothing + return (Just completeChunk) + a : _ -> do + newIncompleteChunk <- evaluate (incompleteChunk <> a) + writeIORef incompleteChunkRef (Just newIncompleteChunk) + loop + _ -> loop + Nothing -> do + writeIORef cachedChunkListRef Nothing + return (Just incompleteChunk) + Nothing -> do + writeIORef cachedChunkListRef newChunkListIfAny + loop + Nothing -> return Nothing + in loop diff --git a/library/Potoki/Core/Transform/Concurrency.hs b/library/Potoki/Core/Transform/Concurrency.hs --- a/library/Potoki/Core/Transform/Concurrency.hs +++ b/library/Potoki/Core/Transform/Concurrency.hs @@ -128,6 +128,66 @@ else return Nothing in atomically (readChan <|> terminate) +concurrentlyInOrder :: NFData b => Int -> Transform a b -> Transform a b +concurrentlyInOrder concurrency (Transform transform) = Transform $ \ (Fetch fetchA) -> liftIO $ do + inputQueue <- newTBQueueIO concurrency + outputSlotQueue <- newTQueueIO + liveWorkersVar <- newTVarIO concurrency + + forkIO $ let + loop = do + fetchAResult <- fetchA + case fetchAResult of + Just a -> do + atomically $ writeTBQueue inputQueue (Just a) + loop + Nothing -> atomically $ replicateM_ concurrency $ writeTBQueue inputQueue Nothing + in loop + + replicateM_ concurrency $ forkIO $ do + outputQueue <- newTQueueIO + needsSwitchVar <- newTVarIO False + + let + localizedFetchA = Fetch $ atomically $ do + needsSwitch <- readTVar needsSwitchVar + if needsSwitch + then writeTQueue outputQueue Nothing + else writeTVar needsSwitchVar True + writeTQueue outputSlotQueue outputQueue + readTBQueue inputQueue + + in do + (Fetch fetchB, finalize) <- case transform localizedFetchA of M.Acquire io -> io + let + loop = do + fetchBResult <- fetchB + case fetchBResult of + Just b -> do + forcedB <- evaluate (force b) + atomically $ writeTQueue outputQueue (Just forcedB) + loop + Nothing -> do + atomically $ do + writeTQueue outputQueue Nothing + modifyTVar' liveWorkersVar pred + finalize + in loop + + return $ Fetch $ atomically $ fix $ \ loop -> mplus + (do + outputQueue <- peekTQueue outputSlotQueue + bIfAny <- readTQueue outputQueue + case bIfAny of + Just b -> return (Just b) + Nothing -> do + readTQueue outputSlotQueue + loop) + (do + liveWorkers <- readTVar liveWorkersVar + guard (liveWorkers <= 0) + return Nothing) + {-| A transform, which fetches the inputs asynchronously on the specified number of threads. -} @@ -158,4 +218,10 @@ concurrentlyWithBatching batching concurrency transform = batch @Vector batching >>> bufferize concurrency >>> unsafeConcurrently concurrency (vector >>> transform >>> batch @Vector batching) >>> + vector + +concurrentlyInOrderWithBatching :: (NFData b) => Int -> Int -> Transform a b -> Transform a b +concurrentlyInOrderWithBatching batching concurrency transform = + batch @Vector batching >>> + concurrentlyInOrder concurrency (vector >>> transform >>> batch @Vector batching) >>> vector diff --git a/potoki-core.cabal b/potoki-core.cabal --- a/potoki-core.cabal +++ b/potoki-core.cabal @@ -1,5 +1,5 @@ name: potoki-core -version: 2.3.0.2 +version: 2.3.1 synopsis: Low-level components of "potoki" description: Provides everything required for building custom instances of @@ -70,6 +70,7 @@ main-is: Main.hs default-extensions: Arrows, BangPatterns, ConstraintKinds, DataKinds, DefaultSignatures, DeriveDataTypeable, DeriveFoldable, DeriveFunctor, DeriveGeneric, DeriveTraversable, EmptyDataDecls, FlexibleContexts, FlexibleInstances, FunctionalDependencies, GADTs, GeneralizedNewtypeDeriving, InstanceSigs, LambdaCase, LiberalTypeSynonyms, MagicHash, MultiParamTypeClasses, MultiWayIf, NoImplicitPrelude, NoMonomorphismRestriction, OverloadedStrings, PatternGuards, ParallelListComp, QuasiQuotes, RankNTypes, RecordWildCards, ScopedTypeVariables, StandaloneDeriving, TemplateHaskell, TupleSections, TypeApplications, TypeFamilies, TypeOperators, UnboxedTuples default-language: Haskell2010 + ghc-options: -threaded "-with-rtsopts=-N" other-modules: Potoki Transform diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -38,6 +38,10 @@ testProperty "extractLinesWithoutTrail" $ forAll chunksGen $ \ chunks -> I.lines (F.concat chunks) === unsafePerformIO (C.produceAndConsume (E.transform A.extractLinesWithoutTrail (E.list chunks)) D.list) + , + testProperty "extractLinesConcurrently" $ forAll chunksGen $ \ chunks -> + F.split 10 (F.concat chunks) === + unsafePerformIO (C.produceAndConsume (E.transform (A.extractLinesConcurrently numCapabilities) (E.list chunks)) D.list) ] , testCase "extractLinesWithoutTrail" $ do diff --git a/test/Transform.hs b/test/Transform.hs --- a/test/Transform.hs +++ b/test/Transform.hs @@ -21,6 +21,23 @@ transform = testGroup "Transform" $ [ + testGroup "concurrentlyInOrder unit tests" $ let + testList (list :: [[Int]]) = + testCase (show list) $ do + actual <- C.produceAndTransformAndConsume (E.list list) (A.concurrentlyInOrder 4 A.list) D.list + assertEqual "" (concat list) actual + in [ + testList [[]] , + testList [[],[]], + testList [[0, 1, 2], []], + testList [[0, 1, 2], [3]], + testList [[0,3,7,-2],[],[],[-8,7,-8,7,3,-4,1,4],[-4,1],[5,-2,3,-2,-5,6,-7],[-8,6,7,1,8,-8,5]] + ] + , + testProperty "concurrentlyInOrder" $ \ (list :: [[Int]]) -> + concat list === + unsafePerformIO (C.produceAndTransformAndConsume (E.list list) (A.concurrentlyInOrder 4 A.list) D.list) + , testProperty "list" $ \ (list :: [[Int]]) -> concat list === unsafePerformIO (C.produceAndTransformAndConsume (E.list list) A.list D.list)