mirror of
https://github.com/sprockets/sprockets-statsd.git
synced 2024-11-28 11:19:52 +00:00
93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
import asyncio
|
|
import logging
|
|
import typing
|
|
|
|
|
|
class Processor(asyncio.Protocol):
|
|
def __init__(self, *, host, port: int = 8125, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.host = host
|
|
self.port = port
|
|
|
|
self.closed = asyncio.Event()
|
|
self.connected = asyncio.Event()
|
|
self.logger = logging.getLogger(__package__).getChild('Processor')
|
|
self.running = False
|
|
self.transport = None
|
|
|
|
self._queue = asyncio.Queue()
|
|
|
|
async def run(self):
|
|
self.running = True
|
|
while self.running:
|
|
try:
|
|
await self._connect_if_necessary()
|
|
await self._process_metric()
|
|
except asyncio.CancelledError:
|
|
self.logger.info('task cancelled, exiting')
|
|
break
|
|
|
|
self.running = False
|
|
self.logger.info('loop finished with %d metrics in the queue',
|
|
self._queue.qsize())
|
|
if self.connected.is_set():
|
|
num_ready = self._queue.qsize()
|
|
self.logger.info('draining %d metrics', num_ready)
|
|
for _ in range(num_ready):
|
|
await self._process_metric()
|
|
self.logger.debug('closing transport')
|
|
self.transport.close()
|
|
|
|
while self.connected.is_set():
|
|
self.logger.debug('waiting on transport to close')
|
|
await asyncio.sleep(0.1)
|
|
|
|
self.logger.info('processor is exiting')
|
|
self.closed.set()
|
|
|
|
async def stop(self):
|
|
self.running = False
|
|
await self.closed.wait()
|
|
|
|
def inject_metric(self, path: str, value: typing.Union[float, int, str],
|
|
type_code: str):
|
|
payload = f'{path}:{value}|{type_code}\n'
|
|
self._queue.put_nowait(payload.encode('utf-8'))
|
|
|
|
def eof_received(self):
|
|
self.logger.warning('received EOF from statsd server')
|
|
self.connected.clear()
|
|
|
|
def connection_made(self, transport: asyncio.Transport):
|
|
server, port = transport.get_extra_info('peername')
|
|
self.logger.info('connected to statsd %s:%s', server, port)
|
|
self.transport = transport
|
|
self.connected.set()
|
|
|
|
def connection_lost(self, exc: typing.Optional[Exception]):
|
|
self.logger.warning('statsd server connection lost')
|
|
self.connected.clear()
|
|
|
|
async def _connect_if_necessary(self, wait_time: float = 0.1):
|
|
try:
|
|
await asyncio.wait_for(self.connected.wait(), wait_time)
|
|
except asyncio.TimeoutError:
|
|
try:
|
|
self.logger.debug('starting connection to %s:%s', self.host,
|
|
self.port)
|
|
await asyncio.get_running_loop().create_connection(
|
|
protocol_factory=lambda: self,
|
|
host=self.host,
|
|
port=self.port)
|
|
except IOError as error:
|
|
self.logger.warning('connection to %s:%s failed: %s',
|
|
self.host, self.port, error)
|
|
|
|
async def _process_metric(self):
|
|
try:
|
|
metric = await asyncio.wait_for(self._queue.get(), 0.1)
|
|
except asyncio.TimeoutError:
|
|
return # nothing to do
|
|
else:
|
|
self.transport.write(metric)
|
|
self._queue.task_done()
|