STM-based awaitMessage.

1 files changed, 18 insertions, 15 deletions

diff --git a/dht/src/Network/Lossless.hs b/dht/src/Network/Lossless.hs
index 861792ab..5a313aed 100644
--- a/dht/src/Network/Lossless.hs
+++ b/dht/src/Network/Lossless.hs
@@ -59,37 +59,38 @@ lossless isLossless encode saddr udp = do
     rloop <- forkIO $ do
         -- This thread enqueues inbound packets or writes them to the oob
         -- channel.
-        myThreadId >>= flip labelThread ("lossless."++show saddr)
-        fix $ \loop -> do
-        awaitMessage udp $ \m -> do
-            m' <- mapM (mapM $ uncurry isLossless) m
+        fix $ \loop -> join $ atomically $ awaitMessage udp $ \m -> do
+            m' <- case m of Terminated   -> return Nothing
+                            ParseError e -> return $ Just (Left e)
+                            Arrival a x  -> Just . Right <$> isLossless x a
             case m' of
                 Nothing -> do
-                    atomically $ writeTChan oob Nothing
+                    atomically $ writeTChan oob Terminated
                     -- Quit thread here.
                 Just (Left e) -> do
-                    atomically $ writeTChan oob (Just $ Left e)
+                    atomically $ writeTChan oob (ParseError e)
                     loop
                 Just (Right event) -> do
                     atomically $ do
                         -- x' <- isLossless xaddr x
                         PB.grokInboundPacket pb event
                         case event of
-                            PacketReceivedLossy {} -> writeTChan oob (Just $ Right $ peReceivedPayload event)
+                            PacketReceivedLossy {} -> writeTChan oob (uncurry (flip Arrival) $ peReceivedPayload event)
                             _                      -> do
                                 report <- pbReport "enqueued" pb
-                                writeTChan oob (Just $ Left report)
+                                writeTChan oob (ParseError report)
                     loop
+    labelThread rloop ("lossless."++show saddr)
     let tr = Transport
-            { awaitMessage = \kont -> do
-                join $ atomically $ orElse
+            { awaitMessage = \kont ->
+                orElse
                     (do x <- readTChan oob
                         return $ kont $! x)
                     (do x <- PB.awaitReadyPacket pb
                         report <- pbReport "dequeued" pb
                         return $ do
-                            atomically $ writeTChan oob (Just $ Left report)
-                            kont $! Just (Right x))
+                            atomically $ writeTChan oob (ParseError report)
+                            kont $! uncurry (flip Arrival) x)
             , sendMessage = \a' x' -> do
                 seqno <- atomically $ do
                     seqno <- PB.nextToSendSequenceNumber pb
@@ -111,9 +112,11 @@ lossless isLossless encode saddr udp = do
                         when isfull retry
                 let sendit = sendMessage udp saddr x
                 maybe sendit (catchIOError sendit) oops
-            , closeTransport = do
-                atomically $ writeTChan oob Nothing -- quit rloop thread
-                closeTransport udp
+            , setActive = \case
+                False -> do
+                    atomically $ writeTChan oob Terminated -- quit rloop thread
+                    setActive udp False
+                True -> return ()
             }
         resend ns = do
             xs <- atomically $ retrieveForResend pb ns