orchestrator.py

#!/usr/bin/env python3

import argparse
import logging
from multiprocessing import pool
import sys

from respdiff import cli, sendrecv
from respdiff.database import LMDB, MetaDatabase


def main():
    cli.setup_logging()
    parser = argparse.ArgumentParser(
        description="read queries from LMDB, send them in parallel to servers "
        "listed in configuration file, and record answers into LMDB"
    )
    cli.add_arg_envdir(parser)
    cli.add_arg_config(parser)
    parser.add_argument(
        "--ignore-timeout",
        action="store_true",
        help="continue despite consecutive timeouts from resolvers",
    )

    args = parser.parse_args()
    sendrecv.module_init(args)

    with LMDB(args.envdir) as lmdb:
        meta = MetaDatabase(lmdb, args.cfg["servers"]["names"], create=True)
        meta.write_version()
        meta.write_start_time()

        lmdb.open_db(LMDB.QUERIES)
        adb = lmdb.open_db(LMDB.ANSWERS, create=True, check_notexists=True)

        qstream = lmdb.key_value_stream(LMDB.QUERIES)
        txn = lmdb.env.begin(adb, write=True)
        try:
            # process queries in parallel
            with pool.Pool(
                processes=args.cfg["sendrecv"]["jobs"], initializer=sendrecv.worker_init
            ) as p:
                i = 0
                for qkey, blob in p.imap_unordered(
                    sendrecv.worker_perform_query, qstream, chunksize=100
                ):
                    i += 1
                    if i % 10000 == 0:
                        logging.info("Received {:d} answers".format(i))
                    txn.put(qkey, blob)
        except KeyboardInterrupt:
            logging.info("SIGINT received, exiting...")
            sys.exit(130)
        except RuntimeError as err:
            logging.error(err)
            sys.exit(1)
        finally:
            # attempt to preserve data if something went wrong (or not)
            logging.debug("Comitting LMDB transaction...")
            txn.commit()
            meta.write_end_time()


if __name__ == "__main__":
    main()