~free.ekanayaka/landscape-client/jaunty-proposed

import thread
import time
import sys
import logging
import bisect
import socket

import gobject

from twisted.test.proto_helpers import FakeDatagramTransport
from twisted.internet.defer import succeed, fail
from twisted.internet.error import DNSLookupError

from landscape.log import format_object


class InvalidID(Exception):
    """Raised when an invalid ID is used with reactor.cancel_call()."""


class CallHookError(Exception):
    """Raised when hooking on a reactor incorrectly."""


class EventID(object):
    """Unique identifier for an event handler."""

    def __init__(self, event_type, pair):
        """
        @param event_type: Name of the event type handled by the handler.
        @param pair: Binary tuple C{(handler, priority)} holding the handler
            function and its priority.
        """
        self._event_type = event_type
        self._pair = pair


class EventHandlingReactorMixin(object):
    """Fire events identified by strings and register handlers for them."""

    def __init__(self):
        super(EventHandlingReactorMixin, self).__init__()
        self._event_handlers = {}

    def call_on(self, event_type, handler, priority=0):
        """Register an event handler.

        @param event_type: The name of the event type to handle.
        @param handler: The function handling the given event type.
        @param priority: The priority of the given handler function.

        @return: The L{EventID} of the registered handler.
        """
        pair = (handler, priority)

        handlers = self._event_handlers.setdefault(event_type, [])
        handlers.append(pair)
        handlers.sort(key=lambda pair: pair[1])

        return EventID(event_type, pair)

    def fire(self, event_type, *args, **kwargs):
        """Fire an event of a given type.

        Call all handlers registered for the given C{event_type}, in order
        of priority.

        @param event_type: The name of the event type to fire.
        @param args: Positional arguments to pass to the registered handlers.
        @param kwargs: Keyword arguments to pass to the registered handlers.
        """
        logging.debug("Started firing %s.", event_type)
        results = []
        for handler, priority in self._event_handlers.get(event_type, ()):
            try:
                logging.debug("Calling %s for %s with priority %d.",
                              format_object(handler), event_type, priority)
                results.append(handler(*args, **kwargs))
            except KeyboardInterrupt:
                logging.exception("Keyboard interrupt while running event "
                                  "handler %s for event type %r with "
                                  "args %r %r.", format_object(handler),
                                  event_type, args, kwargs)
                self.stop()
                raise
            except:
                logging.exception("Error running event handler %s for "
                                  "event type %r with args %r %r.",
                                  format_object(handler), event_type,
                                  args, kwargs)
        logging.debug("Finished firing %s.", event_type)
        return results

    def cancel_call(self, id):
        """Unregister an event handler.

        @param id: the L{EventID} of the handler to unregister.
        """
        if type(id) is EventID:
            self._event_handlers[id._event_type].remove(id._pair)
        else:
            raise InvalidID("EventID instance expected, received %r" % id)


class ThreadedCallsReactorMixin(object):
    """Schedule functions for execution in the main thread or in new ones."""

    def __init__(self):
        super(ThreadedCallsReactorMixin, self).__init__()
        self._threaded_callbacks = []

    def call_in_main(self, f, *args, **kwargs):
        """Schedule a function for execution in the main thread."""
        self._threaded_callbacks.append(lambda: f(*args, **kwargs))

    def call_in_thread(self, callback, errback, f, *args, **kwargs):
        """
        Execute a callable object in a new separate thread.

        @param callback: A function to call in case C{f} was successful, it
            will be passed the return value of C{f}.
        @param errback: A function to call in case C{f} raised an exception,
            it will be pass a C{(type, value, traceback)} tuple giving
            information about the raised exception (see L{sys.exc_info}).

        @note: Both C{callback} and C{errback} will be executed in the
            the parent thread.
        """
        thread.start_new_thread(self._in_thread,
                                (callback, errback, f, args, kwargs))

    def _in_thread(self, callback, errback, f, args, kwargs):
        try:
            result = f(*args, **kwargs)
        except Exception, e:
            exc_info = sys.exc_info()
            if errback is None:
                self.call_in_main(logging.error, e, exc_info=exc_info)
            else:
                self.call_in_main(errback, *exc_info)
        else:
            if callback:
                self.call_in_main(callback, result)

    def _run_threaded_callbacks(self):
        while self._threaded_callbacks:
            try:
                self._threaded_callbacks.pop(0)()
            except Exception, e:
                logging.exception(e)

    def _hook_threaded_callbacks(self):
        id = self.call_every(0.5, self._run_threaded_callbacks)
        self._run_threaded_callbacks_id = id

    def _unhook_threaded_callbacks(self):
        self.cancel_call(self._run_threaded_callbacks_id)


