3020130 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

refactor to use IORef and MVar interchangeably

3020130Harendra Kumar8 years ago

src/Streamly/SVar.hs+49−14

compare with master

unifiedsplit

@@ -71,6 +71,10 @@ module Streamly.SVar

71 , queueEmptyAhead

72 , dequeueAhead

73 , dequeueFromHeap

74+ , dequeueFromHeapSeq

75+ , requeueOnHeapTop

76+ , updateHeapSeq

77+ , withIORef

75 , Rate (..)

79 , Rate (..)

76 , getYieldRateInfo

80 , getYieldRateInfo

@@ -102,8 +106,7 @@ where

102import Control.Concurrent

106import Control.Concurrent

103 (ThreadId, myThreadId, threadDelay, getNumCapabilities, throwTo)

107 (ThreadId, myThreadId, threadDelay, getNumCapabilities, throwTo)

104import Control.Concurrent.MVar

108import Control.Concurrent.MVar

105− (MVar, newEmptyMVar, tryPutMVar, takeMVar, newMVar, readMVar,

109+ (MVar, newEmptyMVar, tryPutMVar, takeMVar, newMVar)

106− modifyMVar, modifyMVar_, withMVar)

107import Control.Exception (SomeException(..), catch, mask, assert, Exception)

110import Control.Exception (SomeException(..), catch, mask, assert, Exception)

108import Control.Monad (when)

111import Control.Monad (when)

109import Control.Monad.Catch (MonadThrow)

112import Control.Monad.Catch (MonadThrow)

@@ -1138,29 +1141,61 @@ dequeueAhead q = liftIO $ do

1138dequeueFromHeap

1141dequeueFromHeap

1139 :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Int)

1142 :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Int)

1140 -> IO (Maybe (Entry Int (AheadHeapEntry t m a)))

1143 -> IO (Maybe (Entry Int (AheadHeapEntry t m a)))

1141−dequeueFromHeap hpRef = do

1144+dequeueFromHeap hpVar =

1142− atomicModifyIORef hpRef $ \hp@(h, snum) -> do

1145+ atomicModifyIORef hpVar $ \pair@(hp, snum) -> do

1143− let r = H.uncons h

1146+ let r = H.uncons hp

1144 case r of

1147 case r of

1145− Nothing -> (hp, Nothing)

