<commit_before>#!/usr/bin/env python

"""Python Capability Container start script"""

__author__ = 'Adam R. Smith, Michael Meisinger'
__license__ = 'Apache 2.0'

import argparse
import yaml
import sys
import traceback
from uuid import uuid4

#
# WARNING - DO NOT IMPORT GEVENT OR PYON HERE. IMPORTS **MUST** BE DONE IN THE MAIN()
# DUE TO DAEMONIZATION.
#
# SEE: http://groups.google.com/group/gevent/browse_thread/thread/6223805ffcd5be22?pli=1
#

version = "2.0"     # TODO: extract this from the code once versioning is automated again
description = '''
pyon (ION capability container) v%s
''' % (version)

def setup_ipython(shell_api=None):
    from IPython.config.loader import Config
    cfg = Config()
    shell_config = cfg.InteractiveShellEmbed
    shell_config.prompt_in1 = '><> '
    shell_config.prompt_in2 = '... '
    shell_config.prompt_out = '--> '
    shell_config.confirm_exit = False

    # monkeypatch the ipython inputhook to be gevent-friendly
    import gevent   # should be auto-monkey-patched by pyon already.
    import select
    import sys
    def stdin_ready():
        infds, outfds, erfds = select.select([sys.stdin], [], [], 0)
        if infds:
            return True
        else:
            return False

    def inputhook_gevent():
        try:
            while not stdin_ready():
                gevent.sleep(0.001)
        except KeyboardInterrupt:
            pass

        return 0

    # install the gevent inputhook
    from IPython.lib.inputhook import inputhook_manager
    inputhook_manager.set_inputhook(inputhook_gevent)
    inputhook_manager._current_gui = 'gevent'

    # First import the embeddable shell class
    from IPython.frontend.terminal.embed import InteractiveShellEmbed

    # Update namespace of interactive shell
    # TODO: Cleanup namespace even further
    if shell_api is not None:
        locals().update(shell_api)

    # Now create an instance of the embeddable shell. The first argument is a
    # string with options exactly as you would type them if you were starting
    # IPython at the system command line. Any parameters you want to define for
    # configuration can thus be specified here.
    ipshell = InteractiveShellEmbed(config=cfg,
                           banner1 = \
"""    ____                                ________  _   __   ____________   ____  ___
   / __ \__  ______  ____              /  _/ __ \/ | / /  / ____/ ____/  / __ \|__ \\
  / /_/ / / / / __ \/ __ \   ______    / // / / /  |/ /  / /   / /      / /_/ /__/ /
 / ____/ /_/ / /_/ / / / /  /_____/  _/ // /_/ / /|  /  / /___/ /___   / _, _// __/
/_/    \__, /\____/_/ /_/           /___/\____/_/ |_/   \____/\____/  /_/ |_|/____/
      /____/""",
                           exit_msg = 'Leaving ION shell, shutting down container.')

    ipshell('Pyon - ION R2 CC interactive IPython shell. Type ionhelp() for help')

# From http://stackoverflow.com/questions/6037503/python-unflatten-dict/6037657#6037657
def unflatten(dictionary):
    resultDict = dict()
    for key, value in dictionary.iteritems():
        parts = key.split(".")
        d = resultDict
        for part in parts[:-1]:
            if part not in d:
                d[part] = dict()
            d = d[part]
        d[parts[-1]] = value
    return resultDict

def main(opts, *args, **kwargs):
    """
    Starts the capability container and processes arguments
    """

    def prepare_container():
        import threading
        threading.current_thread().name = "CC-Main"

        # SIDE EFFECT: The import of pyon.public triggers many module initializers:
        # pyon.core.bootstrap (Config load, logging setup), etc.
        from pyon.public import Container, CFG
        from pyon.util.containers import dict_merge
        from pyon.util.config import Config

        # Check if user opted to override logging config
        if opts.logcfg:
            from pyon.util.config import logging_conf_paths, initialize_logging
            # Re-initialize logging
            logging_conf_paths.append(opts.logcfg)
            initialize_logging()

        # Set that system is not testing. We are running as standalone container
        dict_merge(CFG, {'system':{'testing':False}}, True)

        # Also set the immediate flag, but only if specified - it is an override
        if opts.immediate:
            dict_merge(CFG, {'system':{'immediate':True}}, True)

        # Load any additional config paths and merge them into main config
        if len(opts.config):
            cfg = Config(opts.config)
            dict_merge(CFG, cfg.data, True)

        # Create the container instance
        container = Container(*args, **kwargs)

        return container

    def start_container(container):
        """
        Start container and all internal managers. Returns when ready.
        """
        container.start()

    def do_work(container):
        """
        Performs initial startup actions with the container as requested in arguments.
        Then remains in container shell or infinite wait until container stops.
        Returns when container should stop. Raises an exception if anything failed.
        """
        if opts.proc:
            # Run a one-off process (with the -x argument)
            mod, proc = opts.proc.rsplit('.', 1)
            print "Starting process %s" % opts.proc
            container.spawn_process(proc, mod, proc, config={'process':{'type':'immediate'}})
            # And end
            return

        if opts.rel:
            # Start a rel file
            start_ok = container.start_rel_from_url(opts.rel)
            if not start_ok:
                raise Exception("Cannot start deploy file '%s'" % opts.rel)

        if opts.mx:
            container.spawn_process("ContainerUI", "ion.core.containerui", "ContainerUI")
            print "Container UI started ... listening on http://localhost:8080"

        if not opts.noshell and not opts.daemon:
            # Keep container running while there is an interactive shell
            from pyon.container.shell_api import get_shell_api
            setup_ipython(get_shell_api(container))
        else:
            # Keep container running until process terminated
            container.serve_forever()

    def stop_container(container):
        try:
            if container:
                container.stop()
            return True
        except Exception as ex:
            # We want to make sure to get out here alive
            print "CONTAINER STOP ERROR"
            traceback.print_exc()
            return False

    # ----------------------------------------------------------------------------------
    # Container life cycle

    container = None
    try:
        container = prepare_container()

        start_container(container)
    except Exception as ex:
        print "===== CONTAINER START ERROR -- FAIL ====="
        traceback.print_exc()
        stop_container(container)
        sys.exit(1)

    try:
        do_work(container)
    except Exception as ex:
        stop_container(container)
        print "===== CONTAINER PROCESS START ERROR -- ABORTING ====="
        print ex
        sys.exit(1)

    # Assumption: stop is so robust, it does not fail even if it was only partially started
    stop_ok = stop_container(container)
    if not stop_ok:
        sys.exit(1)

def parse_args(tokens):
    """ Exploit yaml's spectacular type inference (and ensure consistency with config files) """
    args, kwargs = [], {}
    for token in tokens:
        token = token.lstrip('-')
        if '=' in token:
            key,val = token.split('=', 1)
            cfg = unflatten({key: yaml.load(val)})
            kwargs.update(cfg)
        else:
            args.append(yaml.load(token))

    return args, kwargs

def entry():
    #proc_types = GreenProcessSupervisor.type_callables.keys()

    # NOTE: Resist the temptation to add other parameters here! Most container config options
    # should be in the config file (pyon.yml), which can also be specified on the command-line via the extra args

    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('-d', '--daemon', action='store_true')
    parser.add_argument('-n', '--noshell', action='store_true')
    parser.add_argument('-m', '--mx', action='store_true', help='Start a management web UI')
    parser.add_argument('-r', '--rel', type=str, help='Path to a rel file to launch.')
    parser.add_argument('-l', '--logcfg', type=str, help='Path to logging configuration file.')
    parser.add_argument('-x', '--proc', type=str, help='Qualified name of process to start and then exit.')
    parser.add_argument('-i', '--immediate', action='store_true', help='Will exit the container if the only procs started are immediate proc types. Sets CFG system.immediate flag.')
    parser.add_argument('-p', '--pidfile', type=str, help='PID file to use when --daemon specified. Defaults to cc-<rand>.pid')
    parser.add_argument('-c', '--config', action='append', type=str, help='Additional config files to load.', default=[])
    parser.add_argument('-v', '--version', action='version', version='pyon v%s' % (version))
    opts, extra = parser.parse_known_args()
    args, kwargs = parse_args(extra)

    if opts.daemon:
        # TODO: The daemonizing code may need to be moved inside the Container class (so it happens per-process)
        from daemon import DaemonContext
        from lockfile import FileLock

        #logg = open('hi.txt', 'w+')
        #slogg = open('hi2.txt', 'w+')

        # TODO: May need to generate a pidfile based on some parameter or cc name
        pidfile = opts.pidfile or 'cc-%s.pid' % str(uuid4())[0:4]
        with DaemonContext(pidfile=FileLock(pidfile)):#, stdout=logg, stderr=slogg):
            print "Starting ION CC ... deamon=True, opts=%s" % str(opts)
            main(opts, *args, **kwargs)
    else:
        print "Starting ION CC ... deamon=False, opts=%s" % str(opts)
        main(opts, *args, **kwargs)

