Merge pull request #104 from encukou/py3-query-timer

Py3: Add query timer

1 files changed, 6 insertions, 0 deletions

diff --git a/dns/query.py b/dns/query.py
index c416615..1d3954d 100644
--- a/dns/query.py
+++ b/dns/query.py
@@ -210,6 +210,7 @@ def udp(q, where, timeout=None, port=53, af=None, source=None, source_port=0,
         if source is not None:
             s.bind(source)
         _wait_for_writable(s, expiration)
+        begin_time = time.time()
         s.sendto(wire, destination)
         while 1:
             _wait_for_readable(s, expiration)
@@ -223,9 +224,11 @@ def udp(q, where, timeout=None, port=53, af=None, source=None, source_port=0,
                                        '%s instead of %s' % (from_address,
                                                              destination))
     finally:
+        response_time = time.time() - begin_time
         s.close()
     r = dns.message.from_wire(wire, keyring=q.keyring, request_mac=q.mac,
                               one_rr_per_rrset=one_rr_per_rrset)
+    r.time = response_time
     if not q.is_response(r):
         raise BadResponse
     return r
@@ -301,6 +304,7 @@ def tcp(q, where, timeout=None, port=53, af=None, source=None, source_port=0,
     try:
         expiration = _compute_expiration(timeout)
         s.setblocking(0)
+        begin_time = time.time()
         if source is not None:
             s.bind(source)
         _connect(s, destination)
@@ -316,9 +320,11 @@ def tcp(q, where, timeout=None, port=53, af=None, source=None, source_port=0,
         (l,) = struct.unpack("!H", ldata)
         wire = _net_read(s, l, expiration)
     finally:
+        response_time = time.time() - begin_time
         s.close()
     r = dns.message.from_wire(wire, keyring=q.keyring, request_mac=q.mac,
                               one_rr_per_rrset=one_rr_per_rrset)
+    r.time = response_time
     if not q.is_response(r):
         raise BadResponse
     return r