experimental dequeueOrGetMissing function

1 files changed, 18 insertions, 0 deletions

diff --git a/src/Data/PacketQueue.hs b/src/Data/PacketQueue.hs
index f9d9f28f..330934cd 100644
--- a/src/Data/PacketQueue.hs
+++ b/src/Data/PacketQueue.hs
@@ -93,6 +93,24 @@ observeOutOfBand PacketQueue { seqno, qsize, buffend } no = do
     when ( proj < qsize) $ do
     modifyTVar' buffend (\be -> if be - low <= proj then no + 1 else be)
 
+-- | If buffend < seqno then return expected packet numbers for all
+--   the Nothings in the array between them.
+--   Otherwise, behave as 'dequeue' would.
+dequeueOrGetMissing :: PacketQueue a -> STM (Either [Word32] a)
+dequeueOrGetMissing PacketQueue { pktq, seqno, qsize, buffend } = do
+    i0 <- readTVar seqno
+    be <- readTVar buffend
+    if i0 < be
+      then do
+        maybes <- mapM (readArray pktq) (take (fromIntegral qsize) $ map (`mod` qsize) [ be .. i0 ])
+        let nums = map fst . filter (isNothing . snd) $ zip [be ..]  maybes
+        return (Left nums)
+      else do
+        let i = i0 `mod` qsize
+        x <- maybe retry return =<< readArray pktq i
+        writeArray pktq i Nothing
+        modifyTVar' seqno   succ
+        return (Right x)
 
 -- | Retry until the next expected packet is enqueued.  Then return it.
 dequeue :: PacketQueue a -> STM a