if __name__ == '__main__':
    entry()
<commit_msg>Luke: cfg->ipython_cfg<commit_after>#!/usr/bin/env python

"""Python Capability Container start script"""

__author__ = 'Adam R. Smith, Michael Meisinger'
__license__ = 'Apache 2.0'

import argparse
import yaml
import sys
import traceback
from uuid import uuid4

#
# WARNING - DO NOT IMPORT GEVENT OR PYON HERE. IMPORTS **MUST** BE DONE IN THE MAIN()
# DUE TO DAEMONIZATION.
#
# SEE: http://groups.google.com/group/gevent/browse_thread/thread/6223805ffcd5be22?pli=1
#

version = "2.0"     # TODO: extract this from the code once versioning is automated again
description = '''
pyon (ION capability container) v%s
''' % (version)

def setup_ipython(shell_api=None):
    from IPython.config.loader import Config
    ipython_cfg = Config()
    shell_config = ipython_cfg.InteractiveShellEmbed
    shell_config.prompt_in1 = '><> '
    shell_config.prompt_in2 = '... '
    shell_config.prompt_out = '--> '
    shell_config.confirm_exit = False

    # monkeypatch the ipython inputhook to be gevent-friendly
    import gevent   # should be auto-monkey-patched by pyon already.
    import select
    import sys
    def stdin_ready():
        infds, outfds, erfds = select.select([sys.stdin], [], [], 0)
        if infds:
            return True
        else:
            return False

    def inputhook_gevent():
        try:
            while not stdin_ready():
                gevent.sleep(0.001)
        except KeyboardInterrupt:
            pass

        return 0

    # install the gevent inputhook
    from IPython.lib.inputhook import inputhook_manager
    inputhook_manager.set_inputhook(inputhook_gevent)
    inputhook_manager._current_gui = 'gevent'

    # First import the embeddable shell class
    from IPython.frontend.terminal.embed import InteractiveShellEmbed

    # Update namespace of interactive shell
    # TODO: Cleanup namespace even further
    if shell_api is not None:
        locals().update(shell_api)

    # Now create an instance of the embeddable shell. The first argument is a
    # string with options exactly as you would type them if you were starting
    # IPython at the system command line. Any parameters you want to define for
    # configuration can thus be specified here.
    ipshell = InteractiveShellEmbed(config=ipython_cfg,
                           banner1 = \
"""    ____                                ________  _   __   ____________   ____  ___
   / __ \__  ______  ____              /  _/ __ \/ | / /  / ____/ ____/  / __ \|__ \\
  / /_/ / / / / __ \/ __ \   ______    / // / / /  |/ /  / /   / /      / /_/ /__/ /
 / ____/ /_/ / /_/ / / / /  /_____/  _/ // /_/ / /|  /  / /___/ /___   / _, _// __/
/_/    \__, /\____/_/ /_/           /___/\____/_/ |_/   \____/\____/  /_/ |_|/____/
      /____/""",
                           exit_msg = 'Leaving ION shell, shutting down container.')

    ipshell('Pyon - ION R2 CC interactive IPython shell. Type ionhelp() for help')

# From http://stackoverflow.com/questions/6037503/python-unflatten-dict/6037657#6037657
def unflatten(dictionary):
    resultDict = dict()
    for key, value in dictionary.iteritems():
        parts = key.split(".")
        d = resultDict
        for part in parts[:-1]:
            if part not in d:
                d[part] = dict()
            d = d[part]
        d[parts[-1]] = value
    return resultDict

def main(opts, *args, **kwargs):
    """
    Starts the capability container and processes arguments
    """

    def prepare_container():
        import threading
        threading.current_thread().name = "CC-Main"

        # SIDE EFFECT: The import of pyon.public triggers many module initializers:
        # pyon.core.bootstrap (Config load, logging setup), etc.
        from pyon.public import Container, CFG
        from pyon.util.containers import dict_merge
        from pyon.util.config import Config

        # Check if user opted to override logging config
        if opts.logcfg:
            from pyon.util.config import logging_conf_paths, initialize_logging
            # Re-initialize logging
            logging_conf_paths.append(opts.logcfg)
            initialize_logging()

        # Set that system is not testing. We are running as standalone container
        dict_merge(CFG, {'system':{'testing':False}}, True)

        # Also set the immediate flag, but only if specified - it is an override
        if opts.immediate:
            dict_merge(CFG, {'system':{'immediate':True}}, True)

        # Load any additional config paths and merge them into main config
        if len(opts.config):
            ipython_cfg = Config(opts.config)
            dict_merge(CFG, ipython_cfg.data, True)

        # Create the container instance
        container = Container(*args, **kwargs)

        return container

    def start_container(container):
        """
        Start container and all internal managers. Returns when ready.
        """
        container.start()

    def do_work(container):
        """
        Performs initial startup actions with the container as requested in arguments.
        Then remains in container shell or infinite wait until container stops.
        Returns when container should stop. Raises an exception if anything failed.
        """
        if opts.proc:
            # Run a one-off process (with the -x argument)
            mod, proc = opts.proc.rsplit('.', 1)
            print "Starting process %s" % opts.proc
            container.spawn_process(proc, mod, proc, config={'process':{'type':'immediate'}})
            # And end
            return

        if opts.rel:
            # Start a rel file
            start_ok = container.start_rel_from_url(opts.rel)
            if not start_ok:
                raise Exception("Cannot start deploy file '%s'" % opts.rel)

        if opts.mx:
            container.spawn_process("ContainerUI", "ion.core.containerui", "ContainerUI")
            print "Container UI started ... listening on http://localhost:8080"

        if not opts.noshell and not opts.daemon:
            # Keep container running while there is an interactive shell
            from pyon.container.shell_api import get_shell_api
            setup_ipython(get_shell_api(container))
        else:
            # Keep container running until process terminated
            container.serve_forever()

    def stop_container(container):
        try:
            if container:
                container.stop()
            return True
        except Exception as ex:
            # We want to make sure to get out here alive
            print "CONTAINER STOP ERROR"
            traceback.print_exc()
            return False

    # ----------------------------------------------------------------------------------
    # Container life cycle

    container = None
    try:
        container = prepare_container()

        start_container(container)
    except Exception as ex:
        print "===== CONTAINER START ERROR -- FAIL ====="
        traceback.print_exc()
        stop_container(container)
        sys.exit(1)

    try:
        do_work(container)
    except Exception as ex:
        stop_container(container)
        print "===== CONTAINER PROCESS START ERROR -- ABORTING ====="
        print ex
        sys.exit(1)

    # Assumption: stop is so robust, it does not fail even if it was only partially started
    stop_ok = stop_container(container)
    if not stop_ok:
        sys.exit(1)

def parse_args(tokens):
    """ Exploit yaml's spectacular type inference (and ensure consistency with config files) """
    args, kwargs = [], {}
    for token in tokens:
        token = token.lstrip('-')
        if '=' in token:
            key,val = token.split('=', 1)
            ipython_cfg = unflatten({key: yaml.load(val)})
            kwargs.update(ipython_cfg)
        else:
            args.append(yaml.load(token))

    return args, kwargs

