STM-based awaitMessage.

1 files changed, 4 insertions, 3 deletions

diff --git a/dht/src/Network/BitTorrent/MainlineDHT.hs b/dht/src/Network/BitTorrent/MainlineDHT.hs
index ed97ee31..0269268f 100644
--- a/dht/src/Network/BitTorrent/MainlineDHT.hs
+++ b/dht/src/Network/BitTorrent/MainlineDHT.hs
@@ -428,8 +428,9 @@ showPacket f addr flow bs = L8.unpack $ L8.unlines es
 addVerbosity :: Transport err SockAddr ByteString -> Transport err SockAddr ByteString
 addVerbosity tr =
     tr { awaitMessage = \kont -> awaitMessage tr $ \m -> do
-            forM_ m $ mapM_ $ \(msg,addr) -> do
-                dput XBitTorrent (showPacket id addr " --> " msg)
+            case m of
+                Arrival addr msg -> dput XBitTorrent (showPacket id addr " --> " msg)
+                _                -> return ()
             kont m
        , sendMessage = \addr msg -> do
             dput XBitTorrent (showPacket id addr " <-- " msg)
@@ -598,7 +599,7 @@ newClient swarms addr = do
         -- recursive since 'updateRouting' does not invoke 'awaitMessage' which
         -- which was modified by 'onInbound'.  However, I'm going to avoid the
         -- mutual reference just to be safe.
-        outgoingClient = client { clientNet = net { awaitMessage = ($ Nothing) } }
+        outgoingClient = client { clientNet = net { awaitMessage = pure . ($ Terminated) } }
 
         dispatch = DispatchMethods
             { classifyInbound = classify -- :: x -> MessageClass err meth tid addr x