main.py

from concurrent.futures import ThreadPoolExecutor
from src.service import RIS, RouteViews
from src.utils.config import Config
from src.utils.logger import Logger
from src.libs.bmp import BMPv3
from src.seeder import Seeder
from src.sender import Sender
from bs4 import BeautifulSoup
from datetime import datetime
import queue as queueio
import rocksdbpy
import threading
import requests
import asyncio
import signal
import struct
import time
import os

## Constants ##

logger = Logger()

LOG_INTERVAL = 10

## Classes ##

class Relay():
    def __init__(self):
        self.queue = queueio.Queue(maxsize=10000000)
        self.executor = ThreadPoolExecutor(max_workers=4)
        self.db = rocksdbpy.open_default("/var/lib/relay/" + Config.RELAY_NAME)
        self.loop = asyncio.get_event_loop()
        self.events = {
            'seeded': threading.Event(),
            'shutdown': threading.Event(),
        }
        self.memory = {
            'bytes_sent': 0,
            'bytes_received': 0,
        }

    def is_sane(self) -> bool:
        if self.db.get(b'started') == b'\x01':
            if not self.db.get(b'seeded') == b'\x01':
                # Database in inconsistent state.
                # A reset is required.
                return False
        return True

    def _logger(self):
        while True:
            # Queue size.
            queue_size = self.queue.qsize()

            # Check for shutdown event.
            if self.events['shutdown'].is_set():
                break

            # Compute kbit/s.
            kbps_sent = (self.memory['bytes_sent'] * 8) / LOG_INTERVAL / 1000
            kbps_received = (self.memory['bytes_received'] * 8) / LOG_INTERVAL / 1000
            
            logger.info(
                f"receive={kbps_received:.2f} kbps send={kbps_sent:.2f} kbps queued={queue_size}"
            )
        
            # Reset trackers.
            self.memory['bytes_sent'] = 0
            self.memory['bytes_received'] = 0

            # Sleep for a bit.
            time.sleep(LOG_INTERVAL)

    def stop(self):
        self.events['shutdown'].set()
        self.executor.shutdown(wait=False, cancel_futures=True)
        time.sleep(10)
        os._exit(1)

    def start(self):
        try:
            logger.info("Starting")

            if not self.is_sane():
                raise RuntimeError("Database not sane")

            self.loop.add_signal_handler(signal.SIGTERM, self.stop)
            self.loop.add_signal_handler(signal.SIGINT, self.stop)

            # Logger.
            self.loop.run_in_executor(self.executor, self._logger)

            # Start our tasks.
            match Config.RELAY_TYPE:
                case 'ris':
                    # BMPv3 init.
                    message = BMPv3.init_message(
                        router_name=Config.RELAY_NAME,
                        router_descr=f"{Config.RELAY_NAME}.ripe.net"
                    )
                    self.queue.put((message, 0, None, -1, None, -1, -1, False))

                    # Seed the collector.
                    # This is a one-off operation.
                    if self.db.get(b'seeded') != b'\x01':
                        index = f"{Config.RELAY_INDEX}/{datetime.now().year}.{datetime.now().month:02d}"
                        filename = f"bview.{datetime.now().year}{datetime.now().month:02d}{datetime.now().day:02d}.0000.gz"
                        self.loop.run_in_executor(self.executor,
                            lambda: Seeder(
                                queue=self.queue,
                                url=f"{index}/{filename}",
                                events=self.events,
                                memory=self.memory,
                                db=self.db,
                            ).start())
                    else:
                        self.events['seeded'].set()

                    # RIPE RIS (Routing Information Service)
                    # We poll the archive for new updates.
                    self.loop.run_in_executor(self.executor,
                        lambda: RIS(
                            queue=self.queue,
                            events=self.events,
                            memory=self.memory,
                            db=self.db
                        ).start())

                    # Sender.
                    self.loop.run_in_executor(self.executor,
                        lambda: Sender(
                            queue=self.queue,
                            callback=lambda offset, topic, partition, peer, index, size: (
                                self.db.set(
                                    f'offset_{peer}'.encode('utf-8'),
                                    struct.pack('>dQQ', offset, index, size)
                                ) if peer is not None else None
                            ),
                            events=self.events,
                            memory=self.memory,
                            db=self.db
                        ).start())

                case 'routeviews':
                    # BMPv3 init.
                    message = BMPv3.init_message(
                        router_name=Config.RELAY_NAME,
                        router_descr=f"{Config.RELAY_NAME}.routeviews.org"
                    )
                    self.queue.put((message, 0, None, -1, None, -1, -1, False))

                    # Seed the collector.
                    # This is a recurring operation.
                    index = f"{Config.RELAY_INDEX}/{datetime.now().year}.{datetime.now().month:02d}/RIBS"
                    response = requests.get(index, timeout=300)
                    response.raise_for_status()
                    filename = BeautifulSoup(
                        response.text,
                        'html.parser'
                    ).find_all('a')[-1].text
                    self.loop.run_in_executor(self.executor, 
                        lambda: Seeder(
                            queue=self.queue,
                            url=f"{index}/{filename}",
                            events=self.events,
                            memory=self.memory,
                            db=self.db,
                        ).start())

                    # RouteViews (University of Oregon)
                    # We connect to the Kafka stream.
                    self.loop.run_in_executor(self.executor,
                        lambda: RouteViews(
                            queue=self.queue,
                            events=self.events,
                            memory=self.memory,
                            db=self.db
                        ).start())

                    # Sender.
                    self.loop.run_in_executor(self.executor,
                        lambda: Sender(
                            queue=self.queue,
                            callback=lambda offset, topic, partition, peer, index, size: (
                                self.db.set(
                                    f'offset_{topic}_{partition}'.encode('utf-8'),
                                    offset.to_bytes(16, byteorder='big')
                                ) if partition != -1 and topic is not None else None
                            ),
                            events=self.events,
                            memory=self.memory,
                            db=self.db
                        ).start())
                case _:
                    raise ValueError(f"Invalid type: {Config.RELAY_TYPE}")

            # Wait for Event.
            self.events['shutdown'].wait()
        
        except Exception as e:
            logger.error(e, exc_info=True)
            self.stop()
            raise e
        finally:
            # Exit.
            os._exit(1)

## Main ##

async def main():
    Config.validate()
    relay = Relay()
    relay.start()

if __name__ == "__main__":
    asyncio.run(main())