Fix tests broken by retry on reading SBQueue (#8)

src/Control/Monad/Trace.hs

@@ -31,6 +31,7 @@ module Control.Monad.Trace (
   newSBQueueIO,
   isEmptySBQueue,
   readSBQueue,
+  readSBQueueOnce,
   writeSBQueue
 ) where
 
@@ -104,14 +105,24 @@ isEmptySBQueue (SBQueue queue count _capacity) = do
   assert (if isEmpty then numElems == 0 else numElems > 0) $
     return isEmpty
 
--- | Read all the values stored in an 'SBQueue'.
-readSBQueue :: SBQueue a -> STM [a]
-readSBQueue (SBQueue queue count _capacity) = do
+data ShouldRetry = Retry | OnlyOnce
+  deriving (Eq)
+
+readSBQueue' :: ShouldRetry -> SBQueue a -> STM [a]
+readSBQueue' shouldRetry (SBQueue queue count _capacity) = do
   elems <- readTVar queue
-  when (null elems) retry
+  when (null elems && shouldRetry == Retry) retry
   writeTVar queue []
   writeTVar count 0
-  return $ reverse elems
+  pure $ reverse elems
+
+-- | Read all the values stored in an 'SBQueue'. Retry if none are available.
+readSBQueue :: SBQueue a -> STM [a]
+readSBQueue = readSBQueue' Retry
+
+-- | A non-retrying version of readSBQueue
+readSBQueueOnce :: SBQueue a -> STM [a]
+readSBQueueOnce = readSBQueue' OnlyOnce
 
 -- | Write a value to an 'SBQueue'.
 writeSBQueue :: SBQueue a -> a -> STM ()

src/Monitor/Tracing/Local.hs

@@ -34,11 +34,11 @@ collectSpanSamples actn = do
     samplesTC = spanSamples tracer
     pendingTV = pendingSpanCount tracer
   liftIO $ fix $ \loop -> do
-    (mbSample, pending) <- atomically $ (,) <$> readSBQueue samplesTC <*> readTVar pendingTV
+    (mbSample, pending) <- atomically $ (,) <$> readSBQueueOnce samplesTC <*> readTVar pendingTV
     case mbSample of
       (x:xs) -> mapM_ addSample (x:xs) >> loop
       [] | pending > 0 -> do
-        toAdd <- liftIO (atomically $ readSBQueue samplesTC)
+        toAdd <- liftIO (atomically $ readSBQueueOnce samplesTC)
         mapM_ addSample toAdd
         loop
       _ -> pure ()