sourcerouter.py

#!/usr/bin/env python
import threading, time, datetime
import psutil
import json

from kafka import KafkaAdminClient, KafkaConsumer, KafkaProducer


class SourceRouter(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.stop_event = threading.Event()
        self.admin = KafkaAdminClient(bootstrap_servers="amp.iii-ei-stack.com:31090")
        self.producer = KafkaProducer(bootstrap_servers="amp.iii-ei-stack.com:31090")

    def stop(self):
        self.stop_event.set()

    def run(self):
        consumer = KafkaConsumer(
            bootstrap_servers="amp.iii-ei-stack.com:31090",
            auto_offset_reset="latest",
            consumer_timeout_ms=1000,
            value_deserializer=lambda m: json.loads(m),
        )
        origin_topics = [
            i
            for i in consumer.topics()
            if i.split(".")[-1] == "origin" and i.split(".")[0] == "conn"
        ]
        consumer.subscribe(origin_topics)
        while not self.stop_event.is_set():
            for message in consumer:
                payload = message.value.get("payload")
                target_topic = payload.get("topic")
                print("======target_topic======", target_topic)
                prefix = payload.get("prefix", "")
                if prefix:
                    prefix = "{0}{1}".format(prefix, "_")
                dt = datetime.datetime.fromtimestamp(float(message.timestamp) / 1000)
                dt_str = dt.strftime("%Y-%m-%dT%H:%M:%S.%f")
                value_dict = {
                    "d": {
                        payload.get("key"): {
                            prefix + str(key): val
                            for key, val in eval(payload.get("value")).items()
                        }
                    },
                    "ts": dt_str,
                }
                self.producer.send(
                    topic=target_topic, key=message.key, value=json.dumps(value_dict).encode('utf-8')
                )
                if self.stop_event.is_set():
                    break
        consumer.close()


def main():
    while True:
        print('cpu count: {}'.format(psutil.cpu_count()))
        try:
            tasks = [SourceRouter()]
            for t in tasks:
                print('cpu percent: {}'.format(psutil.cpu_percent(interval=0.3)))
                t.start()
            print("======start======", datetime.datetime.now())
            time.sleep(10)
            # Stop threads
            for task in tasks:
                task.stop()
            for task in tasks:
                task.join()
            print("======end======", datetime.datetime.now())
        except Exception as e:
            print("======error======", e)


if __name__ == "__main__":
    main()