1 files changed, 17 insertions, 6 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 41e25486..7244304e 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -78,23 +78,34 @@ partitionTransport :: ((b,a) -> Either (x,xaddr) (y,yaddr))
                       -> ((y,yaddr) -> (b,a))
                       -> Transport err a b
                       -> IO (Transport err xaddr x, Transport err yaddr y)
-partitionTransport parse encodex encodey tr = do
+partitionTransport parse encodex encodey tr =
+    partitionTransportM (return . parse) (return . encodex) (return . encodey) tr
+
+-- | Paritions a 'Transport' into two higher-level transports.  Note: An 'MVar'
+-- is used to share the same underlying socket, so be sure to fork a thread for
+-- both returned 'Transport's to avoid hanging.
+partitionTransportM :: ((b,a) -> IO (Either (x,xaddr) (y,yaddr))) 
+                      -> ((x,xaddr) -> IO (b,a))
+                      -> ((y,yaddr) -> IO (b,a))
+                      -> Transport err a b
+                      -> IO (Transport err xaddr x, Transport err yaddr y)
+partitionTransportM parse encodex encodey tr = do
     mvar <- newEmptyMVar
     let xtr = tr { awaitMessage = \kont -> fix $ \again -> do
                     awaitMessage tr $ \m -> case m of
-                        Just (Right msg) -> either (kont . Just . Right)
-                                                   (\y -> putMVar mvar y >> again)
-                                                $ parse msg
+                        Just (Right msg) -> parse msg >>=
+                                                either (kont . Just . Right)
+                                                       (\y -> putMVar mvar y >> again)
                         Just (Left e)    -> kont $ Just (Left e)
                         Nothing          -> kont Nothing
                  , sendMessage = \addr' msg' -> do
-                    let (msg,addr) = encodex (msg',addr')
+                    (msg,addr) <- encodex (msg',addr')
                     sendMessage tr addr msg
                  }
         ytr = Transport
                 { awaitMessage = \kont -> takeMVar mvar >>= kont . Just . Right
                 , sendMessage = \addr' msg' -> do
-                    let (msg,addr) = encodey (msg',addr')
+                    (msg,addr) <- encodey (msg',addr')
                     sendMessage tr addr msg
                 , closeTransport = return ()
                 }