6b01f91 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add a stream retry combinator

6b01f91Adithya Kumar5 years ago

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

compare with master

unifiedsplit

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

19 , finally

20 , ghandle

21 , handle

22+ , retry

22 )

23 )

23where

24where

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

26import Control.Monad.Catch (MonadCatch)

27import Control.Monad.Catch (MonadCatch)

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

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

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

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

30+import Data.Map.Strict (Map)

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

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

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

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

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

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

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

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

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

218handle handler xs =

220handle handler xs =

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

221 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

20 , finally

21 , ghandle

22 , handle

23+ , retry

23 )

24 )

24where

25where

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

29import Control.Monad.Catch (MonadCatch)

30import Control.Monad.Catch (MonadCatch)

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

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

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

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

33+import Data.Map.Strict (Map)

32import GHC.Exts (inline)

34import GHC.Exts (inline)

33import Streamly.Internal.Data.IOFinalizer

35import Streamly.Internal.Data.IOFinalizer

34 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)

36 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)

36import qualified Control.Monad.Catch as MC

38import qualified Control.Monad.Catch as MC

39+import qualified Data.Map.Strict as Map

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

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

39import Streamly.Internal.Data.SVar

42import Streamly.Internal.Data.SVar

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

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

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

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

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

359 Stop -> return Stop

362 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