1146− Just (ent@(Entry seqNo _ev), hp') ->

1148+ Just (ent@(Entry seqNo _ev), hp') | seqNo == snum ->

1147− if (seqNo == snum)

1148− then ((hp', seqNo), Just ent)

1149+ ((hp', snum), Just ent)

1149− else (hp, Nothing)

1150+ _ -> (pair, Nothing)

1151+

1152+withIORef :: IORef a -> (a -> IO b) -> IO b

1153+withIORef ref f = readIORef ref >>= f

1154+

1155+atomicModifyIORef_ :: IORef a -> (a -> a) -> IO ()

1156+atomicModifyIORef_ ref f =

1157+ atomicModifyIORef ref $ \x -> (f x, ())

1150

1158

1151{-# INLINE canDequeueFromHeap #-}

1159{-# INLINE canDequeueFromHeap #-}

1152canDequeueFromHeap

1160canDequeueFromHeap

1153 :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Int)

1161 :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Int)

1154 -> IO Bool

1162 -> IO Bool

1155canDequeueFromHeap hpVar =

1163canDequeueFromHeap hpVar =

1156− withIORef hpVar $ \pair@(hp, snum) -> do

1164+ withIORef hpVar $ \(hp, snum) -> do

1157 let r = H.uncons hp

1165 let r = H.uncons hp

1158 case r of

1166 case r of

1159− Just (ent@(Entry seqNo _ev), hp') -> return $ seqNo == snum

1167+ Just ((Entry seqNo _ev), _) -> return $ seqNo == snum

1160 _ -> return False

1168 _ -> return False

1161− where

1162

1169

1163− withIORef ref f = readIORef ref >>= f

1170+{-# INLINE dequeueFromHeapSeq #-}

1171+dequeueFromHeapSeq

1172+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Int)

1173+ -> Int

1174+ -> IO (Maybe (Entry Int (AheadHeapEntry t m a)))

1175+dequeueFromHeapSeq hpVar snum =

1176+ atomicModifyIORef hpVar $ \(hp, _) -> do

1177+ let r = H.uncons hp

1178+ case r of

1179+ Just (ent@(Entry seqNo _ev), hp') | seqNo == snum ->

1180+ ((hp', snum), Just ent)

1181+ _ -> ((hp, snum), Nothing)

1182+

1183+{-# INLINE requeueOnHeapTop #-}

1184+requeueOnHeapTop

1185+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Int)

1186+ -> Entry Int (AheadHeapEntry t m a)

1187+ -> Int

1188+ -> IO ()

1189+requeueOnHeapTop hpVar ent seqNo =

1190+ atomicModifyIORef_ hpVar $ \(hp, _) -> (H.insert ent hp, seqNo)

1191+

1192+{-# INLINE updateHeapSeq #-}

1193+updateHeapSeq

1194+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Int)

1195+ -> Int

1196+ -> IO ()

1197+updateHeapSeq hpVar seqNo =

1198+ atomicModifyIORef_ hpVar $ \(hp, _) -> (hp, seqNo)

1164

1199

1165-------------------------------------------------------------------------------

1200-------------------------------------------------------------------------------

1166-- WAhead

1201-- WAhead

src/Streamly/Streams/Ahead.hs+52−52

@@ -169,23 +169,23 @@ preStopCheck ::

169preStopCheck sv heap = do

170 -- check the stop condition under a lock before actually

171 -- stopping so that the whole herd does not stop at once.

172+ withIORef heap $ \(hp, _) -> do

173+ heapOk <- underMaxHeap sv hp

172− takeMVar (workerStopMVar sv)

174+ takeMVar (workerStopMVar sv)

173− let stop = do

175+ let stop = do

174− putMVar (workerStopMVar sv) ()

176+ putMVar (workerStopMVar sv) ()

175− return True

177+ return True

176− continue = do

178+ continue = do

177− putMVar (workerStopMVar sv) ()

179+ putMVar (workerStopMVar sv) ()

178− return False

180+ return False

179− (hp, _) <- readIORef heap

180− heapOk <- underMaxHeap sv hp

181− if heapOk

181+ if heapOk

182− then

182+ then

183− case yieldRateInfo sv of

183+ case yieldRateInfo sv of

184− Nothing -> continue

184+ Nothing -> continue

185− Just yinfo -> do

185+ Just yinfo -> do

186− rateOk <- isBeyondMaxRate sv yinfo

186+ rateOk <- isBeyondMaxRate sv yinfo

187− if rateOk then continue else stop

187+ if rateOk then continue else stop

188− else stop

188+ else stop

189

190processHeap :: MonadIO m

191 => IORef ([Stream m a], Int)

@@ -206,19 +206,11 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry

206 if stopIt

207 then liftIO $ do

208 -- put the entry back in the heap and stop

209− atomicModifyIORef heap $ \(h, _) ->

210− ((H.insert (Entry seqNo ent) h, seqNo), ())

209+ liftIO $ requeueOnHeapTop heap (Entry seqNo ent) seqNo

211 sendStop sv winfo

210 sendStop sv winfo

212 else runStreamWithYieldLimit True seqNo r

211 else runStreamWithYieldLimit True seqNo r

213

212

214 loopHeap seqNo ent = do

213 loopHeap seqNo ent = do

215−#ifdef DIAGNOSTICS

216− liftIO $ do

217− maxHp <- readIORef (maxHeapSize $ svarStats sv)

218− (hp, _) <- readIORef heap

219− when (H.size hp > maxHp) $ writeIORef (maxHeapSize $ svarStats sv)

220− (H.size hp)

221−#endif

222 case ent of

214 case ent of

223 AheadEntryPure a -> do

215 AheadEntryPure a -> do

224 -- Use 'send' directly so that we do not account this in worker

216 -- Use 'send' directly so that we do not account this in worker

@@ -233,10 +225,7 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry

233 else runStreamWithYieldLimit True seqNo r

225 else runStreamWithYieldLimit True seqNo r

234

226

235 nextHeap prevSeqNo = do

227 nextHeap prevSeqNo = do

236− -- XXX use "dequeueIfSeqential prevSeqNo" instead of always

237− -- updating the sequence number in heap.

238− liftIO $ atomicModifyIORef heap $ \(h, _) -> ((h, prevSeqNo + 1), ())

228+ ent <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)

239− ent <- liftIO $ dequeueFromHeap heap

240 case ent of

229 case ent of

241 Just (Entry seqNo hent) -> loopHeap seqNo hent

230 Just (Entry seqNo hent) -> loopHeap seqNo hent

242 Nothing -> do

231 Nothing -> do

@@ -291,8 +280,8 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry

291 (singleStreamFromHeap seqNo)

280 (singleStreamFromHeap seqNo)

292 (yieldStreamFromHeap seqNo)

281 (yieldStreamFromHeap seqNo)

293 else liftIO $ do

282 else liftIO $ do

294− atomicModifyIORef heap $ \(h, _) ->

295− ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())

283+ let ent = Entry seqNo (AheadEntryStream r)

284+ liftIO $ requeueOnHeapTop heap ent seqNo

296 incrementYieldLimit sv

285 incrementYieldLimit sv

297 sendStop sv winfo

286 sendStop sv winfo

298

287

@@ -315,6 +304,8 @@ drainHeap q heap st sv winfo = do

315 Just (Entry seqNo hent) ->

304 Just (Entry seqNo hent) ->

316 processHeap q heap st sv winfo hent seqNo True

305 processHeap q heap st sv winfo hent seqNo True

317

306

307+data HeapStatus = HContinue | HStop

308+

318processWithoutToken :: MonadIO m

309processWithoutToken :: MonadIO m

319 => IORef ([Stream m a], Int)

310 => IORef ([Stream m a], Int)

320 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)

311 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)

@@ -340,25 +331,37 @@ processWithoutToken q heap st sv winfo m sno = do

340 -- modification, otherwise contention and retries can make a thread

331 -- modification, otherwise contention and retries can make a thread

341 -- context switch and throw it behind other threads which come later in

332 -- context switch and throw it behind other threads which come later in

342 -- sequence.

333 -- sequence.

343− hp <- liftIO $ atomicModifyIORef heap $ \(h, snum) ->

334+ newHp <- liftIO $ atomicModifyIORef heap $ \(hp, snum) ->

344− ((H.insert (Entry seqNo ent) h, snum), h)

335+ let hp' = H.insert (Entry seqNo ent) hp

336+ in ((hp', snum), hp')

345

337

338+#ifdef DIAGNOSTICS

339+ liftIO $ do

340+ maxHp <- readIORef (maxHeapSize $ svarStats sv)

341+ when (H.size newHp > maxHp) $

342+ writeIORef (maxHeapSize $ svarStats sv) (H.size newHp)

343+#endif

346− heapOk <- liftIO $ underMaxHeap sv hp

344+ heapOk <- liftIO $ underMaxHeap sv newHp

347− let keepDraining = drainHeap q heap st sv winfo

345+ let drainAndStop = drainHeap q heap st sv winfo

348 mainLoop = workLoopAhead q heap st sv winfo

346 mainLoop = workLoopAhead q heap st sv winfo

349− if heapOk

347+ status <-

350− then

351 case yieldRateInfo sv of

348 case yieldRateInfo sv of

352− Nothing -> mainLoop

349+ Nothing -> return HContinue

353 Just yinfo -> do

350 Just yinfo -> do

354 case winfo of

351 case winfo of

355 Just info -> do

352 Just info -> do

356 rateOk <- liftIO $ workerRateControl sv yinfo info

353 rateOk <- liftIO $ workerRateControl sv yinfo info

357 if rateOk

354 if rateOk

358− then mainLoop

355+ then return HContinue

359− else keepDraining

356+ else return HStop

357+ Nothing -> return HContinue

358+

359+ if heapOk

360+ then

361+ case status of

360− Nothing -> mainLoop

362+ HContinue -> mainLoop

361− else keepDraining

363+ HStop -> drainAndStop

364+ else drainAndStop

362

365

363 singleToHeap seqNo a = toHeap seqNo (AheadEntryPure a)

366 singleToHeap seqNo a = toHeap seqNo (AheadEntryPure a)

364 yieldToHeap seqNo a r = toHeap seqNo (AheadEntryStream (a `K.cons` r))

367 yieldToHeap seqNo a r = toHeap seqNo (AheadEntryStream (a `K.cons` r))

@@ -388,7 +391,7 @@ processWithToken q heap st sv winfo action sno = do

388 if continue

391 if continue

389 then loopWithToken seqNo

392 then loopWithToken seqNo

390 else do

393 else do

391− liftIO $ atomicModifyIORef heap $ \(h, _) -> ((h, seqNo + 1), ())

394+ liftIO $ updateHeapSeq heap (seqNo + 1)

392 drainHeap q heap st sv winfo

395 drainHeap q heap st sv winfo

393

396

394 -- XXX use a wrapper function around stop so that we never miss

397 -- XXX use a wrapper function around stop so that we never miss

@@ -406,8 +409,8 @@ processWithToken q heap st sv winfo action sno = do

406 (singleOutput seqNo)

409 (singleOutput seqNo)

407 (yieldOutput seqNo)

410 (yieldOutput seqNo)

408 else do

411 else do

409− liftIO $ atomicModifyIORef heap $ \(h, _) ->

410− ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())

412+ let ent = Entry seqNo (AheadEntryStream r)

413+ liftIO $ requeueOnHeapTop heap ent seqNo

411 liftIO $ incrementYieldLimit sv

414 liftIO $ incrementYieldLimit sv

412 drainHeap q heap st sv winfo

415 drainHeap q heap st sv winfo

413

416

@@ -415,8 +418,7 @@ processWithToken q heap st sv winfo action sno = do

415 work <- dequeueAhead q

418 work <- dequeueAhead q

416 case work of

419 case work of

417 Nothing -> do

420 Nothing -> do

418− liftIO $ atomicModifyIORef heap $ \(h, _) ->

419− ((h, prevSeqNo + 1), ())

421+ liftIO $ updateHeapSeq heap (prevSeqNo + 1)

420 workLoopAhead q heap st sv winfo

422 workLoopAhead q heap st sv winfo

421

423

422 Just (m, seqNo) -> do

424 Just (m, seqNo) -> do

@@ -432,8 +434,7 @@ processWithToken q heap st sv winfo action sno = do

432 (singleOutput seqNo)

434 (singleOutput seqNo)

433 (yieldOutput seqNo)

435 (yieldOutput seqNo)

434 else do

436 else do

435− liftIO $ atomicModifyIORef heap $ \(h, _) ->

436− ((h, prevSeqNo + 1), ())

437+ liftIO $ updateHeapSeq heap (prevSeqNo + 1)

437 liftIO (incrementYieldLimit sv)

438 liftIO (incrementYieldLimit sv)

438 -- To avoid a race when another thread puts something

439 -- To avoid a race when another thread puts something

439 -- on the heap and goes away, the consumer will not get

440 -- on the heap and goes away, the consumer will not get

@@ -445,8 +446,7 @@ processWithToken q heap st sv winfo action sno = do

445 liftIO $ reEnqueueAhead sv q m

446 liftIO $ reEnqueueAhead sv q m

446 workLoopAhead q heap st sv winfo

447 workLoopAhead q heap st sv winfo

447 else do

448 else do

448− liftIO $ atomicModifyIORef heap $ \(h, _) ->

449− ((h, prevSeqNo + 1), ())

449+ liftIO $ updateHeapSeq heap (prevSeqNo + 1)

450 liftIO $ reEnqueueAhead sv q m

451 liftIO $ incrementYieldLimit sv

452 drainHeap q heap st sv winfo