def entry():
    #proc_types = GreenProcessSupervisor.type_callables.keys()

    # NOTE: Resist the temptation to add other parameters here! Most container config options
    # should be in the config file (pyon.yml), which can also be specified on the command-line via the extra args

    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('-d', '--daemon', action='store_true')
    parser.add_argument('-n', '--noshell', action='store_true')
    parser.add_argument('-m', '--mx', action='store_true', help='Start a management web UI')
    parser.add_argument('-r', '--rel', type=str, help='Path to a rel file to launch.')
    parser.add_argument('-l', '--logcfg', type=str, help='Path to logging configuration file.')
    parser.add_argument('-x', '--proc', type=str, help='Qualified name of process to start and then exit.')
    parser.add_argument('-i', '--immediate', action='store_true', help='Will exit the container if the only procs started are immediate proc types. Sets CFG system.immediate flag.')
    parser.add_argument('-p', '--pidfile', type=str, help='PID file to use when --daemon specified. Defaults to cc-<rand>.pid')
    parser.add_argument('-c', '--config', action='append', type=str, help='Additional config files to load.', default=[])
    parser.add_argument('-v', '--version', action='version', version='pyon v%s' % (version))
    opts, extra = parser.parse_known_args()
    args, kwargs = parse_args(extra)

    if opts.daemon:
        # TODO: The daemonizing code may need to be moved inside the Container class (so it happens per-process)
        from daemon import DaemonContext
        from lockfile import FileLock

        #logg = open('hi.txt', 'w+')
        #slogg = open('hi2.txt', 'w+')

        # TODO: May need to generate a pidfile based on some parameter or cc name
        pidfile = opts.pidfile or 'cc-%s.pid' % str(uuid4())[0:4]
        with DaemonContext(pidfile=FileLock(pidfile)):#, stdout=logg, stderr=slogg):
            print "Starting ION CC ... deamon=True, opts=%s" % str(opts)
            main(opts, *args, **kwargs)
    else:
        print "Starting ION CC ... deamon=False, opts=%s" % str(opts)
        main(opts, *args, **kwargs)

if __name__ == '__main__':
    entry()
<|endoftext|>


<commit_before># pylint: skip-file

#
# The Qubes OS Project, http://www.qubes-os.org
#
# Copyright (C) 2010  Rafal Wojtczuk  <rafal@invisiblethingslab.com>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, see <https://www.gnu.org/licenses/>.
#
#
import configparser
import socketserver
import logging
import logging.handlers
import os
import socket
import sys
import threading

import xen.lowlevel.xs

import qubes.qmemman
import qubes.qmemman.algo
import qubes.utils

SOCK_PATH = '/var/run/qubes/qmemman.sock'

system_state = qubes.qmemman.SystemState()
global_lock = threading.Lock()
# If XS_Watcher will
# handle meminfo event before @introduceDomain, it will use
# incomplete domain list for that and may redistribute memory
# allocated to some VM, but not yet used (see #1389).
# To fix that, system_state should be updated (refresh domain
# list) before processing other changes, every time some process requested
# memory for a new VM, before releasing the lock. Then XS_Watcher will check
# this flag before processing other event.
force_refresh_domain_list = False

def only_in_first_list(l1, l2):
    ret = []
    for i in l1:
        if not i in l2:
            ret.append(i)
    return ret

def get_domain_meminfo_key(domain_id):
    return '/local/domain/'+domain_id+'/memory/meminfo'

class WatchType(object):
    def __init__(self, fn, param):
        self.fn = fn
        self.param = param

class XS_Watcher(object):
    def __init__(self):
        self.log = logging.getLogger('qmemman.daemon.xswatcher')
        self.log.debug('XS_Watcher()')

        self.handle = xen.lowlevel.xs.xs()
        self.handle.watch('@introduceDomain', WatchType(
            XS_Watcher.domain_list_changed, False))
        self.handle.watch('@releaseDomain', WatchType(
            XS_Watcher.domain_list_changed, False))
        self.watch_token_dict = {}

    def domain_list_changed(self, refresh_only=False):
        """
        Check if any domain was created/destroyed. If it was, update
        appropriate list. Then redistribute memory.

        :param refresh_only If True, only refresh domain list, do not
        redistribute memory. In this mode, caller must already hold
        global_lock.
        """
        self.log.debug('domain_list_changed(only_refresh={!r})'.format(
            refresh_only))

        got_lock = False
        if not refresh_only:
            self.log.debug('acquiring global_lock')
            global_lock.acquire()
            got_lock = True
            self.log.debug('global_lock acquired')
        try:
            curr = self.handle.ls('', '/local/domain')
            if curr is None:
                return

            # check if domain is really there, it may happen that some empty
            # directories are left in xenstore
            curr = list(filter(
                lambda x:
                self.handle.read('',
                                 '/local/domain/{}/domid'.format(x)
                                 ) is not None,
                curr
            ))
            self.log.debug('curr={!r}'.format(curr))

            for i in only_in_first_list(curr, self.watch_token_dict.keys()):
                # new domain has been created
                watch = WatchType(XS_Watcher.meminfo_changed, i)
                self.watch_token_dict[i] = watch
                self.handle.watch(get_domain_meminfo_key(i), watch)
                system_state.add_domain(i)

            for i in only_in_first_list(self.watch_token_dict.keys(), curr):
                # domain destroyed
                self.handle.unwatch(get_domain_meminfo_key(i), self.watch_token_dict[i])
                self.watch_token_dict.pop(i)
                system_state.del_domain(i)
        except:
            self.log.exception('Updating domain list failed')
        finally:
            if got_lock:
                global_lock.release()
                self.log.debug('global_lock released')

        if not refresh_only:
            try:
                system_state.do_balance()
            except:
                self.log.exception('do_balance() failed')

    def meminfo_changed(self, domain_id):
        self.log.debug('meminfo_changed(domain_id={!r})'.format(domain_id))
        untrusted_meminfo_key = self.handle.read(
            '', get_domain_meminfo_key(domain_id))
        if untrusted_meminfo_key == None or untrusted_meminfo_key == b'':
            return

        self.log.debug('acquiring global_lock')
        global_lock.acquire()
        self.log.debug('global_lock acquired')
        try:
            global force_refresh_domain_list
            if force_refresh_domain_list:
                self.domain_list_changed(refresh_only=True)
                force_refresh_domain_list = False
            if domain_id not in self.watch_token_dict:
                # domain just destroyed
                return

            system_state.refresh_meminfo(domain_id, untrusted_meminfo_key)
        except:
            self.log.exception('Updating meminfo for %d failed', domain_id)
        finally:
            global_lock.release()
            self.log.debug('global_lock released')

    def watch_loop(self):
        self.log.debug('watch_loop()')
        while True:
            result = self.handle.read_watch()
            self.log.debug('watch_loop result={!r}'.format(result))
            token = result[1]
            token.fn(self, token.param)

class QMemmanReqHandler(socketserver.BaseRequestHandler):
    """
    The RequestHandler class for our server.

    It is instantiated once per connection to the server, and must
    override the handle() method to implement communication to the
    client.
    """

    def handle(self):
        self.log = logging.getLogger('qmemman.daemon.reqhandler')

        got_lock = False
        try:
            # self.request is the TCP socket connected to the client
            while True:
                self.data = self.request.recv(1024).strip()
                self.log.debug('data={!r}'.format(self.data))
                if len(self.data) == 0:
                    self.log.info('client disconnected, resuming membalance')
                    if got_lock:
                        global force_refresh_domain_list
                        force_refresh_domain_list = True
                    return

                # XXX something is wrong here: return without release?
                if got_lock:
                    self.log.warning('Second request over qmemman.sock?')
                    return

                self.log.debug('acquiring global_lock')
                global_lock.acquire()
                self.log.debug('global_lock acquired')

                got_lock = True
                if system_state.do_balloon(int(self.data.decode('ascii'))):
                    resp = b"OK\n"
                else:
                    resp = b"FAIL\n"
                self.log.debug('resp={!r}'.format(resp))
                self.request.send(resp)
        except BaseException as e:
            self.log.exception(
                "exception while handling request: {!r}".format(e))
        finally:
            if got_lock:
                global_lock.release()
                self.log.debug('global_lock released')

parser = qubes.tools.QubesArgumentParser(want_app=False)

parser.add_argument('--config', '-c', metavar='FILE',
    action='store', default='/etc/qubes/qmemman.conf',
    help='qmemman config file')

parser.add_argument('--foreground',
    action='store_true', default=False,
    help='do not close stdio')