class ReactorID(object):

    def __init__(self, timeout):
        self._timeout = timeout


class Reactor(EventHandlingReactorMixin,
              ThreadedCallsReactorMixin):

    def __init__(self):
        super(Reactor, self).__init__()
        self._context = gobject.MainContext()
        self._mainloop = gobject.MainLoop(context=self._context)

    def call_later(self, timeout, function, *args, **kwargs):
        def fake_function():
            function(*args, **kwargs)
            return False
        timeout = gobject.Timeout(int(timeout*1000))
        timeout.set_callback(fake_function)
        timeout.attach(self._context)
        return ReactorID(timeout)

    def cancel_call(self, id):
        if type(id) is ReactorID:
            id._timeout.destroy()
        else:
            super(Reactor, self).cancel_call(id)

    def call_every(self, timeout, function, *args, **kwargs):
        def fake_function():
            function(*args, **kwargs)
            return True
        timeout = gobject.Timeout(int(timeout*1000))
        timeout.set_callback(fake_function)
        timeout.attach(self._context)
        return ReactorID(timeout)

    def run(self):
        self.fire("run")
        self._hook_threaded_callbacks()
        self._mainloop.run()
        self._unhook_threaded_callbacks()
        self.fire("stop")

    def stop(self):
        self._mainloop.quit()


class FakeReactorID(object):

    def __init__(self, data):
        self.active = True
        self._data = data


class FakeReactor(EventHandlingReactorMixin,
                  ThreadedCallsReactorMixin):
    """
    @ivar udp_transports: dict of {port: (protocol, transport)}
    @ivar hosts: Dict of {hostname: ip}. Users should populate this
        and L{resolve} will use it.
    """
    def __init__(self):
        super(FakeReactor, self).__init__()
        self._current_time = 0
        self._calls = []
        self.udp_transports = {}
        self.hosts = {}

    def time(self):
        return float(self._current_time)

    def call_later(self, seconds, f, *args, **kwargs):
        scheduled_time = self._current_time + seconds
        call = (scheduled_time, f, args, kwargs)
        bisect.insort_left(self._calls, call)
        return FakeReactorID(call)

    def cancel_call(self, id):
        if type(id) is FakeReactorID:
            if id._data in self._calls:
                self._calls.remove(id._data)
            id.active = False
        else:
            super(FakeReactor, self).cancel_call(id)

    def call_every(self, seconds, f, *args, **kwargs):
        def fake():
            # update the call so that cancellation will continue
            # working with the same ID. And do it *before* the call
            # because the call might cancel it!
            call._data = self.call_later(seconds, fake)._data
            try:
                f(*args, **kwargs)
            except:
                if call.active:
                    self.cancel_call(call)
                raise
        call = self.call_later(seconds, fake)
        return call

    def call_in_thread(self, callback, errback, f, *args, **kwargs):
        self._in_thread(callback, errback, f, args, kwargs)

        # Running threaded callbacks here doesn't reflect reality, since
        # they're usually run while the main reactor loop is active.
        # At the same time, this is convenient as it means we don't need
        # to run the the reactor with all registered handlers to test for
        # actions performed on completion of specific events (e.g. firing
        # exchange will fire exchange-done when ready). IOW, it's easier
        # to test things synchronously.
        self._run_threaded_callbacks()

    def advance(self, seconds):
        """Advance this reactor C{seconds} into the future.

        This is the preferred method for advancing time in your unit tests.
        """
        while (self._calls and self._calls[0][0]
               <= self._current_time + seconds):
            call = self._calls.pop(0)
            # If we find a call within the time we're advancing,
            # before calling it, let's advance the time *just* to
            # when that call is expecting to be run, so that if it
            # schedules any calls itself they will be relative to
            # the correct time.
            seconds -= call[0] - self._current_time
            self._current_time = call[0]
            try:
                call[1](*call[2], **call[3])
            except Exception, e:
                logging.exception(e)
        self._current_time += seconds

    def run(self):
        """Continuously advance this reactor until reactor.stop() is called."""
        self.fire("run")
        self._running = True
        while self._running:
            self.advance(self._calls[0][0])
        self.fire("stop")

    def stop(self):
        self._running = False

    def listen_udp(self, port, protocol):
        """
        Connect the given protocol with a fake transport, and keep the
        transport in C{self.udp_transports}.
        """
        transport = FakeDatagramTransport()
        self.udp_transports[port] = (protocol, transport)
        protocol.makeConnection(transport)


    def resolve(self, hostname):
        """Look up the hostname in C{self.hosts}.

        @return: A Deferred resulting in the IP address.
        """
        try:
            # is it an IP address?
            socket.inet_aton(hostname)
        except socket.error: # no
            if hostname in self.hosts:
                return succeed(self.hosts[hostname])
            else:
                return fail(DNSLookupError(hostname))
        else: # yes
            return succeed(hostname)



