hat.drivers.tcp

  1"""Asyncio TCP wrapper"""
  2
  3import asyncio
  4import collections
  5import functools
  6import logging
  7import sys
  8import typing
  9
 10from hat import aio
 11from hat import util
 12
 13from hat.drivers import ssl
 14
 15
 16mlog: logging.Logger = logging.getLogger(__name__)
 17"""Module logger"""
 18
 19
 20class Address(typing.NamedTuple):
 21    host: str
 22    port: int
 23
 24
 25class ConnectionInfo(typing.NamedTuple):
 26    name: str | None
 27    local_addr: Address
 28    remote_addr: Address
 29
 30
 31class ServerInfo(typing.NamedTuple):
 32    name: str | None
 33    addresses: list[Address]
 34
 35
 36ConnectionCb: typing.TypeAlias = aio.AsyncCallable[['Connection'], None]
 37"""Connection callback"""
 38
 39
 40async def connect(addr: Address,
 41                  *,
 42                  name: str | None = None,
 43                  input_buffer_limit: int = 64 * 1024,
 44                  **kwargs
 45                  ) -> 'Connection':
 46    """Create TCP connection
 47
 48    Argument `addr` specifies remote server listening address.
 49
 50    Argument `name` defines connection name available in property `info`.
 51
 52    Argument `input_buffer_limit` defines number of bytes in input buffer
 53    that whill temporary pause data receiving. Once number of bytes
 54    drops bellow `input_buffer_limit`, data receiving is resumed. If this
 55    argument is ``0``, data receive pausing is disabled.
 56
 57    Additional arguments are passed directly to `asyncio.create_connection`.
 58
 59    """
 60    loop = asyncio.get_running_loop()
 61    create_protocol = functools.partial(Protocol, None, name,
 62                                        input_buffer_limit)
 63    _, protocol = await loop.create_connection(create_protocol,
 64                                               addr.host, addr.port,
 65                                               **kwargs)
 66    return Connection(protocol)
 67
 68
 69async def listen(connection_cb: ConnectionCb,
 70                 addr: Address,
 71                 *,
 72                 name: str | None = None,
 73                 bind_connections: bool = False,
 74                 input_buffer_limit: int = 64 * 1024,
 75                 **kwargs
 76                 ) -> 'Server':
 77    """Create listening server
 78
 79    Argument `name` defines server name available in property `info`. This
 80    name is used for all incomming connections.
 81
 82    If `bind_connections` is ``True``, closing server will close all open
 83    incoming connections.
 84
 85    Argument `input_buffer_limit` is associated with newly created connections
 86    (see `connect`).
 87
 88    Additional arguments are passed directly to `asyncio.create_server`.
 89
 90    """
 91    server = Server()
 92    server._connection_cb = connection_cb
 93    server._bind_connections = bind_connections
 94    server._async_group = aio.Group()
 95    server._log = mlog
 96
 97    on_connection = functools.partial(server.async_group.spawn,
 98                                      server._on_connection)
 99    create_protocol = functools.partial(Protocol, on_connection, name,
100                                        input_buffer_limit)
101
102    loop = asyncio.get_running_loop()
103    server._srv = await loop.create_server(create_protocol, addr.host,
104                                           addr.port, **kwargs)
105
106    server.async_group.spawn(aio.call_on_cancel, server._on_close)
107
108    try:
109        socknames = (socket.getsockname() for socket in server._srv.sockets)
110        addresses = [Address(*sockname[:2]) for sockname in socknames]
111        server._info = ServerInfo(name=name,
112                                  addresses=addresses)
113        server._log = _create_server_logger_adapter(server._info)
114
115    except Exception:
116        await aio.uncancellable(server.async_close())
117        raise
118
119    server._log.debug('listening for incomming connections')
120
121    return server
122
123
124class Server(aio.Resource):
125    """TCP listening server
126
127    Closing server will cancel all running `connection_cb` coroutines.
128
129    """
130
131    @property
132    def async_group(self) -> aio.Group:
133        """Async group"""
134        return self._async_group
135
136    @property
137    def info(self) -> ServerInfo:
138        """Server info"""
139        return self._info
140
141    async def _on_close(self):
142        self._srv.close()
143
144        if self._bind_connections or sys.version_info[:2] < (3, 12):
145            await self._srv.wait_closed()
146
147    async def _on_connection(self, protocol):
148        self._log.debug('new incomming connection')
149
150        conn = Connection(protocol)
151
152        try:
153            await aio.call(self._connection_cb, conn)
154
155            if self._bind_connections:
156                await conn.wait_closing()
157
158            else:
159                conn = None
160
161        except Exception as e:
162            self._log.warning('connection callback error: %s', e, exc_info=e)
163
164        finally:
165            if conn:
166                await aio.uncancellable(conn.async_close())
167
168
169class Connection(aio.Resource):
170    """TCP connection"""
171
172    def __init__(self, protocol: 'Protocol'):
173        self._protocol = protocol
174        self._async_group = aio.Group()
175        self._log = _create_connection_logger_adapter(protocol.info)
176
177        self._log.debug('connection established')
178
179        self.async_group.spawn(aio.call_on_cancel, protocol.async_close)
180        self.async_group.spawn(aio.call_on_done, protocol.wait_closed(),
181                               self.close)
182
183    @property
184    def async_group(self) -> aio.Group:
185        """Async group"""
186        return self._async_group
187
188    @property
189    def info(self) -> ConnectionInfo:
190        """Connection info"""
191        return self._protocol.info
192
193    @property
194    def ssl_object(self) -> ssl.SSLObject | ssl.SSLSocket | None:
195        """SSL Object"""
196        return self._protocol.ssl_object
197
198    async def write(self, data: util.Bytes):
199        """Write data
200
201        This coroutine will wait until `data` can be added to output buffer.
202
203        """
204        if not self.is_open:
205            raise ConnectionError()
206
207        await self._protocol.write(data)
208
209    async def drain(self):
210        """Drain output buffer"""
211        await self._protocol.drain()
212
213    async def read(self, n: int = -1) -> util.Bytes:
214        """Read up to `n` bytes
215
216        If EOF is detected and no new bytes are available, `ConnectionError`
217        is raised.
218
219        """
220        return await self._protocol.read(n)
221
222    async def readexactly(self, n: int) -> util.Bytes:
223        """Read exactly `n` bytes
224
225        If exact number of bytes could not be read, `ConnectionError` is
226        raised.
227
228        """
229        return await self._protocol.readexactly(n)
230
231    def clear_input_buffer(self) -> int:
232        """Clear input buffer
233
234        Returns number of bytes cleared from buffer.
235
236        """
237        return self._protocol.clear_input_buffer()
238
239
240class Protocol(asyncio.Protocol):
241    """Asyncio protocol implementation"""
242
243    def __init__(self,
244                 on_connected: typing.Callable[['Protocol'], None] | None,
245                 name: str | None,
246                 input_buffer_limit: int):
247        self._on_connected = on_connected
248        self._name = name
249        self._input_buffer_limit = input_buffer_limit
250        self._loop = asyncio.get_running_loop()
251        self._input_buffer = util.BytesBuffer()
252        self._transport = None
253        self._read_queue = None
254        self._write_queue = None
255        self._drain_futures = None
256        self._closed_futures = None
257        self._info = None
258        self._ssl_object = None
259        self._log = mlog
260
261    @property
262    def info(self) -> ConnectionInfo:
263        return self._info
264
265    @property
266    def ssl_object(self) -> ssl.SSLObject | ssl.SSLSocket | None:
267        return self._ssl_object
268
269    def connection_made(self, transport: asyncio.Transport):
270        self._transport = transport
271        self._read_queue = collections.deque()
272        self._closed_futures = collections.deque()
273
274        try:
275            sockname = transport.get_extra_info('sockname')
276            peername = transport.get_extra_info('peername')
277            self._info = ConnectionInfo(
278                name=self._name,
279                local_addr=Address(sockname[0], sockname[1]),
280                remote_addr=Address(peername[0], peername[1]))
281
282            self._log = _create_connection_logger_adapter(self._info)
283
284            self._ssl_object = transport.get_extra_info('ssl_object')
285
286            if self._on_connected:
287                self._on_connected(self)
288
289        except Exception:
290            transport.abort()
291            return
292
293    def connection_lost(self, exc: Exception | None):
294        self._transport = None
295        write_queue, self._write_queue = self._write_queue, None
296        drain_futures, self._drain_futures = self._drain_futures, None
297        closed_futures, self._closed_futures = self._closed_futures, None
298
299        self.eof_received()
300
301        while write_queue:
302            _, future = write_queue.popleft()
303            if not future.done():
304                future.set_exception(ConnectionError())
305
306        while drain_futures:
307            future = drain_futures.popleft()
308            if not future.done():
309                future.set_result(None)
310
311        while closed_futures:
312            future = closed_futures.popleft()
313            if not future.done():
314                future.set_result(None)
315
316        self._log.debug('connection closed')
317
318    def pause_writing(self):
319        self._log.debug('pause writing')
320
321        self._write_queue = collections.deque()
322        self._drain_futures = collections.deque()
323
324    def resume_writing(self):
325        self._log.debug('resume writing')
326
327        write_queue, self._write_queue = self._write_queue, None
328        drain_futures, self._drain_futures = self._drain_futures, None
329
330        while self._write_queue is None and write_queue:
331            data, future = write_queue.popleft()
332            if future.done():
333                continue
334
335            self._log.debug('writing %s bytes', len(data))
336            self._transport.write(data)
337            future.set_result(None)
338
339        if write_queue:
340            write_queue.extend(self._write_queue)
341            self._write_queue = write_queue
342
343            drain_futures.extend(self._drain_futures)
344            self._drain_futures = drain_futures
345
346            return
347
348        while drain_futures:
349            future = drain_futures.popleft()
350            if not future.done():
351                future.set_result(None)
352
353    def data_received(self, data: util.Bytes):
354        self._log.debug('received %s bytes', len(data))
355
356        self._input_buffer.add(data)
357        self._process_input_buffer()
358
359    def eof_received(self):
360        self._log.debug('eof received')
361
362        while self._read_queue:
363            exact, n, future = self._read_queue.popleft()
364            if future.done():
365                continue
366
367            if exact and n <= len(self._input_buffer):
368                future.set_result(self._input_buffer.read(n))
369
370            elif not exact and self._input_buffer:
371                future.set_result(self._input_buffer.read(n))
372
373            else:
374                future.set_exception(ConnectionError())
375
376        self._read_queue = None
377
378    async def write(self, data: util.Bytes):
379        if self._transport is None:
380            raise ConnectionError()
381
382        if self._write_queue is None:
383            self._log.debug('writing %s bytes', len(data))
384            self._transport.write(data)
385            return
386
387        future = self._loop.create_future()
388        self._write_queue.append((data, future))
389        await future
390
391    async def drain(self):
392        if self._drain_futures is None:
393            return
394
395        future = self._loop.create_future()
396        self._drain_futures.append(future)
397        await future
398
399    async def read(self, n: int) -> util.Bytes:
400        if n == 0:
401            return b''
402
403        if self._input_buffer and not self._read_queue:
404            data = self._input_buffer.read(n)
405            self._process_input_buffer()
406            return data
407
408        if self._read_queue is None:
409            raise ConnectionError()
410
411        future = self._loop.create_future()
412        future.add_done_callback(self._on_read_future_done)
413        self._read_queue.append((False, n, future))
414        return await future
415
416    async def readexactly(self, n: int) -> util.Bytes:
417        if n == 0:
418            return b''
419
420        if n <= len(self._input_buffer) and not self._read_queue:
421            data = self._input_buffer.read(n)
422            self._process_input_buffer()
423            return data
424
425        if self._read_queue is None:
426            raise ConnectionError()
427
428        future = self._loop.create_future()
429        future.add_done_callback(self._on_read_future_done)
430        self._read_queue.append((True, n, future))
431        self._process_input_buffer()
432        return await future
433
434    def clear_input_buffer(self) -> int:
435        count = self._input_buffer.clear()
436        self._transport.resume_reading()
437        return count
438
439    async def async_close(self):
440        if self._transport is not None:
441            self._transport.close()
442
443        await self.wait_closed()
444
445    async def wait_closed(self):
446        if self._closed_futures is None:
447            return
448
449        future = self._loop.create_future()
450        self._closed_futures.append(future)
451        await future
452
453    def _on_read_future_done(self, future):
454        if not self._read_queue:
455            return
456
457        if not future.cancelled():
458            return
459
460        for _ in range(len(self._read_queue)):
461            i = self._read_queue.popleft()
462            if not i[2].done():
463                self._read_queue.append(i)
464
465        self._process_input_buffer()
466
467    def _process_input_buffer(self):
468        while self._input_buffer and self._read_queue:
469            exact, n, future = self._read_queue.popleft()
470            if future.done():
471                continue
472
473            if not exact:
474                future.set_result(self._input_buffer.read(n))
475
476            elif n <= len(self._input_buffer):
477                future.set_result(self._input_buffer.read(n))
478
479            else:
480                self._read_queue.appendleft((exact, n, future))
481                break
482
483        if not self._transport:
484            return
485
486        pause = (self._input_buffer_limit > 0 and
487                 len(self._input_buffer) > self._input_buffer_limit and
488                 not self._read_queue)
489
490        if pause:
491            self._transport.pause_reading()
492
493        else:
494            self._transport.resume_reading()
495
496
497def _create_server_logger_adapter(info):
498    extra = {'meta': {'type': 'TcpServer',
499                      'name': info.name,
500                      'addresses': [{'host': addr.host,
501                                     'port': addr.port}
502                                    for addr in info.addresses]}}
503
504    return logging.LoggerAdapter(mlog, extra)
505
506
507def _create_connection_logger_adapter(info):
508    extra = {'meta': {'type': 'TcpConnection',
509                      'name': info.name,
510                      'local_addr': {'host': info.local_addr.host,
511                                     'port': info.local_addr.port},
512                      'remote_addr': {'host': info.remote_addr.host,
513                                      'port': info.remote_addr.port}}}
514
515    return logging.LoggerAdapter(mlog, extra)