hat.drivers.tcp

200class Server(aio.Resource):
201    """TCP or Unix Domain Socket listening server
202
203    Closing server will cancel all running `connection_cb` coroutines.
204
205    """
206
207    @property
208    def async_group(self) -> aio.Group:
209        """Async group"""
210        return self._async_group
211
212    @property
213    def info(self) -> ServerInfo:
214        """Server info"""
215        return self._info
216
217    async def _on_close(self):
218        self._srv.close()
219
220        if self._bind_connections or sys.version_info[:2] < (3, 12):
221            await self._srv.wait_closed()
222
223    async def _on_connection(self, protocol):
224        self._log.debug('new incomming connection')
225
226        conn = _create_connection(protocol)
227
228        try:
229            await aio.call(self._connection_cb, conn)
230
231            if self._bind_connections:
232                await conn.wait_closing()
233
234            else:
235                conn = None
236
237        except Exception as e:
238            self._log.warning('connection callback error: %s', e, exc_info=e)
239
240        finally:
241            if conn:
242                await aio.uncancellable(conn.async_close())

257class Connection(aio.Resource):
258    """TCP or Unix Domain Socket connection"""
259
260    @property
261    def async_group(self) -> aio.Group:
262        """Async group"""
263        return self._async_group
264
265    @property
266    def info(self) -> ConnectionInfo:
267        """Connection info"""
268        return self._protocol.info
269
270    @property
271    def ssl_object(self) -> ssl.SSLObject | ssl.SSLSocket | None:
272        """SSL Object"""
273        return self._protocol.ssl_object
274
275    async def write(self, data: util.Bytes):
276        """Write data
277
278        This coroutine will wait until `data` can be added to output buffer.
279
280        """
281        if not self.is_open:
282            raise ConnectionError()
283
284        await self._protocol.write(data)
285
286    async def drain(self):
287        """Drain output buffer"""
288        await self._protocol.drain()
289
290    async def read(self, n: int = -1) -> util.Bytes:
291        """Read up to `n` bytes
292
293        If EOF is detected and no new bytes are available, `ConnectionError`
294        is raised.
295
296        """
297        return await self._protocol.read(n)
298
299    async def readexactly(self, n: int) -> util.Bytes:
300        """Read exactly `n` bytes
301
302        If exact number of bytes could not be read, `ConnectionError` is
303        raised.
304
305        """
306        return await self._protocol.readexactly(n)
307
308    def clear_input_buffer(self) -> int:
309        """Clear input buffer
310
311        Returns number of bytes cleared from buffer.
312
313        """
314        return self._protocol.clear_input_buffer()

 96async def connect(addr: StreamAddress,
 97                  *,
 98                  name: str | None = None,
 99                  input_buffer_limit: int = 64 * 1024,
100                  **kwargs
101                  ) -> 'Connection':
102    """Create TCP or Unix Domain Socket connection
103
104    Argument `addr` specifies remote server listening address.
105
106    Argument `name` defines connection name available in property `info`.
107
108    Argument `input_buffer_limit` defines number of bytes in input buffer
109    that whill temporary pause data receiving. Once number of bytes
110    drops bellow `input_buffer_limit`, data receiving is resumed. If this
111    argument is ``0``, data receive pausing is disabled.
112
113    Additional arguments are passed directly to `asyncio.create_connection` or
114    `asyncio.create_unix_connection`.
115
116    """
117    loop = asyncio.get_running_loop()
118    stream_type = _get_stream_type(addr)
119    create_protocol = functools.partial(_Protocol, stream_type, None, name,
120                                        input_buffer_limit)
121
122    if stream_type == StreamType.TCP:
123        _, protocol = await loop.create_connection(create_protocol,
124                                                   addr.host, addr.port,
125                                                   **kwargs)
126
127    elif stream_type == StreamType.UNIX:
128        _, protocol = await loop.create_unix_connection(create_protocol, addr,
129                                                        **kwargs)
130
131    else:
132        raise ValueError('unsupported stream type')
133
134    return _create_connection(protocol)

137async def listen(connection_cb: ConnectionCb,
138                 addr: StreamAddress,
139                 *,
140                 name: str | None = None,
141                 bind_connections: bool = False,
142                 input_buffer_limit: int = 64 * 1024,
143                 **kwargs
144                 ) -> 'Server':
145    """Create listening TCP or Unix Domain Socket server
146
147    Argument `name` defines server name available in property `info`. This
148    name is used for all incomming connections.
149
150    If `bind_connections` is ``True``, closing server will close all open
151    incoming connections.
152
153    Argument `input_buffer_limit` is associated with newly created connections
154    (see `connect`).
155
156    Additional arguments are passed directly to `asyncio.create_server` or
157    `asyncio.create_unix_server`.
158
159    """
160    server = Server()
161    server._connection_cb = connection_cb
162    server._bind_connections = bind_connections
163    server._stream_type = _get_stream_type(addr)
164    server._async_group = aio.Group()
165    server._log = mlog
166
167    loop = asyncio.get_running_loop()
168    on_connection = functools.partial(server.async_group.spawn,
169                                      server._on_connection)
170    create_protocol = functools.partial(_Protocol, server._stream_type,
171                                        on_connection, name,
172                                        input_buffer_limit)
173
174    if server._stream_type == StreamType.TCP:
175        server._srv = await loop.create_server(create_protocol,
176                                               addr.host, addr.port, **kwargs)
177
178    elif server._stream_type == StreamType.UNIX:
179        server._srv = await loop.create_unix_server(create_protocol,
180                                                    addr, **kwargs)
181
182    else:
183        raise ValueError('unsupported stream type')
184
185    server.async_group.spawn(aio.call_on_cancel, server._on_close)
186
187    try:
188        server._info = _get_server_info(server._stream_type, name, server._srv)
189        server._log = _create_server_logger(server._info)
190
191    except Exception:
192        await aio.uncancellable(server.async_close())
193        raise
194
195    server._log.debug('listening for incomming connections')
196
197    return server