Skip to content

Commit

Permalink
hstream-store: minor improvements
Browse files Browse the repository at this point in the history
  • Loading branch information
4eUeP committed Feb 22, 2024
1 parent 9fbdf4c commit b5a5de7
Showing 1 changed file with 21 additions and 4 deletions.
25 changes: 21 additions & 4 deletions hstream-store/HStream/Store/Stream.hs
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ module HStream.Store.Stream

import Control.Exception (catch, try)
import Control.Monad (filterM, forM, (<=<))
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Bifunctor (bimap)
import Data.Bits (bit)
import Data.Default (def)
Expand All @@ -183,6 +184,7 @@ import qualified Data.Text as T
import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Data.Vector.Algorithms.Intro as V
import qualified Data.Vector.Generic.Mutable as VGM
import Foreign.C (CSize)
import Foreign.ForeignPtr (withForeignPtr)
import GHC.Generics (Generic)
Expand Down Expand Up @@ -539,10 +541,9 @@ listStreamPartitionsOrdered
listStreamPartitionsOrdered client streamid = do
dir_path <- getStreamDirPath streamid
keys <- LD.logDirLogsNames =<< LD.getLogDirectory client dir_path
ps <- forM (V.fromList keys) $ \key -> do
!mvec <- generateFromListM keys $ \key -> do
logId <- getUnderlyingLogId client streamid (Just key)
pure (key, logId)
!mvec <- V.unsafeThaw ps
V.sortBy (\e1 e2 -> compare (snd e1) (snd e2)) mvec
V.unsafeFreeze mvec

Expand All @@ -555,10 +556,9 @@ listStreamPartitionsOrderedByName
listStreamPartitionsOrderedByName client streamid = do
dir_path <- getStreamDirPath streamid
keys <- LD.logDirLogsNames =<< LD.getLogDirectory client dir_path
ps <- forM (V.fromList keys) $ \key -> do
!mvec <- generateFromListM keys $ \key -> do
logId <- getUnderlyingLogId client streamid (Just key)
pure (key, logId)
!mvec <- V.unsafeThaw ps
V.sortBy (\e1 e2 -> compare (fst e1) (fst e2)) mvec
V.unsafeFreeze mvec

Expand Down Expand Up @@ -955,6 +955,23 @@ thawToPosixPath name = P.PosixString (cbytes2sbsUnsafe name)
{-# INLINABLE thawToPosixPath #-}
#endif

-- Variant of 'VGM.generateM'
--
-- https://hackage.haskell.org/package/vector-0.13.0.0/docs/Data-Vector-Generic-Mutable.html#v:generateM
generateFromListM
:: (PrimMonad m, VGM.MVector v b)
=> [a] -> (a -> m b) -> m (v (PrimState m) b)
generateFromListM es f
| length es <= 0 = VGM.new 0
| otherwise = do
let n = length es
vec <- VGM.new n
let loop (i, xs) | i >= n = return vec
| otherwise = do VGM.unsafeWrite vec i =<< f (head xs)
loop (i + 1, tail xs)
loop (0, es)
{-# INLINE generateFromListM #-}

-------------------------------------------------------------------------------

#undef HSTREAM_USE_LOCAL_STREAM_CACHE

0 comments on commit b5a5de7

Please sign in to comment.