trace-forward: keep helper functions simple and all logic in the same place

Author: fmaste

trace-forward/src/Trace/Forward/Utils/TraceObject.hs

@@ -101,13 +101,35 @@ writeToSinkSTM qTVar disconnectedSize connectedSize wasUsedTVar traceObject = do
     (q', res) <- ((,) <$> isFullTBQueue q <*> isEmptyTBQueue q) >>= \case
       -- The `TBQueue` is full.
       (True, _)    -> do
-                          flushQueue q
+                          -- Only get the length when you are sure it's needed.
+                          qLen <- lengthTBQueue q
+                          if fromIntegral qLen == connectedSize
+                          then do
+                            -- The queue is full, but if it's a small queue, we
+                            -- can switch it to a big one and give a chance not
+                            -- to flush items to stdout yet.
+                            q' <- growQueue q
+                            pure (q', [])
+                          else do
+                            -- The big queue is full, will flush it to stdout.
+                            res <- flushTBQueue q
+                            pure (q, res)
       -- The `TBQueue` is not full and the `TBQueue` is empty.
       (_,    True) -> do
                           wasUsed <- readTVar wasUsedTVar
                           q' <- if wasUsed
-                                then shrinkQueue q
-                                else pure q
+                                then do
+                                  -- Get the length when sure it's needed.
+                                  qLen <- lengthTBQueue q
+                                  if fromIntegral qLen == disconnectedSize
+                                  -- If the sink was used, now we can shrink it.
+                                  then switchQueue connectedSize
+                                  -- Queue is already of the desired capacity.
+                                  else pure q
+                                else do
+                                  -- The sink was not used, nothing to do here.
+                                  pure q
+                          -- This case never flushes.
                           pure (q', [])
       -- The `TBQueue` is not full and the `TBQueue` is not empty.
       (_,    _)    -> do
@@ -118,32 +140,14 @@ writeToSinkSTM qTVar disconnectedSize connectedSize wasUsedTVar traceObject = do
 
  where
 
-  -- The queue is full, but if it's a small queue, we can switch it
-  -- to a big one and give a chance not to flush items to stdout yet.
-  flushQueue q = do
-    qLen <- lengthTBQueue q
-    if fromIntegral qLen == connectedSize
-      then do
-        -- The small queue is full, so we have to switch to a big one and
-        -- then flush collected items from the small queue and store them in
-        -- a big one.
-
-        acceptedItems <- -- trace ("growQueue disconnected" ++ show disconnectedSize) $
-                          flushTBQueue q
-        bigQ <- switchQueue disconnectedSize
-        mapM_ (writeTBQueue bigQ) acceptedItems
-        pure (bigQ, [])
-      else do
-        -- The big queue is full, we have to flush it to stdout.
-        res <- flushTBQueue q
-        pure (q, res)
-
-  -- if the sink was used and it
-  shrinkQueue q = do
-    qLen <- lengthTBQueue q
-    if fromIntegral qLen == disconnectedSize
-    then switchQueue connectedSize
-    else pure q
+  -- The small queue is full, so we have to switch to a big one and
+  -- then flush collected items from the small queue and store them in
+  -- a big one.
+  growQueue q = do
+    acceptedItems <- flushTBQueue q
+    bigQ <- switchQueue disconnectedSize
+    mapM_ (writeTBQueue bigQ) acceptedItems
+    pure bigQ
 
   switchQueue size = do
     q' <- newTBQueue (fromIntegral size)