def main():
    args = parser.parse_args()

    # setup logging
    ha_syslog = logging.handlers.SysLogHandler('/dev/log')
    ha_syslog.setFormatter(
        logging.Formatter('%(name)s[%(process)d]: %(message)s'))
    logging.root.addHandler(ha_syslog)

    if args.foreground:
        ha_stderr = logging.StreamHandler(sys.stderr)
        ha_file.setFormatter(
            logging.Formatter('%(asctime)s %(name)s[%(process)d]: %(message)s'))
        logging.root.addHandler(ha_stderr)

    sys.stdin.close()

    log = logging.getLogger('qmemman.daemon')

    config = configparser.SafeConfigParser({
            'vm-min-mem': str(qubes.qmemman.algo.MIN_PREFMEM),
            'dom0-mem-boost': str(qubes.qmemman.algo.DOM0_MEM_BOOST),
            'cache-margin-factor': str(qubes.qmemman.algo.CACHE_FACTOR)
            })
    config.read(args.config)

    if config.has_section('global'):
        qubes.qmemman.algo.MIN_PREFMEM = \
            qubes.utils.parse_size(config.get('global', 'vm-min-mem'))
        qubes.qmemman.algo.DOM0_MEM_BOOST = \
            qubes.utils.parse_size(config.get('global', 'dom0-mem-boost'))
        qubes.qmemman.algo.CACHE_FACTOR = \
            config.getfloat('global', 'cache-margin-factor')
        loglevel = config.getint('global', 'log-level')
        logging.root.setLevel(loglevel)

    log.info('MIN_PREFMEM={algo.MIN_PREFMEM}'
        ' DOM0_MEM_BOOST={algo.DOM0_MEM_BOOST}'
        ' CACHE_FACTOR={algo.CACHE_FACTOR}'.format(
            algo=qubes.qmemman.algo))

    try:
        os.unlink(SOCK_PATH)
    except:
        pass

    log.debug('instantiating server')
    os.umask(0)

    # Initialize the connection to Xen and to XenStore
    system_state.init()

    server = socketserver.UnixStreamServer(SOCK_PATH, QMemmanReqHandler)
    os.umask(0o077)

    # notify systemd
    nofity_socket = os.getenv('NOTIFY_SOCKET')
    if nofity_socket:
        log.debug('notifying systemd')
        s = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
        if nofity_socket.startswith('@'):
            nofity_socket = '\0%s' % nofity_socket[1:]
        s.connect(nofity_socket)
        s.sendall(b"READY=1")
        s.close()

    threading.Thread(target=server.serve_forever).start()
    XS_Watcher().watch_loop()
<commit_msg>Enable pylint for qubes.tools.qmemmand<commit_after>#
# The Qubes OS Project, http://www.qubes-os.org
#
# Copyright (C) 2010  Rafal Wojtczuk  <rafal@invisiblethingslab.com>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, see <https://www.gnu.org/licenses/>.
#
#
# pylint: disable=global-statement

import configparser
import socketserver
import logging
import logging.handlers
import os
import socket
import sys
import threading
from dataclasses import dataclass
from typing import Callable, Any

import xen.lowlevel.xs  # pylint: disable=import-error

import qubes.qmemman
import qubes.qmemman.algo
import qubes.utils

SOCK_PATH = '/var/run/qubes/qmemman.sock'

system_state = qubes.qmemman.SystemState()
global_lock = threading.Lock()
# If XSWatcher will
# handle meminfo event before @introduceDomain, it will use
# incomplete domain list for that and may redistribute memory
# allocated to some VM, but not yet used (see #1389).
# To fix that, system_state should be updated (refresh domain
# list) before processing other changes, every time some process requested
# memory for a new VM, before releasing the lock. Then XS_Watcher will check
# this flag before processing other event.
force_refresh_domain_list = False

def only_in_first_list(list1, list2):
    ret = []
    for i in list1:
        if i not in list2:
            ret.append(i)
    return ret

def get_domain_meminfo_key(domain_id):
    return '/local/domain/'+domain_id+'/memory/meminfo'

@dataclass
class WatchType:
    func: Callable
    param: Any

class XSWatcher:
    def __init__(self):
        self.log = logging.getLogger('qmemman.daemon.xswatcher')
        self.log.debug('XS_Watcher()')

        self.handle = xen.lowlevel.xs.xs()
        self.handle.watch('@introduceDomain', WatchType(
            XSWatcher.domain_list_changed, False))
        self.handle.watch('@releaseDomain', WatchType(
            XSWatcher.domain_list_changed, False))
        self.watch_token_dict = {}

    def domain_list_changed(self, refresh_only=False):
        """
        Check if any domain was created/destroyed. If it was, update
        appropriate list. Then redistribute memory.

        :param refresh_only If True, only refresh domain list, do not
        redistribute memory. In this mode, caller must already hold
        global_lock.
        """
        self.log.debug('domain_list_changed(only_refresh={!r})'.format(
            refresh_only))

        got_lock = False
        if not refresh_only:
            self.log.debug('acquiring global_lock')
            global_lock.acquire()
            got_lock = True
            self.log.debug('global_lock acquired')
        try:
            curr = self.handle.ls('', '/local/domain')
            if curr is None:
                return

            # check if domain is really there, it may happen that some empty
            # directories are left in xenstore
            curr = list(filter(
                lambda x:
                self.handle.read('',
                                 '/local/domain/{}/domid'.format(x)
                                 ) is not None,
                curr
            ))
            self.log.debug('curr={!r}'.format(curr))

            for i in only_in_first_list(curr, self.watch_token_dict.keys()):
                # new domain has been created
                watch = WatchType(XSWatcher.meminfo_changed, i)
                self.watch_token_dict[i] = watch
                self.handle.watch(get_domain_meminfo_key(i), watch)
                system_state.add_domain(i)

            for i in only_in_first_list(self.watch_token_dict.keys(), curr):
                # domain destroyed
                self.handle.unwatch(get_domain_meminfo_key(i),
                                    self.watch_token_dict[i])
                self.watch_token_dict.pop(i)
                system_state.del_domain(i)
        except:  # pylint: disable=bare-except
            self.log.exception('Updating domain list failed')
        finally:
            if got_lock:
                global_lock.release()
                self.log.debug('global_lock released')

        if not refresh_only:
            try:
                system_state.do_balance()
            except:  # pylint: disable=bare-except
                self.log.exception('do_balance() failed')

    def meminfo_changed(self, domain_id):
        self.log.debug('meminfo_changed(domain_id={!r})'.format(domain_id))
        untrusted_meminfo_key = self.handle.read(
            '', get_domain_meminfo_key(domain_id))
        if untrusted_meminfo_key is None or untrusted_meminfo_key == b'':
            return

        self.log.debug('acquiring global_lock')
        global_lock.acquire()
        self.log.debug('global_lock acquired')
        try:
            global force_refresh_domain_list
            if force_refresh_domain_list:
                self.domain_list_changed(refresh_only=True)
                force_refresh_domain_list = False
            if domain_id not in self.watch_token_dict:
                # domain just destroyed
                return

            system_state.refresh_meminfo(domain_id, untrusted_meminfo_key)
        except:  # pylint: disable=bare-except
            self.log.exception('Updating meminfo for %d failed', domain_id)
        finally:
            global_lock.release()
            self.log.debug('global_lock released')

    def watch_loop(self):
        self.log.debug('watch_loop()')
        while True:
            result = self.handle.read_watch()
            self.log.debug('watch_loop result={!r}'.format(result))
            token = result[1]
            token.func(self, token.param)

class QMemmanReqHandler(socketserver.BaseRequestHandler):
    """
    The RequestHandler class for our server.

    It is instantiated once per connection to the server, and must
    override the handle() method to implement communication to the
    client.
    """

    def handle(self):
        self.log = logging.getLogger('qmemman.daemon.reqhandler')

        got_lock = False
        try:
            # self.request is the TCP socket connected to the client
            while True:
                self.data = self.request.recv(1024).strip()
                self.log.debug('data={!r}'.format(self.data))
                if len(self.data) == 0:
                    self.log.info('client disconnected, resuming membalance')
                    if got_lock:
                        global force_refresh_domain_list
                        force_refresh_domain_list = True
                    return

                # XXX something is wrong here: return without release?
                if got_lock:
                    self.log.warning('Second request over qmemman.sock?')
                    return

                self.log.debug('acquiring global_lock')
                global_lock.acquire()
                self.log.debug('global_lock acquired')

                got_lock = True
                if system_state.do_balloon(int(self.data.decode('ascii'))):
                    resp = b"OK\n"
                else:
                    resp = b"FAIL\n"
                self.log.debug('resp={!r}'.format(resp))
                self.request.send(resp)
        except BaseException as e:
            self.log.exception(
                "exception while handling request: {!r}".format(e))
        finally:
            if got_lock:
                global_lock.release()
                self.log.debug('global_lock released')

