1 files changed, 43 insertions, 3 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 4e110ec3..0fbbc929 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -134,6 +134,35 @@ partitionTransportM parse encodex tr = do
                 }
     return (xtr, ytr)
 
+partitionAndForkTransport ::
+                      (dst -> msg -> IO ())
+                      -> ((b,a) -> IO (Either (x,xaddr) (b,a)))
+                      -> ((x,xaddr) -> IO (Maybe (Either (msg,dst) (b,a))))
+                      -> Transport err a b
+                      -> IO (Transport err xaddr x, Transport err a b)
+partitionAndForkTransport forkedSend parse encodex tr = do
+    mvar <- newEmptyMVar
+    let xtr = tr { awaitMessage = \kont -> fix $ \again -> do
+                    awaitMessage tr $ \m -> case m of
+                        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
+                    msg_addr <- encodex (msg',addr')
+                    case msg_addr of
+                        Just (Right (b,a))    -> sendMessage tr a b
+                        Just (Left (msg,dst)) -> forkedSend dst msg
+                        Nothing               -> return ()
+                 }
+        ytr = Transport
+                { awaitMessage   = \kont -> takeMVar mvar >>= kont . Just . Right
+                , sendMessage    = sendMessage tr
+                , closeTransport = return ()
+                }
+    return (xtr, ytr)
+
 -- |
 --      * f add x --> Nothing, consume x
 --                --> Just id, leave x to a different handler
@@ -376,16 +405,27 @@ transactionMethods ::
     TableMethods t tid  -- ^ Table methods to lookup values by /tid/.
     -> (g -> (tid,g))   -- ^ Generate a new unique /tid/ value and update the generator state /g/.
     -> TransactionMethods (g,t (MVar x)) tid addr x
-transactionMethods (TableMethods insert delete lookup) generate = TransactionMethods
+transactionMethods methods generate = transactionMethods' id tryPutMVar methods generate
+
+-- | Like 'transactionMethods' but allows extra information to be stored in the
+-- table of pending transactions.  This also enables multiple 'Client's to
+-- share a single transaction table.
+transactionMethods' ::
+    (MVar x -> a)          -- ^ store MVar into table entry
+    -> (a -> x -> IO void) -- ^ load MVar from table entry
+    -> TableMethods t tid  -- ^ Table methods to lookup values by /tid/.
+    -> (g -> (tid,g))      -- ^ Generate a new unique /tid/ value and update the generator state /g/.
+    -> TransactionMethods (g,t a) tid addr x
+transactionMethods' store load (TableMethods insert delete lookup) generate = TransactionMethods
     { dispatchCancel = \tid (g,t) -> return (g, delete tid t)
     , dispatchRegister = \v _ (g,t) ->
         let (tid,g') = generate g
-            t' = insert tid v t
+            t' = insert tid (store v) t
         in return ( tid, (g',t') )
     , dispatchResponse = \tid x (g,t) ->
         case lookup tid t of
             Just v -> let t' = delete tid t
-                      in return ((g,t'),void $ tryPutMVar v x)
+                      in return ((g,t'),void $ load v x)
             Nothing -> return ((g,t), return ())
     }