6b01f91 · streamly · git.iko.soy

git.iko.soy/streamly

Add a stream retry combinator

commit6b01f917f5161cceef0f82bc56ca58b70e22ef10 authorAdithya Kumar <adithya@composewell.com> dateSun 1 Aug 2021, 23:49 (5 years ago) parent69e3b2c

2 files changed+100

compare with master

unifiedsplit

src/Streamly/Internal/Data/Stream/IsStream/Exception.hs+45 src/Streamly/Internal/Data/Stream/StreamD/Exception.hs+55

src/Streamly/Internal/Data/Stream/IsStream/Exception.hs+45

@@ -19,6 +19,7 @@ module Streamly.Internal.Data.Stream.IsStream.Exception

1919 , finally

2020 , ghandle

2121 , handle

22+ , retry

2223 )

2324where

2425

@@ -26,6 +27,7 @@ import Control.Exception (Exception)

2627import Control.Monad.Catch (MonadCatch)

2728import Control.Monad.IO.Class (MonadIO(..))

2829import Control.Monad.Trans.Control (MonadBaseControl)

30+import Data.Map.Strict (Map)

2931import Streamly.Internal.Data.Stream.StreamD (toStreamD)

3032import Streamly.Internal.Data.Stream.StreamK (IsStream)

3133import Streamly.Internal.Data.SVar (MonadAsync)

@@ -217,3 +219,46 @@ handle :: (IsStream t, MonadCatch m, Exception e)

217219 => (e -> t m a) -> t m a -> t m a

218220handle handler xs =

219221 D.fromStreamD $ D.handle (D.toStreamD . handler) $ D.toStreamD xs

222+

223+

224+-- | @retry@ takes 3 arguments

225+--

226+-- 1. A map @m@ whose keys are exceptions and values are the number of times to

227+-- retry the action given that the exception occurs.

228+--

229+-- 2. A handler @han@ that decides how to handle an exception when the exception

230+-- cannot be retried.

231+--

232+-- 3. The stream itself that we want to run this mechanism on.

233+--

234+-- When evaluating a stream if an exception occurs,

235+--

236+-- 1. The stream evaluation aborts

237+--

238+-- 2. The exception is looked up in @m@

239+--

240+-- a. If the exception exists and the mapped value is > 0 then,

241+--

242+-- i. The value is decreased by 1.

243+--

244+-- ii. The stream is resumed from where the exception was called, retrying

245+-- the action.

246+--

247+-- b. If the exception exists and the mapped value is == 0 then the stream

248+-- evaluation stops.

249+--

250+-- c. If the exception does not exist then we handle the exception using

251+-- @han@.

252+--

253+-- /Internal/

254+--

255+{-# INLINE retry #-}

256+retry :: (IsStream t, MonadCatch m, Exception e, Ord e)

257+ => Map e Int

258+ -- ^ map from exception to retry count

259+ -> (e -> t m a)

260+ -- ^ default handler for those exceptions that are not in the map

261+ -> t m a

262+ -> t m a

263+retry emap handler inp =

264+ D.fromStreamD $ D.retry emap (D.toStreamD . handler) $ D.toStreamD inp

src/Streamly/Internal/Data/Stream/StreamD/Exception.hs+55

@@ -20,6 +20,7 @@ module Streamly.Internal.Data.Stream.StreamD.Exception

2020 , finally

2121 , ghandle

2222 , handle

23+ , retry

2324 )

2425where

2526

@@ -29,11 +30,13 @@ import Control.Exception (Exception, SomeException, mask_)

2930import Control.Monad.Catch (MonadCatch)

3031import Control.Monad.IO.Class (MonadIO(..))

3132import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)

33+import Data.Map.Strict (Map)

3234import GHC.Exts (inline)

3335import Streamly.Internal.Data.IOFinalizer

3436 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)

3537

3638import qualified Control.Monad.Catch as MC

39+import qualified Data.Map.Strict as Map

3740

3841import Streamly.Internal.Data.Stream.StreamD.Type

3942import Streamly.Internal.Data.SVar

@@ -357,3 +360,55 @@ _handle f (Stream step state) = Stream step' (Left state)

357360 Yield x s -> return $ Yield x (Right (Stream step1 s))

358361 Skip s -> return $ Skip (Right (Stream step1 s))

359362 Stop -> return Stop

363+

364+data RetryState emap s1 s2

365+ = RetryWithMap emap s1

366+ | RetryDefault s2

367+

368+-- | See 'Streamly.Internal.Data.Stream.IsStream.retry'

369+--

370+{-# INLINE_NORMAL retry #-}

371+retry

372+ :: forall e m a. (Exception e, Ord e, MonadCatch m)

373+ => Map e Int

374+ -- ^ map from exception to retry count

375+ -> (e -> Stream m a)

376+ -- ^ default handler for those exceptions that are not in the map

377+ -> Stream m a

378+ -> Stream m a

379+retry emap0 defaultHandler (Stream step0 state0) = Stream step state

380+

381+ where

382+

383+ state = RetryWithMap emap0 state0

384+

385+ {-# INLINE_LATE step #-}

386+ step gst (RetryWithMap emap st) = do

387+ eres <- MC.try $ step0 gst st

388+ case eres of

389+ Left e -> handler e emap st

390+ Right res ->

391+ return

392+ $ case res of

393+ Yield x st1 -> Yield x $ RetryWithMap emap st1

394+ Skip st1 -> Skip $ RetryWithMap emap st1

395+ Stop -> Stop

396+ step gst (RetryDefault (UnStream step1 state1)) = do

397+ res <- step1 gst state1

398+ return

399+ $ case res of

400+ Yield x st1 -> Yield x $ RetryDefault (Stream step1 st1)

401+ Skip st1 -> Skip $ RetryDefault (Stream step1 st1)

402+ Stop -> Stop

403+

404+ {-# INLINE handler #-}

405+ handler e emap st =

406+ return

407+ $ Skip

408+ $ case Map.lookup e emap of

409+ Just i

410+ | i > 0 ->

411+ let emap1 = Map.insert e (i - 1) emap

412+ in RetryWithMap emap1 st

413+ | otherwise -> RetryDefault $ defaultHandler e

414+ Nothing -> RetryDefault $ defaultHandler e