parser = qubes.tools.QubesArgumentParser(want_app=False)

parser.add_argument('--config', '-c', metavar='FILE',
    action='store', default='/etc/qubes/qmemman.conf',
    help='qmemman config file')

parser.add_argument('--foreground',
    action='store_true', default=False,
    help='do not close stdio')

def main():
    args = parser.parse_args()

    # setup logging
    ha_syslog = logging.handlers.SysLogHandler('/dev/log')
    ha_syslog.setFormatter(
        logging.Formatter('%(name)s[%(process)d]: %(message)s'))
    logging.root.addHandler(ha_syslog)

    if args.foreground:
        ha_stderr = logging.StreamHandler(sys.stderr)
        ha_file.setFormatter(
            logging.Formatter('%(asctime)s %(name)s[%(process)d]: %(message)s'))
        logging.root.addHandler(ha_stderr)

    sys.stdin.close()

    log = logging.getLogger('qmemman.daemon')

    config = configparser.SafeConfigParser({
            'vm-min-mem': str(qubes.qmemman.algo.MIN_PREFMEM),
            'dom0-mem-boost': str(qubes.qmemman.algo.DOM0_MEM_BOOST),
            'cache-margin-factor': str(qubes.qmemman.algo.CACHE_FACTOR)
            })
    config.read(args.config)

    if config.has_section('global'):
        qubes.qmemman.algo.MIN_PREFMEM = \
            qubes.utils.parse_size(config.get('global', 'vm-min-mem'))
        qubes.qmemman.algo.DOM0_MEM_BOOST = \
            qubes.utils.parse_size(config.get('global', 'dom0-mem-boost'))
        qubes.qmemman.algo.CACHE_FACTOR = \
            config.getfloat('global', 'cache-margin-factor')
        loglevel = config.getint('global', 'log-level')
        logging.root.setLevel(loglevel)

    log.info('MIN_PREFMEM={algo.MIN_PREFMEM}'
        ' DOM0_MEM_BOOST={algo.DOM0_MEM_BOOST}'
        ' CACHE_FACTOR={algo.CACHE_FACTOR}'.format(
            algo=qubes.qmemman.algo))

    try:
        os.unlink(SOCK_PATH)
    except FileNotFoundError:
        pass

    log.debug('instantiating server')
    os.umask(0)

    # Initialize the connection to Xen and to XenStore
    system_state.init()

    server = socketserver.UnixStreamServer(SOCK_PATH, QMemmanReqHandler)
    os.umask(0o077)

    # notify systemd
    nofity_socket = os.getenv('NOTIFY_SOCKET')
    if nofity_socket:
        log.debug('notifying systemd')
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
        if nofity_socket.startswith('@'):
            nofity_socket = '\0%s' % nofity_socket[1:]
        sock.connect(nofity_socket)
        sock.sendall(b"READY=1")
        sock.close()

    threading.Thread(target=server.serve_forever).start()
    XSWatcher().watch_loop()
<|endoftext|>


<commit_before># coding: utf-8
from ui.mainwindow import Ui_MainWindow
from feedreader import FeedReader
from PyQt4 import QtGui, QtCore
from . import model
from .settings import SettingsDialog

class MainWindow(Ui_MainWindow):
    def setupUi(self, window):
        Ui_MainWindow.setupUi(self, window)
        header = QtGui.QHeaderView(QtCore.Qt.Horizontal)
        header.setResizeMode(QtGui.QHeaderView.Stretch)
        self.episodeTable.setHorizontalHeader(header)
        self.fetchFeedButton.clicked.connect(self._load_feed)
        self.downloadButton.clicked.connect(self._download)
        self.deleteButton.clicked.connect(self._delete_row)
        self.saveQueueButton.clicked.connect(self._save_queue)
        self.settings = SettingsDialog()
        self.settingsButton.clicked.connect(self.settings.show)

        QtGui.qApp.queuehandler.finishedLoading.connect(self._set_model)
        QtGui.qApp.queuehandler.load()

        QtGui.qApp.downloader.fileDownloaded.connect(self._downloaded)
        QtGui.qApp.downloader.notFound.connect(self._notfound)

    def _download(self):
        rows = sorted(set([item.row() for item in
            self.episodeTable.selectedIndexes()]))
        l_model = self.episodeTable.model()

        items = {}

        if len(rows) > 0:
            for row in rows:
                items[l_model.item(row, model.COLUMN_CRC).text()] = \
                        l_model.item(row, model.COLUMN_ANIME).text()
        else:
            for i in xrange(0, l_model.rowCount()):
                items[l_model.item(i, model.COLUMN_CRC).text()] = \
                        l_model.item(i, model.COLUMN_ANIME).text()

        self.downloader.enqueue(items)

    def _downloaded(self, crc):
        l_model = self.episodeTable.model()
        for i in xrange(0, l_model.rowCount()):
                if crc == l_model.data(l_model.index(i, model.COLUMN_CRC)).toString():
                    l_model.removeRows(i, 1)

    def _notfound(self, crc):
        l_model = self.episodeTable.model()
        for i in xrange(0, l_model.rowCount()):
            if crc == l_model.data(l_model.index(i, model.COLUMN_CRC)).toString():
                for k in xrange(0, model.COLUMNS):
                    l_model.setData(l_model.index(i,k),
                            QtCore.QVariant(QtGui.QBrush(QtCore.Qt.red)),
                            QtCore.Qt.BackgroundRole)

    def _load_feed(self):
        loader = FeedReader()
        loader.finishedLoading.connect(self._dicts_to_tree)
        loader.load()

    def _dicts_to_tree(self, dicts):
    # TODO move into the feedreader
        d_model = model.make_model(len(dicts))
        for index, item in enumerate(dicts):
            d_model.setData(d_model.index(index, 0), QtCore.QVariant(item["anime"]))
            if "episode" in item:
                d_model.setData(d_model.index(index, 1),
                        QtCore.QVariant(item["episode"]))
            if "group" in item:
                d_model.setData(d_model.index(index, 3),
                        QtCore.QVariant(item["group"]))
            d_model.setData(d_model.index(index, 2), QtCore.QVariant(item["crc"]))
        self._set_model(d_model)

    def _set_model(self, new_model):
        new_model.sort(model.COLUMN_ANIME)
        self.episodeTable.setModel(new_model)
        self.episodeTable.resizeColumnsToContents()

    def _delete_row(self):
        rows = sorted(set([item.row() for item in
            self.episodeTable.selectedIndexes()]))
        model = self.episodeTable.model()
        for index, row in enumerate(rows):
            model.removeRows(row - index, 1)

    def _save_queue(self):
        model = self.episodeTable.model()
        self.queuehandler.save(model)
<commit_msg>use QtGui.qApp some more<commit_after># coding: utf-8
from ui.mainwindow import Ui_MainWindow
from feedreader import FeedReader
from PyQt4 import QtGui, QtCore
from . import model
from .settings import SettingsDialog

