1 files changed, 18 insertions, 15 deletions

diff --git a/dht/src/Network/QueryResponse.hs b/dht/src/Network/QueryResponse.hs
index 9c33b911..5fcd1989 100644
--- a/dht/src/Network/QueryResponse.hs
+++ b/dht/src/Network/QueryResponse.hs
@@ -217,11 +217,11 @@ asyncQuery_ (Client net d err pending whoami _) meth q addr0 withResponse = do
     now <- getPOSIXTime
     (tid,addr,expiry) <- atomically $ do
         tbl <- readTVar pending
-        ((tid,addr,expiry), tbl') <- dispatchRegister (tableMethods d)
-                (methodTimeout meth)
-                now
+        (addr,expiry) <- methodTimeout meth addr0
+        (tid, tbl') <- dispatchRegister (tableMethods d)
+                (now + microsecondsDiff expiry)
                 (withResponse . fmap (unwrapResponse meth))
-                addr0
+                addr -- XXX: Should be addr0 or addr?
                 tbl
         -- (addr,expiry) <- methodTimeout meth tid addr0
         writeTVar pending tbl'
@@ -365,7 +365,7 @@ data MethodSerializer tid addr x meth a b = MethodSerializer
     { -- | Returns the microseconds to wait for a response to this query being
       -- sent to the given address.  The /addr/ may also be modified to add
       -- routing information.
-      methodTimeout :: tid -> addr -> STM (addr,Int)
+      methodTimeout :: addr -> STM (addr,Int)
       -- | A method identifier used for error reporting.  This needn't be the
       -- same as the /meth/ argument to 'MethodHandler', but it is suggested.
     , method :: meth
@@ -385,21 +385,25 @@ data MethodSerializer tid addr x meth a b = MethodSerializer
 --
 -- The type variable /d/ is used to represent the current state of the
 -- transaction generator and the table of pending transactions.
-data TransactionMethods d tid addr x = TransactionMethods
+data TransactionMethods d qid addr x = TransactionMethods
     {
       -- | Before a query is sent, this function stores an 'MVar' to which the
-      -- response will be written too.  The returned /tid/ is a transaction id
+      -- response will be written too.  The returned /qid/ is a transaction id
       -- that can be used to forget the 'MVar' if the remote peer is not
       -- responding.
-      dispatchRegister :: (tid -> addr -> STM (addr,Int)) -> POSIXTime -> (Maybe x -> IO ()) -> addr -> d -> STM ((tid,addr,Int), d)
+      dispatchRegister :: POSIXTime -- time of expiry
+                          -> (Maybe x -> IO ()) -- callback upon response (or timeout)
+                          -> addr
+                          -> d
+                          -> STM (qid, d)
       -- | This method is invoked when an incoming packet /x/ indicates it is
-      -- a response to the transaction with id /tid/.  The returned IO action
+      -- a response to the transaction with id /qid/.  The returned IO action
       -- will write the packet to the correct 'MVar' thus completing the
       -- dispatch.
-    , dispatchResponse :: tid -> x -> d -> STM (d, IO ())
+    , dispatchResponse :: qid -> x -> d -> STM (d, IO ())
       -- | When a timeout interval elapses, this method is called to remove the
       -- transaction from the table.
-    , dispatchCancel :: tid -> d -> STM d
+    , dispatchCancel :: qid -> d -> STM d
     }
 
 -- | Construct 'TransactionMethods' methods out of 3 lookup table primitives and a
@@ -424,11 +428,10 @@ transactionMethods' ::
     -> 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 = \getTimeout now v a0 (g,t) -> do
+    , dispatchRegister = \nowPlusExpiry v a (g,t) -> do
         let (tid,g') = generate g
-        (a,expiry) <- getTimeout tid a0
-        let t' = insert tid (store v) (now + microsecondsDiff expiry) t
-        return ( (tid,a,expiry), (g',t') )
+        let t' = insert tid (store v) nowPlusExpiry t -- (now + microsecondsDiff expiry) t
+        return ( tid, (g',t') )
     , dispatchResponse = \tid x (g,t) ->
         case lookup tid t of
             Just v -> let t' = delete tid t