def __init__(
self,
host: str,
port: int,
protocol_flag: ProtocolFlag,
debug: bool,
timeout: int,
create_timeout: int,
):
self._debug: bool = debug
self._protocol_flag: ProtocolFlag = protocol_flag
self._create_timeout: int = create_timeout
self._connection_info = f"{protocol_flag}://{host}:{port}"
_loop = get_event_loop()
if protocol_flag == ProtocolFlag.udp:
self._connection: DatagramProtocol = DatagramProtocol(timeout=timeout)
self._connection_proxy = _loop.create_datagram_endpoint(lambda: self._connection, remote_addr=(host, port))
elif protocol_flag == ProtocolFlag.tcp:
self._connection: TcpProtocol = TcpProtocol(timeout=timeout)
self._connection_proxy = _loop.create_connection(lambda: self._connection, host=host, port=port)
else:
raise ConnectionError(f"Not support protocol:{protocol_flag}")
self.sendto = self._sendto if not debug else self._sendto_debug
def timeit(self, key: str, sample_rate: Union[int, float, None] = None) -> Iterator[None]:
"""
Context manager for easily timing methods.
"""
_loop: "asyncio.AbstractEventLoop" = get_event_loop()
started_at: float = _loop.time()
yield
value: float = _loop.time() - started_at
self.timer(key, int(value * 1000), sample_rate)
async def connect(self) -> None:
_loop = get_event_loop()
if self._protocol_flag == ProtocolFlag.udp:
connection_proxy = _loop.create_datagram_endpoint(
lambda: self._connection, remote_addr=(self._host, self._port)
)
elif self._protocol_flag == ProtocolFlag.tcp:
connection_proxy = _loop.create_connection(lambda: self._connection, host=self._host, port=self._port)
else:
raise ConnectionError(f"Not support protocol:{self._protocol_flag}")
try:
await asyncio.wait_for(connection_proxy, timeout=self._create_timeout)
except asyncio.TimeoutError as e:
raise TimeoutError(f"create connection:{self._connection_info} timeout") from e
logging.debug(f"create connection:{self._connection_info}")
