brozzler/bin/brozzler-worker

#!/usr/bin/env python
#
# brozzler-worker - main entrypoint for brozzler, gets sites and pages to
# brozzle from rethinkdb, brozzles them
#
# Copyright (C) 2014-2016 Internet Archive
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#


import argparse
import os
import sys
import logging
import brozzler
import brozzler.worker
import threading
import time
import signal
import pprint
import traceback
import rethinkstuff
import warnings
import requests

arg_parser = argparse.ArgumentParser(prog=os.path.basename(__file__),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
arg_parser.add_argument('--rethinkdb-servers', dest='rethinkdb_servers', default="localhost",
        help='rethinkdb servers, e.g. db0.foo.org,db0.foo.org:38015,db1.foo.org')
arg_parser.add_argument('--rethinkdb-db', dest='rethinkdb_db', default="brozzler",
        help='rethinkdb database name')
arg_parser.add_argument('-e', '--executable', dest='chrome_exe', default='chromium-browser',
        help='executable to use to invoke chrome')
arg_parser.add_argument('-n', '--max-browsers', dest='max_browsers', default='1',
        help='max number of chrome instances simultaneously browsing pages')
arg_parser.add_argument('-v', '--verbose', dest='log_level',
        action="store_const", default=logging.INFO, const=logging.DEBUG)
arg_parser.add_argument('--version', action='version',
        version="brozzler {} - {}".format(brozzler.__version__, os.path.basename(__file__)))
args = arg_parser.parse_args(args=sys.argv[1:])

logging.basicConfig(stream=sys.stdout, level=args.log_level,
        format='%(asctime)s %(process)d %(levelname)s %(threadName)s %(name)s.%(funcName)s(%(filename)s:%(lineno)d) %(message)s')
logging.getLogger("requests.packages.urllib3").setLevel(logging.WARN)
warnings.simplefilter("ignore", category=requests.packages.urllib3.exceptions.InsecureRequestWarning)
warnings.simplefilter("ignore", category=requests.packages.urllib3.exceptions.InsecurePlatformWarning)

def sigterm(signum, frame):
    raise brozzler.ShutdownRequested('shutdown requested (caught SIGTERM)')
def sigint(signum, frame):
    raise brozzler.ShutdownRequested('shutdown requested (caught SIGINT)')

def dump_state(signum, frame):
    pp = pprint.PrettyPrinter(indent=4)
    state_strs = []

    for th in threading.enumerate():
        state_strs.append(str(th))
        stack = traceback.format_stack(sys._current_frames()[th.ident])
        state_strs.append("".join(stack))

    logging.warn("dumping state (caught signal {})\n{}".format(
        signum, "\n".join(state_strs)))

signal.signal(signal.SIGQUIT, dump_state)
signal.signal(signal.SIGTERM, sigterm)
signal.signal(signal.SIGINT, sigint)

r = rethinkstuff.Rethinker(args.rethinkdb_servers.split(","), args.rethinkdb_db)
frontier = brozzler.RethinkDbFrontier(r)
service_registry = rethinkstuff.ServiceRegistry(r)
worker = brozzler.worker.BrozzlerWorker(
        frontier, service_registry, max_browsers=int(args.max_browsers),
        chrome_exe=args.chrome_exe)

worker_thread = worker.start()

try:
    while worker_thread.is_alive():
        time.sleep(0.5)
    logging.critical("worker thread has died, shutting down")
except brozzler.ShutdownRequested as e:
    pass
finally:
    worker.shutdown_now()
    for th in threading.enumerate():
        if th != threading.current_thread():
            th.join()

logging.info("all done, exiting")