Make aio server client_timeout an idle timeout

tests/test_aio.py

@@ -375,6 +375,59 @@ class TestAIOFramedCompactSSL(SSLServerMixin, TestAIOFramedCompact):
     pass
 
 
+@pytest.mark.parametrize("use_unix_socket", [False, True])
+@pytest.mark.asyncio
+async def test_server_client_timeout_is_idle_timeout(use_unix_socket):
+    """client_timeout bounds idle time between reads, not connection age."""
+    if use_unix_socket:
+        server_kwargs = {"unix_socket": "/tmp/aio_thriftpy_idle_timeout.sock"}
+    else:
+        server_kwargs = {"host": "127.0.0.1", "port": 0}
+    server = make_aio_server(
+        addressbook.AddressBookService,
+        Dispatcher(),
+        client_timeout=500,
+        loop=asyncio.new_event_loop(),
+        **server_kwargs,
+    )
+    _start_in_thread(server)
+    try:
+        if use_unix_socket:
+            client_kwargs = server_kwargs
+        else:
+            port = server.trans.raw_sock.getsockname()[1]
+            client_kwargs = {"host": "127.0.0.1", "port": port}
+
+        c = await make_aio_client(addressbook.AddressBookService,
+                                  **client_kwargs)
+        try:
+            # a long lived connection that keeps making requests stays open
+            for _ in range(5):
+                await asyncio.sleep(0.1)
+                await c.ping()
+            # a request whose handler runs longer than client_timeout
+            # is not cut short
+            assert await c.sleep(1000) is True
+        finally:
+            c.close()
+
+        # an idle connection is closed by the server after client_timeout,
+        # well before the 3s default that would apply if the value were
+        # dropped on the way to the server socket
+        if use_unix_socket:
+            reader, writer = await asyncio.open_unix_connection(
+                server_kwargs["unix_socket"])
+        else:
+            reader, writer = await asyncio.open_connection(
+                "127.0.0.1", port)
+        try:
+            assert await asyncio.wait_for(reader.read(1), 2) == b""
+        finally:
+            writer.close()
+    finally:
+        server.loop.call_soon_threadsafe(server.loop.stop)
+
+
 @pytest.mark.asyncio
 async def test_client_connect_timeout():
     with pytest.raises(TTransportException):

thriftpy2/contrib/aio/rpc.py

@@ -81,7 +81,8 @@ def make_server(
     processor = TAsyncProcessor(service, handler)
 
     if unix_socket:
-        server_socket = TAsyncServerSocket(unix_socket=unix_socket)
+        server_socket = TAsyncServerSocket(
+            unix_socket=unix_socket, client_timeout=client_timeout)
         if certfile:
             warnings.warn("SSL only works with host:port, not unix_socket.")
     elif host and port is not None:

thriftpy2/contrib/aio/socket.py

@@ -233,7 +233,9 @@ def __init__(self, host=None, port=None, unix_socket=None,
         @param unix_socket(str) The filename of a unix socket to connect to
         @param socket_family(str) socket.AF_INET or socket.AF_INET6. only
             take effect when using host/port
-        @param client_timeout   client socket timeout
+        @param client_timeout   idle timeout in ms for each client
+            connection, applied to every read and flush like the
+            sync server does with settimeout. None disables it.
         @param backlog          backlog for server socket
         @param certfile(str)        The server cert pem filename
         @param keyfile(str)         The server cert key filename
@@ -314,11 +316,9 @@ def _create_client_connected_cb(self, callback):
 
         async def client_connected_cb(reader, writer):
             try:
-                await asyncio.wait_for(
-                    callback(StreamHandler(reader, writer)),
-                    self.client_timeout
-                )
-            except asyncio.exceptions.TimeoutError:
+                await callback(
+                    StreamHandler(reader, writer, self.client_timeout))
+            finally:
                 writer.close()
 
         return client_connected_cb
@@ -335,12 +335,20 @@ def close(self):
 
 
 class StreamHandler(TAsyncTransportBase):
-    def __init__(self, reader, writer):
+    def __init__(self, reader, writer, timeout=None):
+        """
+        @param timeout  idle timeout in seconds for read and flush, None
+            means wait forever.
+        """
         self.reader, self.writer = reader, writer
+        self.timeout = timeout
 
     async def read(self, sz):
         try:
-            buff = await self.reader.read(sz)
+            buff = await asyncio.wait_for(self.reader.read(sz), self.timeout)
+        except asyncio.TimeoutError:
+            raise TTransportException(type=TTransportException.TIMED_OUT,
+                                      message='TSocket read timed out')
         except socket.error as e:
             if e.errno == errno.ECONNRESET and MAC_OR_BSD:
                 # freebsd and Mach don't follow POSIX semantic of recv
@@ -362,7 +370,14 @@ def write(self, buf):
         self.writer.write(buf)
 
     async def flush(self):
-        await self.writer.drain()
+        try:
+            await asyncio.wait_for(self.writer.drain(), self.timeout)
+        except asyncio.TimeoutError:
+            # close() would keep the fd open until the peer drains the
+            # buffered data, drop the connection right away instead
+            self.writer.transport.abort()
+            raise TTransportException(type=TTransportException.TIMED_OUT,
+                                      message='TSocket write timed out')
 
     def close(self):
         try: