async def send_all(sock, data): total_sent = 0 while total_sent < len(data): sent = await sock.send(data[total_sent:]) if sent == 0: raise trio.BrokenResourceError('socket connection broken') total_sent += sent
def close(self): """Close this receive channel object synchronously. All channel objects have an asynchronous `~.AsyncResource.aclose` method. Memory channels can also be closed synchronously. This has the same effect on the channel and other tasks using it, but `close` is not a trio checkpoint. This simplifies cleaning up in cancelled tasks. Using ``with receive_channel:`` will close the channel object on leaving the with block. """ if self._closed: return self._closed = True for task in self._tasks: trio.lowlevel.reschedule(task, Error(trio.ClosedResourceError())) del self._state.receive_tasks[task] self._tasks.clear() self._state.open_receive_channels -= 1 if self._state.open_receive_channels == 0: assert not self._state.receive_tasks for task in self._state.send_tasks: task.custom_sleep_data._tasks.remove(task) trio.lowlevel.reschedule(task, Error(trio.BrokenResourceError())) self._state.send_tasks.clear() self._state.data.clear()
def _translate_socket_errors_to_stream_errors(): try: yield except OSError as exc: if exc.errno in _closed_stream_errnos: raise trio.ClosedResourceError( "this socket was already closed") from None else: raise trio.BrokenResourceError( "socket connection broken: {}".format(exc)) from exc
def _translate_socket_errors_to_stream_errors(): try: yield except OSError as exc: if exc.errno in {errno.EBADF, errno.ENOTSOCK}: # EBADF on Unix, ENOTSOCK on Windows raise trio.ClosedResourceError( "this socket was already closed") from None else: raise trio.BrokenResourceError( "socket connection broken: {}".format(exc)) from exc
async def aclose(self): if self._closed: await trio.hazmat.checkpoint() return self._closed = True for task in self._tasks: trio.hazmat.reschedule(task, Error(trio.ClosedResourceError())) del self._state.receive_tasks[task] self._tasks.clear() self._state.open_receive_channels -= 1 if self._state.open_receive_channels == 0: assert not self._state.receive_tasks for task in self._state.send_tasks: task.custom_sleep_data._tasks.remove(task) trio.hazmat.reschedule(task, Error(trio.BrokenResourceError())) self._state.send_tasks.clear() self._state.data.clear() await trio.hazmat.checkpoint()
async def listen(self, table: str, *symbols: Optional[Sequence[str]]): """ Subscribe to a channel and optionally one or more specific symbols. Returns an async generator that yields messages from the subscribed channel. """ if self._websocket.closed is not None: raise trio.BrokenResourceError('Connection is closed.') listeners = [(table, )] if not symbols else [(table, symbol) for symbol in symbols] args = [] for listener in listeners: if self._subscriptions[listener] == 0: args.append(listener[0] if not symbols else ':'.join(listener)) self._subscriptions[listener] += 1 await self._send_channel.send({'op': 'subscribe', 'args': args}) async with self._pipeline.tap() as aiter: async for item, item_symbol, item_table, _ in aiter: # Lock list of listeners while sending if item_table == table and (not symbols or item_symbol in symbols): yield item log.debug('Listener detached from table: %s, symbol: %s', table, symbols) if self._websocket.closed: return args = [] for listener in listeners: self._subscriptions[listener] -= 1 if self._subscriptions[listener] == 0: log.debug( 'No more listeners on table: %s, symbol: %s. Unsubscribing.', table, symbols) args.append(listener[0] if not symbols else ':'.join(listener)) if args: await self._send_channel.send({'op': 'unsubscribe', 'args': args})
async def _broken_stream(*args, **kwargs): raise trio.BrokenResourceError()