class TwistedReactor(EventHandlingReactorMixin,
                     ThreadedCallsReactorMixin):
    """Wrap and add functionalities to the Twisted C{reactor}."""

    def __init__(self):
        from twisted.internet import reactor
        from twisted.internet.task import LoopingCall
        self._LoopingCall = LoopingCall
        self._reactor = reactor
        self._cleanup()

        super(TwistedReactor, self).__init__()

    def _cleanup(self):
        # Since the reactor is global, we should clean it up when we
        # initialize one of our wrappers.
        for call in self._reactor.getDelayedCalls():
            if call.active():
                call.cancel()

    def call_later(self, *args, **kwargs):
        """Call a function later.

        Simply call C{callLater(*args, **kwargs)} and return its result.

        @see: L{twisted.internet.interfaces.IReactorTime.callLater}.

        """
        return self._reactor.callLater(*args, **kwargs)

    def call_every(self, seconds, f, *args, **kwargs):
        """Call a function repeatedly.

        Create a new L{twisted.internet.task.LoopingCall} object and
        start it.

        @return: the created C{LoopingCall} object.
        """
        lc = self._LoopingCall(f, *args, **kwargs)
        lc.start(seconds, now=False)
        return lc

    def cancel_call(self, id):
        """Cancel a scheduled function or event handler.

        @param id: The function call or handler to remove. It can be an
            L{EventID}, a L{LoopingCall} or a C{IDelayedCall}, as returned
            by L{call_on}, L{call_every} and L{call_later} respectively.
        """
        if isinstance(id, EventID):
            return EventHandlingReactorMixin.cancel_call(self, id)
        if isinstance(id, self._LoopingCall):
            return id.stop()
        if id.active():
            id.cancel()

    def call_in_main(self, f, *args, **kwargs):
        """Cause a function to be executed by the reactor thread.

        @param f: The callable object to execute.
        @param args: The arguments to call it with.
        @param kwargs: The keyword arguments to call it with.

        @see: L{twisted.internet.interfaces.IReactorThreads.callFromThread}
        """
        self._reactor.callFromThread(f, *args, **kwargs)

    def run(self):
        """Start the reactor, a C{"run"} event will be fired."""

        self.fire("run")
        self._reactor.run()
        self.fire("stop")

    def stop(self):
        """Stop the reactor, a C{"stop"} event will be fired."""

        self._reactor.crash()
        self._cleanup()

    def time(self):
        """Get current time.

        @see L{time.time}
        """
        return time.time()


    def listen_udp(self, port, protocol):
        """Connect the given protocol with a UDP transport.

        @see L{twisted.internet.interfaces.IReactorUDP.listenUDP}.
        """
        return self._reactor.listenUDP(port, protocol)

    def resolve(self, host):
        """Look up the IP of the given host.

        @return: A L{Deferred} resulting in the hostname.

        @see L{twisted.internet.interfaces.IReactorCore.resolve}.

        """
        return self._reactor.resolve(host)