1 files changed, 17 insertions, 10 deletions

diff --git a/dht/src/Network/Lossless.hs b/dht/src/Network/Lossless.hs
index 5a313aed..079f4d07 100644
--- a/dht/src/Network/Lossless.hs
+++ b/dht/src/Network/Lossless.hs
@@ -9,6 +9,7 @@
 {-# LANGUAGE TupleSections #-}
 module Network.Lossless where
 
+import Control.Concurrent.STM
 import Control.Concurrent.STM.TChan
 import Control.Monad
 import Control.Monad.STM
@@ -42,7 +43,8 @@ data OutgoingInfo y = OutgoingInfo
 
 -- | Obtain a reliable transport form an unreliable one.
 lossless :: Show addr =>
-    (x -> addr -> IO (PacketInboundEvent (x',addr')))       -- ^ Used to classify newly arrived packets.
+    String -- ^ Label for debugging.
+    -> (x -> addr -> IO (PacketInboundEvent (x',addr')))    -- ^ Used to classify newly arrived packets.
     -> (SequenceInfo -> x' -> addr' -> IO (OutgoingInfo y)) -- ^ Used to encode and classify outbound packets.
     -> addr                                                 -- ^ The remote address for this session.
     -> TransportA String addr x y                           -- ^ An unreliable lossy transport.
@@ -51,7 +53,7 @@ lossless :: Show addr =>
           , [Word32] -> IO ()         -- ^ Use this to request lost packets be re-sent.
           , IO ([Word32],Word32)      -- ^ Use this to discover missing packets to request.
           )
-lossless isLossless encode saddr udp = do
+lossless lbl isLossless encode saddr udp = do
     pb <- atomically newPacketBuffer
     oob <- atomically newTChan -- Out-of-band channel, these packets (or
                                -- errors) bypass the packet buffer to be
@@ -80,11 +82,14 @@ lossless isLossless encode saddr udp = do
                                 report <- pbReport "enqueued" pb
                                 writeTChan oob (ParseError report)
                     loop
-    labelThread rloop ("lossless."++show saddr)
+    labelThread rloop ("lossless."++lbl)
+    term <- newTVarIO retry -- In case awaitMessage is called multiple times beyond termination,
+                            -- we will use this STM action stop it from waiting on the oob TChan.
+                            -- XXX: This shouldn't be neccessary and might be costly.
     let tr = Transport
             { awaitMessage = \kont ->
                 orElse
-                    (do x <- readTChan oob
+                    (do x <- readTChan oob `orElse` join (readTVar term)
                         return $ kont $! x)
                     (do x <- PB.awaitReadyPacket pb
                         report <- pbReport "dequeued" pb
@@ -100,13 +105,13 @@ lossless isLossless encode saddr udp = do
                 (isfull,nn) <-
                     if islossy
                         then do
-                            dput XNetCrypto $ shows saddr $ " <-- Lossy packet " ++ show seqno
+                            dput XNetCrypto $ mappend lbl $ " <-- Lossy packet " ++ show seqno
                             return (False,(0,0)) -- avoid updating seqno on lossy packets.
                         else do
-                            dput XNetCrypto $ shows saddr $ " <-- Lossless packet " ++ show seqno
+                            dput XNetCrypto $ mappend lbl $ " <-- Lossless packet " ++ show seqno
                             atomically $ PB.grokOutboundPacket pb (PacketSent (sequenceNumber seqno) x)
                 when isfull $ do
-                    dput XNetCrypto $ shows saddr $ " <-- Outbound queue is full! Retrying... " ++ show (nn,seqno)
+                    dput XNetCrypto $ mappend lbl $ " <-- Outbound queue is full! Retrying... " ++ show (nn,seqno)
                     atomically $ do
                         (isfull,_) <- PB.grokOutboundPacket pb (PacketSent (sequenceNumber seqno) x)
                         when isfull retry
@@ -114,14 +119,16 @@ lossless isLossless encode saddr udp = do
                 maybe sendit (catchIOError sendit) oops
             , setActive = \case
                 False -> do
-                    atomically $ writeTChan oob Terminated -- quit rloop thread
+                    atomically $ do
+                        writeTChan oob Terminated -- quit rloop thread
+                        writeTVar term (return Terminated)
                     setActive udp False
                 True -> return ()
             }
         resend ns = do
             xs <- atomically $ retrieveForResend pb ns
-            dput XNetCrypto $ shows saddr $ " <-- Resending " ++ show (length xs) ++ " packets."
+            dput XNetCrypto $ mappend lbl $ " <-- Resending " ++ show (length xs) ++ " packets."
             forM_ xs $ \x -> do
-                dput XNetCrypto $ shows saddr $ " <-- Resending packet."
+                dput XNetCrypto $ mappend lbl $ " <-- Resending packet."
                 sendMessage udp saddr . snd $ x
     return (tr, resend, atomically $ PB.packetNumbersToRequest pb)