class MainWindow(Ui_MainWindow):
    def setupUi(self, window):
        Ui_MainWindow.setupUi(self, window)
        header = QtGui.QHeaderView(QtCore.Qt.Horizontal)
        header.setResizeMode(QtGui.QHeaderView.Stretch)
        self.episodeTable.setHorizontalHeader(header)
        self.fetchFeedButton.clicked.connect(self._load_feed)
        self.downloadButton.clicked.connect(self._download)
        self.deleteButton.clicked.connect(self._delete_row)
        self.saveQueueButton.clicked.connect(self._save_queue)
        self.settings = SettingsDialog()
        self.settingsButton.clicked.connect(self.settings.show)

        QtGui.qApp.queuehandler.finishedLoading.connect(self._set_model)
        QtGui.qApp.queuehandler.load()

        QtGui.qApp.downloader.fileDownloaded.connect(self._downloaded)
        QtGui.qApp.downloader.notFound.connect(self._notfound)

    def _download(self):
        rows = sorted(set([item.row() for item in
            self.episodeTable.selectedIndexes()]))
        l_model = self.episodeTable.model()

        items = {}

        if len(rows) > 0:
            for row in rows:
                items[l_model.item(row, model.COLUMN_CRC).text()] = \
                        l_model.item(row, model.COLUMN_ANIME).text()
        else:
            for i in xrange(0, l_model.rowCount()):
                items[l_model.item(i, model.COLUMN_CRC).text()] = \
                        l_model.item(i, model.COLUMN_ANIME).text()

        QtGui.qApp.downloader.enqueue(items)

    def _downloaded(self, crc):
        l_model = self.episodeTable.model()
        for i in xrange(0, l_model.rowCount()):
                if crc == l_model.data(l_model.index(i, model.COLUMN_CRC)).toString():
                    l_model.removeRows(i, 1)

    def _notfound(self, crc):
        l_model = self.episodeTable.model()
        for i in xrange(0, l_model.rowCount()):
            if crc == l_model.data(l_model.index(i, model.COLUMN_CRC)).toString():
                for k in xrange(0, model.COLUMNS):
                    l_model.setData(l_model.index(i,k),
                            QtCore.QVariant(QtGui.QBrush(QtCore.Qt.red)),
                            QtCore.Qt.BackgroundRole)

    def _load_feed(self):
        loader = FeedReader()
        loader.finishedLoading.connect(self._dicts_to_tree)
        loader.load()

    def _dicts_to_tree(self, dicts):
    # TODO move into the feedreader
        d_model = model.make_model(len(dicts))
        for index, item in enumerate(dicts):
            d_model.setData(d_model.index(index, 0), QtCore.QVariant(item["anime"]))
            if "episode" in item:
                d_model.setData(d_model.index(index, 1),
                        QtCore.QVariant(item["episode"]))
            if "group" in item:
                d_model.setData(d_model.index(index, 3),
                        QtCore.QVariant(item["group"]))
            d_model.setData(d_model.index(index, 2), QtCore.QVariant(item["crc"]))
        self._set_model(d_model)

    def _set_model(self, new_model):
        new_model.sort(model.COLUMN_ANIME)
        self.episodeTable.setModel(new_model)
        self.episodeTable.resizeColumnsToContents()

    def _delete_row(self):
        rows = sorted(set([item.row() for item in
            self.episodeTable.selectedIndexes()]))
        model = self.episodeTable.model()
        for index, row in enumerate(rows):
            model.removeRows(row - index, 1)

    def _save_queue(self):
        model = self.episodeTable.model()
        QtGui.qApp.queuehandler.save(model)
<|endoftext|>


<commit_before>from dataloader.model import LanguageModel
from dataloader.sim_regression import SimRegression
from typing import List

class QueryExpander:
    def __init__(self, model: LanguageModel):
        self.model = model
    def ExpandQuery(self, query: List[str], topn = 1):
        expanded = list(query)
        for term in query:
            most_similars = self.model.find_most_similar_words(term, topn)
            for similar, score in most_similars:
                if similar not in expanded:
                    expanded.append(similar)
        return expanded
<commit_msg>wzory 7 i 8, cosinus-sigmoid<commit_after>from dataloader.model import LanguageModel
from dataloader.sim_regression import SimRegression
from typing import List

class QueryExpander:
    def __init__(self, model: LanguageModel):
        self.model = model

    def p7(self, word1: str, word2: str, topn: int):
        similars = self.model.find_most_similar_words(word2, topn)
        numerator = self.model.similarity(word1, word2)
        denominator = 0
        for word3, score in similars:
            denominator += self.model.similarity(word3, word2)
        if denominator != 0:
            return numerator/denominator

    def p8(self, word: str, query: List[str], topn: int):
        prob = 0
        query_length = len(query)
        for term in query:
            term_count = 0
            for term2 in query:
                if term2 == term:
                    term_count += 1
            prob += self.p7(word, term, topn) * term_count/query_length
        return prob

    def ExpandQuery(self, query: List[str], limit = 20):
        expanded = list(query)
        tmp = dict()
        for term in query:
            tmp[term] = self.p8(term, query, limit)
        #return expanded
        return tmp<|endoftext|>


<commit_before>import json
import redis
import time

from queue import Queue, Empty
from threading import Thread
from contextlib import contextmanager
from redis.exceptions import ConnectionError as RedisConnectionError

class ConnectionError(Exception):
    '''
    Represents connection errors that can occur when talking to the ArchiveBot
    control node, Redis or otherwise.
    '''

    pass

@contextmanager
def conn(controller):
    try:
        if not controller.connected():
            controller.connect()
        yield
    except RedisConnectionError as e:
        controller.disconnect()
        raise ConnectionError(str(e)) from e

def candidate_queues(named_queues, pipeline_nick, ao_only):
    '''
    Generates names of queues that this pipeline will check for work.
    '''

    def applies(q):
        return q.replace('pending:', '') in pipeline_nick

    if ao_only:
        return ['pending-ao']
    else:
        matches = [q for q in named_queues if applies(q)]
        matches.append('pending-ao')
        matches.append('pending')

        return matches

