#!/usr/bin/env python3 import argparse import database import json import logging import sys import typing import multiprocessing NUMBER_THREADS = 2 BLOCK_SIZE = 100 # select, confirm, write FUNCTION_MAP: typing.Any = { 'a': ( database.Database.get_ip4, database.Database.get_domain_in_zone, database.Database.set_hostname, ), 'cname': ( database.Database.get_domain, database.Database.get_domain_in_zone, database.Database.set_hostname, ), 'ptr': ( database.Database.get_domain, database.Database.get_ip4_in_network, database.Database.set_ip4address, ), } class Reader(multiprocessing.Process): def __init__(self, lines_queue: multiprocessing.Queue, write_queue: multiprocessing.Queue, index: int = 0): super(Reader, self).__init__() self.log = logging.getLogger(f'rd{index:03d}') self.lines_queue = lines_queue self.write_queue = write_queue self.index = index def run(self) -> None: self.db = database.Database(write=False) self.db.log = logging.getLogger(f'db{self.index:03d}') self.db.enter_step('line_wait') block: typing.List[str] try: for block in iter(self.lines_queue.get, None): for line in block: dtype, updated, name, value = line self.db.enter_step('feed_switch') select, confirm, write = FUNCTION_MAP[dtype] for rule in select(self.db, value): if not any(confirm(self.db, name)): self.db.enter_step('wait_put') self.write_queue.put((write, name, updated)) self.db.enter_step('line_wait') except KeyboardInterrupt: self.log.error('Interrupted') self.db.enter_step('end') self.db.close() class Writer(multiprocessing.Process): def __init__(self, write_queue: multiprocessing.Queue, ): super(Writer, self).__init__() self.log = logging.getLogger(f'wr ') self.write_queue = write_queue def run(self) -> None: self.db = database.Database(write=True) self.db.log = logging.getLogger(f'dbw ') self.db.enter_step('line_wait') block: typing.List[str] try: fun: typing.Callable name: str updated: int for fun, name, updated in iter(self.write_queue.get, None): self.db.enter_step('exec') fun(self.db, name, updated) self.db.enter_step('line_wait') except KeyboardInterrupt: self.log.error('Interrupted') self.db.enter_step('end') self.db.close() if __name__ == '__main__': # Parsing arguments log = logging.getLogger('feed_dns') parser = argparse.ArgumentParser( description="TODO") parser.add_argument( # '-i', '--input', type=argparse.FileType('rb'), default=sys.stdin.buffer, '-i', '--input', type=argparse.FileType('r'), default=sys.stdin, help="TODO") args = parser.parse_args() DB = database.Database(write=False) # Not needed, just for timing DB.log = logging.getLogger('db ') lines_queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=100) write_queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=100) DB.enter_step('proc_create') readers: typing.List[Reader] = list() for w in range(NUMBER_THREADS): readers.append(Reader(lines_queue, write_queue, w)) writer = Writer(write_queue) DB.enter_step('proc_start') for reader in readers: reader.start() writer.start() try: block: typing.List[str] = list() DB.enter_step('iowait') for line in args.input: DB.enter_step('block_append') DB.enter_step('feed_json_parse') data = json.loads(line) line = (data['type'], int(data['timestamp']), data['name'], data['value']) block.append(line) if len(block) >= BLOCK_SIZE: DB.enter_step('wait_put') lines_queue.put(block) block = list() DB.enter_step('iowait') DB.enter_step('wait_put') lines_queue.put(block) DB.enter_step('end_put') for _ in range(NUMBER_THREADS): lines_queue.put(None) write_queue.put(None) DB.enter_step('proc_join') for reader in readers: reader.join() writer.join() except KeyboardInterrupt: log.error('Interrupted') DB.close()