async def recv_paint(paint: RedPaint, consumer: aiokafka.AIOKafkaConsumer, wait_notice): last_time = time.monotonic() wait_ct = ( 1 # if count gets too high, stop generating new events. Something is broken. ) while True: if time.monotonic() - last_time > wait_notice * wait_ct: if paint.is_sending(): log.info(f"No valid message in {wait_notice}*{wait_ct} seconds") wait_ct += 1 if wait_ct > 10: await paint.pause_sending() try: msg = await asyncio.wait_for(consumer.getone(), 1) if paint.validate((msg.key, msg.value)): last_time = time.monotonic() wait_ct = 1 paint.start_sending() except asyncio.TimeoutError as e: # We are expecting this when there is no big flow of messages through the system pass
class AIOKafkaRPCClient(object): log = logging.getLogger(__name__) def __init__(self, kafka_servers='localhost:9092', in_topic='aiokafkarpc_in', out_topic='aiokafkarpc_out', out_partitions=(0, ), translation_table=[], *, loop): self.call = CallObj(self._call_wrapper) self._topic_in = in_topic self._loop = loop self._waiters = {} self._out_topic = out_topic self._out_partitions = out_partitions default, ext_hook = get_msgpack_hooks(translation_table) self.__consumer = AIOKafkaConsumer( self._out_topic, loop=loop, bootstrap_servers=kafka_servers, group_id=None, key_deserializer=lambda x: x.decode("utf-8"), value_deserializer=lambda x: msgpack.unpackb( x, ext_hook=ext_hook, encoding="utf-8")) self.__producer = AIOKafkaProducer( bootstrap_servers=kafka_servers, loop=loop, key_serializer=lambda x: x.encode("utf-8"), value_serializer=lambda x: msgpack.packb(x, default=default)) @asyncio.coroutine def run(self): yield from self.__producer.start() yield from self.__consumer.start() # FIXME manual partition assignment does not work correctly in aiokafka # self.__consumer.assign( # [TopicPartition(self._out_topic, p) for p in self._out_partitions]) # # ensure that topic partitions exists for tp in self.__consumer.assignment(): yield from self.__consumer.position(tp) self._consume_task = self._loop.create_task(self.__consume_routine()) @asyncio.coroutine def close(self, timeout=10): yield from self.__producer.stop() if self._waiters: yield from asyncio.wait(self._waiters.values(), loop=self._loop, timeout=timeout) self._consume_task.cancel() try: yield from self._consume_task except asyncio.CancelledError: pass yield from self.__consumer.stop() for fut in self._waiters.values(): fut.set_exception(asyncio.TimeoutError()) def _call_wrapper(self, method): @asyncio.coroutine def rpc_call(*args, **kw_args): call_id = uuid.uuid4().hex ptid = random.choice(self._out_partitions) request = (method, args, kw_args, ptid) fut = asyncio.Future(loop=self._loop) fut.add_done_callback(lambda fut: self._waiters.pop(call_id)) self._waiters[call_id] = fut try: yield from self.__producer.send(self._topic_in, request, key=call_id) except Exception as err: self.log.error("send RPC request failed: %s", err) self._waiters[call_id].set_exception(err) return (yield from self._waiters[call_id]) return rpc_call @asyncio.coroutine def __consume_routine(self): while True: message = yield from self.__consumer.getone() call_id = message.key response = message.value self.call = CallObj(self._call_wrapper) fut = self._waiters.get(call_id) if fut is None: continue if "error" in response: self.log.debug(response.get("stacktrace")) fut.set_exception(RPCError(response["error"])) else: fut.set_result(response["result"])
class AIOKafkaRPC(object): log = logging.getLogger(__name__) def __init__(self, rpc_obj, kafka_servers='localhost:9092', in_topic='aiokafkarpc_in', out_topic='aiokafkarpc_out', translation_table=[], *, loop): self._tasks = {} self._loop = loop self._topic_out = out_topic self._rpc_obj = rpc_obj self._res_queue = asyncio.Queue(loop=loop) default, ext_hook = get_msgpack_hooks(translation_table) self.__consumer = AIOKafkaConsumer( in_topic, loop=loop, bootstrap_servers=kafka_servers, group_id=in_topic + '-group', key_deserializer=lambda x: x.decode("utf-8"), value_deserializer=lambda x: msgpack.unpackb( x, ext_hook=ext_hook, encoding="utf-8")) self.__producer = AIOKafkaProducer( bootstrap_servers=kafka_servers, loop=loop, key_serializer=lambda x: x.encode("utf-8"), value_serializer=lambda x: msgpack.packb(x, default=default)) @asyncio.coroutine def run(self): yield from self.__producer.start() yield from self.__consumer.start() self._consume_task = self._loop.create_task(self.__consume_routine()) self._produce_task = self._loop.create_task(self.__produce_routine()) @asyncio.coroutine def close(self, timeout=10): self._consume_task.cancel() try: yield from self._consume_task except asyncio.CancelledError: pass if self._tasks: yield from asyncio.wait(self._tasks.values(), loop=self._loop, timeout=timeout) self._res_queue.put_nowait((None, None, None)) yield from self._produce_task yield from self.__producer.stop() yield from self.__consumer.stop() def __send_result(self, call_id, ptid, result=None, err=None): if err is not None: out = io.StringIO() traceback.print_exc(file=out) ret = {"error": repr(err), "stacktrace": out.getvalue()} else: ret = {"result": result} self._res_queue.put_nowait((call_id, ptid, ret)) def __on_done_task(self, call_id, ptid, task): self._tasks.pop(call_id) try: result = task.result() self.__send_result(call_id, ptid, result) except Exception as err: self.__send_result(call_id, ptid, err=err) @asyncio.coroutine def __produce_routine(self): while True: call_id, ptid, res = yield from self._res_queue.get() if call_id is None: break try: yield from self.__producer.send(self._topic_out, res, key=call_id, partition=ptid) except KafkaError as err: self.log.error("send RPC response failed: %s", err) except Exception as err: self.__send_result(call_id, err=err) @asyncio.coroutine def __consume_routine(self): while True: message = yield from self.__consumer.getone() call_id = message.key method_name, args, kw_args, ptid = message.value try: method = getattr(self._rpc_obj, method_name) if not asyncio.iscoroutinefunction(method): raise RuntimeError( "'{}' should be a coroutine".format(method_name)) task = self._loop.create_task(method(*args, **kw_args)) task.add_done_callback( functools.partial(self.__on_done_task, call_id, ptid)) self._tasks[call_id] = task except Exception as err: self.__send_result(call_id, ptid, err=err)