class Control(object):
    '''
    Handles communication to and from the ArchiveBot control server.

    If a message cannot be processed due to a connection error, the Redis
    connection is closed and deleted.  An archivebot.control.ConnectionError
    is also raised.
    '''

    def __init__(self, redis_url, log_channel, pipeline_channel):
        self.log_channel = log_channel
        self.pipeline_channel = pipeline_channel
        self.items_downloaded_outstanding = 0
        self.items_queued_outstanding = 0
        self.redis_url = redis_url
        self.log_queue = Queue()
        self.bytes_downloaded_queue = Queue()
        self.item_count_queue = Queue()

        self.connect()

        #log_thread is joined in finish_logging()
        self.ending = False
        self.log_thread = Thread(target=self.ship_logs)
        self.log_thread.start()

    def connected(self):
        return self.redis is not None

    def connect(self):
        if self.redis_url is None:
            raise ConnectionError('self.redis_url not set')

        self.redis = redis.StrictRedis.from_url(self.redis_url,
                                                decode_responses=True)

        self.register_scripts()

    def disconnect(self):
        self.redis = None

    def register_scripts(self):
        self.mark_done_script = self.redis.register_script(MARK_DONE_SCRIPT)
        self.mark_aborted_script = self.redis.register_script(MARK_ABORTED_SCRIPT)
        self.log_script = self.redis.register_script(LOGGER_SCRIPT)

    def all_named_pending_queues(self):
        with conn(self):
            pipelines = set()

            for name in self.redis.scan_iter('pending:*'):
                pipelines.add(name)

            return pipelines

    def reserve_job(self, pipeline_id, pipeline_nick, ao_only):
        named_queues = self.all_named_pending_queues()

        for queue in candidate_queues(named_queues, pipeline_nick, ao_only):
            ident = self.dequeue_item(queue)

            if ident:
                return self.complete_reservation(ident, pipeline_id)

        return None, None

    def dequeue_item(self, queue):
        with conn(self):
            return self.redis.rpoplpush(queue, 'working')

    def complete_reservation(self, ident, pipeline_id):
        with conn(self):
            self.redis.hmset(ident, dict(
                started_at=time.time(),
                pipeline_id=pipeline_id
            ))

            return ident, self.redis.hgetall(ident)

    def heartbeat(self, ident):
        try:
            with conn(self):
                self.redis.hincrby(ident, 'heartbeat', 1)
        except ConnectionError:
            pass

    def is_aborted(self, ident):
        with conn(self):
            return self.redis.hget(ident, 'aborted')

    def mark_done(self, item, expire_time):
        with conn(self):
            self.mark_done_script(keys=[item['ident']], args=[expire_time,
                self.log_channel, int(time.time()), json.dumps(item['info']),
                                                              item['log_key']])

    def mark_aborted(self, ident):
        with conn(self):
            self.mark_aborted_script(keys=[ident], args=[self.log_channel])

    def update_bytes_downloaded(self, ident, size):
        self.bytes_downloaded_queue.put({'ident': ident,
                                         'bytes': size
                                        })

    def update_items_downloaded(self, count):
        self.items_downloaded_outstanding += count

    def update_items_queued(self, count):
        self.items_queued_outstanding += count

    def flush_item_counts(self, ident):
        self.item_count_queue.put({'ident': ident,
                                   'items_downloaded':
                                       self.items_downloaded_outstanding,
                                   'items_queued': self.items_queued_outstanding
                                  })
        self.items_downloaded_outstanding = 0
        self.items_queued_outstanding = 0

    def pipeline_report(self, pipeline_id, report):
        try:
            with conn(self):
                self.redis.hmset(pipeline_id, report)
                self.redis.sadd('pipelines', pipeline_id)
                self.redis.publish(self.pipeline_channel, pipeline_id)
        except ConnectionError:
            pass

    def finish_logging(self):
        self.ending = True
        self.log_thread.join()

    def unregister_pipeline(self, pipeline_id):
        try:
            with conn(self):
                self.redis.delete(pipeline_id)
                self.redis.srem('pipelines', pipeline_id)
                self.redis.publish(self.pipeline_channel, pipeline_id)
        except ConnectionError:
            pass

    # This function is a thread used to asynchronously ship logs to redis for
    # this job, in a daemonic thread
    def ship_logs(self):
        bytes_entries = {}
        counts_entries = {}
        shipping_count = 0

        with self.redis.pipeline(transaction=False) as pipe:
            while not (self.ending and
                       self.log_queue.empty() and
                       self.bytes_downloaded_queue.empty() and
                       self.item_count_queue.empty()):
                try:
                    # Ship a log entry
                    try:
                        entry = self.log_queue.get(timeout=5)
                        with conn(self):
                            self.log_script(keys=entry['keys'], args=entry['args'], client=pipe)
                        self.log_queue.task_done()
                    except Empty:
                        pass
                    except ConnectionError as exception: # If we can't ship the log entry, discard
                        self.log_queue.task_done()
                        raise exception

                    # Aggregate counts to ship when logs do
                    try:
                        entry = self.bytes_downloaded_queue.get(block=False)
                        if entry['ident'] in bytes_entries:
                            bytes_entries[entry['ident']] += int(entry['bytes'])
                        else:
                            bytes_entries[entry['ident']] = int(entry['bytes'])
                        self.bytes_downloaded_queue.task_done()
                    except Empty:
                        pass

                    try:
                        entry = self.item_count_queue.get(block=False)
                        if entry['ident'] in counts_entries:
                            counts_entries[entry['ident']][0] += int(entry['items_downloaded'])
                            counts_entries[entry['ident']][1] += int(entry['items_queued'])
                        else:
                            counts_entries[entry['ident']] = [
                                int(entry['items_downloaded']),
                                int(entry['items_queued'])
                            ]
                        self.item_count_queue.task_done()
                    except Empty:
                        pass

                    # If we have accreted enough or the queue is empty, commit logs and counts
                    # The magic constant against which to compare shipping_count should be
                    # selected such that about that many log entries might be shipped every
                    # round-trip to the dashboard, under congested conditions
                    if self.log_queue.empty() or shipping_count >= 64:
                        for ident, count in bytes_entries:
                            pipe.hincrby(ident, 'bytes_downloaded', count)
                            bytes_entries = {}

                        for ident, data in counts_entries:
                            pipe.hincrby(ident, 'items_downloaded', data[0])
                            pipe.hincrby(ident, 'items_queued', data[1])
                            counts_entries = {}

                        pipe.execute()
                        shipping_count = 0

                except ConnectionError:
                    pass

    def log(self, packet, ident, log_key):
        self.log_queue.put({'keys': [ident],
                            'args': [json.dumps(packet), self.log_channel, log_key]
                           })

    def get_url_file(self, ident):
        try:
            with conn(self):
                return self.redis.hget(ident, 'url_file')
        except ConnectionError:
            pass

    def get_settings(self, ident):
        with conn(self):
            data = self.redis.hmget(ident, 'delay_min', 'delay_max',
                                    'concurrency',
                                    'settings_age',
                                    'abort_requested',
                                    'suppress_ignore_reports',
                                    'ignore_patterns_set_key')

            result = dict(
                delay_min=data[0],
                delay_max=data[1],
                concurrency=data[2],
                age=data[3],
                abort_requested=data[4],
                suppress_ignore_reports=data[5]
                )

            if data[6]:
                result['ignore_patterns'] = self.redis.smembers(data[6])
            else:
                result['ignore_patterns'] = []

            return result

# ------------------------------------------------------------------------------

MARK_DONE_SCRIPT = '''
local ident = KEYS[1]
local expire_time = ARGV[1]
local log_channel = ARGV[2]
local finished_at = ARGV[3]
local info = ARGV[4]
local log_key = ARGV[5]

redis.call('hmset', ident, 'finished_at', finished_at)
redis.call('lrem', 'working', 1, ident)

local was_aborted = redis.call('hget', ident, 'aborted')

-- If the job was aborted, we ignore the given expire time.  Instead, we set a
-- much shorter expire time -- one that's long enough for (most) subscribers
-- to read a message, but short enough to not cause undue suffering in the
-- case of retrying an aborted job.
if was_aborted then
    redis.call('incr', 'jobs_aborted')
    redis.call('expire', ident, 5)
    redis.call('expire', log_key, 5)
    redis.call('expire', ident..'_ignores', 5)
else
    redis.call('incr', 'jobs_completed')
    redis.call('expire', ident, expire_time)
    redis.call('expire', log_key, expire_time)
    redis.call('expire', ident..'_ignores', expire_time)
end

redis.call('rpush', 'finish_notifications', info)
redis.call('publish', log_channel, ident)
'''

MARK_ABORTED_SCRIPT = '''
local ident = KEYS[1]
local log_channel = ARGV[1]

redis.call('hset', ident, 'aborted', 'true')
redis.call('publish', log_channel, ident)
'''

LOGGER_SCRIPT = '''
local ident = KEYS[1]
local message = ARGV[1]
local log_channel = ARGV[2]
local log_key = ARGV[3]

local nextseq = redis.call('hincrby', ident, 'log_score', 1)

redis.call('zadd', log_key, nextseq, message)
redis.call('publish', log_channel, ident)
'''

# vim:ts=4:sw=4:et:tw=78
<commit_msg>Correct dict iteration error<commit_after>import json
import redis
import time

from queue import Queue, Empty
from threading import Thread
from contextlib import contextmanager
from redis.exceptions import ConnectionError as RedisConnectionError

class ConnectionError(Exception):
    '''
    Represents connection errors that can occur when talking to the ArchiveBot
    control node, Redis or otherwise.
    '''

    pass

@contextmanager
def conn(controller):
    try:
        if not controller.connected():
            controller.connect()
        yield
    except RedisConnectionError as e:
        controller.disconnect()
        raise ConnectionError(str(e)) from e

def candidate_queues(named_queues, pipeline_nick, ao_only):
    '''
    Generates names of queues that this pipeline will check for work.
    '''

    def applies(q):
        return q.replace('pending:', '') in pipeline_nick

    if ao_only:
        return ['pending-ao']
    else:
        matches = [q for q in named_queues if applies(q)]
        matches.append('pending-ao')
        matches.append('pending')

        return matches

class Control(object):
    '''
    Handles communication to and from the ArchiveBot control server.

    If a message cannot be processed due to a connection error, the Redis
    connection is closed and deleted.  An archivebot.control.ConnectionError
    is also raised.
    '''

    def __init__(self, redis_url, log_channel, pipeline_channel):
        self.log_channel = log_channel
        self.pipeline_channel = pipeline_channel
        self.items_downloaded_outstanding = 0
        self.items_queued_outstanding = 0
        self.redis_url = redis_url
        self.log_queue = Queue()
        self.bytes_downloaded_queue = Queue()
        self.item_count_queue = Queue()

        self.connect()

        #log_thread is joined in finish_logging()
        self.ending = False
        self.log_thread = Thread(target=self.ship_logs)
        self.log_thread.start()

    def connected(self):
        return self.redis is not None

    def connect(self):
        if self.redis_url is None:
            raise ConnectionError('self.redis_url not set')

        self.redis = redis.StrictRedis.from_url(self.redis_url,
                                                decode_responses=True)

        self.register_scripts()

    def disconnect(self):
        self.redis = None

    def register_scripts(self):
        self.mark_done_script = self.redis.register_script(MARK_DONE_SCRIPT)
        self.mark_aborted_script = self.redis.register_script(MARK_ABORTED_SCRIPT)
        self.log_script = self.redis.register_script(LOGGER_SCRIPT)

    def all_named_pending_queues(self):
        with conn(self):
            pipelines = set()

            for name in self.redis.scan_iter('pending:*'):
                pipelines.add(name)

            return pipelines

    def reserve_job(self, pipeline_id, pipeline_nick, ao_only):
        named_queues = self.all_named_pending_queues()

        for queue in candidate_queues(named_queues, pipeline_nick, ao_only):
            ident = self.dequeue_item(queue)

            if ident:
                return self.complete_reservation(ident, pipeline_id)

        return None, None

    def dequeue_item(self, queue):
        with conn(self):
            return self.redis.rpoplpush(queue, 'working')

    def complete_reservation(self, ident, pipeline_id):
        with conn(self):
            self.redis.hmset(ident, dict(
                started_at=time.time(),
                pipeline_id=pipeline_id
            ))

            return ident, self.redis.hgetall(ident)

    def heartbeat(self, ident):
        try:
            with conn(self):
                self.redis.hincrby(ident, 'heartbeat', 1)
        except ConnectionError:
            pass

    def is_aborted(self, ident):
        with conn(self):
            return self.redis.hget(ident, 'aborted')

    def mark_done(self, item, expire_time):
        with conn(self):
            self.mark_done_script(keys=[item['ident']], args=[expire_time,
                self.log_channel, int(time.time()), json.dumps(item['info']),
                                                              item['log_key']])

    def mark_aborted(self, ident):
        with conn(self):
            self.mark_aborted_script(keys=[ident], args=[self.log_channel])

    def update_bytes_downloaded(self, ident, size):
        self.bytes_downloaded_queue.put({'ident': ident,
                                         'bytes': size
                                        })

    def update_items_downloaded(self, count):
        self.items_downloaded_outstanding += count

    def update_items_queued(self, count):
        self.items_queued_outstanding += count

    def flush_item_counts(self, ident):
        self.item_count_queue.put({'ident': ident,
                                   'items_downloaded':
                                       self.items_downloaded_outstanding,
                                   'items_queued': self.items_queued_outstanding
                                  })
        self.items_downloaded_outstanding = 0
        self.items_queued_outstanding = 0

    def pipeline_report(self, pipeline_id, report):
        try:
            with conn(self):
                self.redis.hmset(pipeline_id, report)
                self.redis.sadd('pipelines', pipeline_id)
                self.redis.publish(self.pipeline_channel, pipeline_id)
        except ConnectionError:
            pass

    def finish_logging(self):
        self.ending = True
        self.log_thread.join()

    def unregister_pipeline(self, pipeline_id):
        try:
            with conn(self):
                self.redis.delete(pipeline_id)
                self.redis.srem('pipelines', pipeline_id)
                self.redis.publish(self.pipeline_channel, pipeline_id)
        except ConnectionError:
            pass

    # This function is a thread used to asynchronously ship logs to redis for
    # this job, in a daemonic thread
    def ship_logs(self):
        bytes_entries = {}
        counts_entries = {}
        shipping_count = 0

        with self.redis.pipeline(transaction=False) as pipe:
            while not (self.ending and
                       self.log_queue.empty() and
                       self.bytes_downloaded_queue.empty() and
                       self.item_count_queue.empty()):
                try:
                    # Ship a log entry
                    try:
                        entry = self.log_queue.get(timeout=5)
                        with conn(self):
                            self.log_script(keys=entry['keys'], args=entry['args'], client=pipe)
                        self.log_queue.task_done()
                    except Empty:
                        pass
                    except ConnectionError as exception: # If we can't ship the log entry, discard
                        self.log_queue.task_done()
                        raise exception

                    # Aggregate counts to ship when logs do
                    try:
                        entry = self.bytes_downloaded_queue.get(block=False)
                        if entry['ident'] in bytes_entries:
                            bytes_entries[entry['ident']] += int(entry['bytes'])
                        else:
                            bytes_entries[entry['ident']] = int(entry['bytes'])
                        self.bytes_downloaded_queue.task_done()
                    except Empty:
                        pass

                    try:
                        entry = self.item_count_queue.get(block=False)
                        if entry['ident'] in counts_entries:
                            counts_entries[entry['ident']][0] += int(entry['items_downloaded'])
                            counts_entries[entry['ident']][1] += int(entry['items_queued'])
                        else:
                            counts_entries[entry['ident']] = [
                                int(entry['items_downloaded']),
                                int(entry['items_queued'])
                            ]
                        self.item_count_queue.task_done()
                    except Empty:
                        pass

                    # If we have accreted enough or the queue is empty, commit logs and counts
                    # The magic constant against which to compare shipping_count should be
                    # selected such that about that many log entries might be shipped every
                    # round-trip to the dashboard, under congested conditions
                    if self.log_queue.empty() or shipping_count >= 64:
                        for ident, count in bytes_entries.iteritems():
                            pipe.hincrby(ident, 'bytes_downloaded', count)
                            bytes_entries = {}

                        for ident, data in counts_entries.iteritems():
                            pipe.hincrby(ident, 'items_downloaded', data[0])
                            pipe.hincrby(ident, 'items_queued', data[1])
                            counts_entries = {}

                        pipe.execute()
                        shipping_count = 0

                except ConnectionError:
                    pass

    def log(self, packet, ident, log_key):
        self.log_queue.put({'keys': [ident],
                            'args': [json.dumps(packet), self.log_channel, log_key]
                           })

    def get_url_file(self, ident):
        try:
            with conn(self):
                return self.redis.hget(ident, 'url_file')
        except ConnectionError:
            pass

    def get_settings(self, ident):
        with conn(self):
            data = self.redis.hmget(ident, 'delay_min', 'delay_max',
                                    'concurrency',
                                    'settings_age',
                                    'abort_requested',
                                    'suppress_ignore_reports',
                                    'ignore_patterns_set_key')

            result = dict(
                delay_min=data[0],
                delay_max=data[1],
                concurrency=data[2],
                age=data[3],
                abort_requested=data[4],
                suppress_ignore_reports=data[5]
                )

            if data[6]:
                result['ignore_patterns'] = self.redis.smembers(data[6])
            else:
                result['ignore_patterns'] = []

            return result

# ------------------------------------------------------------------------------

MARK_DONE_SCRIPT = '''
local ident = KEYS[1]
local expire_time = ARGV[1]
local log_channel = ARGV[2]
local finished_at = ARGV[3]
local info = ARGV[4]
local log_key = ARGV[5]

redis.call('hmset', ident, 'finished_at', finished_at)
redis.call('lrem', 'working', 1, ident)

local was_aborted = redis.call('hget', ident, 'aborted')

-- If the job was aborted, we ignore the given expire time.  Instead, we set a
-- much shorter expire time -- one that's long enough for (most) subscribers
-- to read a message, but short enough to not cause undue suffering in the
-- case of retrying an aborted job.
if was_aborted then
    redis.call('incr', 'jobs_aborted')
    redis.call('expire', ident, 5)
    redis.call('expire', log_key, 5)
    redis.call('expire', ident..'_ignores', 5)
else
    redis.call('incr', 'jobs_completed')
    redis.call('expire', ident, expire_time)
    redis.call('expire', log_key, expire_time)
    redis.call('expire', ident..'_ignores', expire_time)
end

redis.call('rpush', 'finish_notifications', info)
redis.call('publish', log_channel, ident)
'''

MARK_ABORTED_SCRIPT = '''
local ident = KEYS[1]
local log_channel = ARGV[1]

redis.call('hset', ident, 'aborted', 'true')
redis.call('publish', log_channel, ident)
'''

LOGGER_SCRIPT = '''
local ident = KEYS[1]
local message = ARGV[1]
local log_channel = ARGV[2]
local log_key = ARGV[3]

local nextseq = redis.call('hincrby', ident, 'log_score', 1)

redis.call('zadd', log_key, nextseq, message)
redis.call('publish', log_channel, ident)
'''

# vim:ts=4:sw=4:et:tw=78
<|endoftext|>