From 7ad7a29a8f08ff815164731eec50ff1cba46b07a Mon Sep 17 00:00:00 2001 From: Nathan Fritz Date: Wed, 13 Oct 2010 18:15:21 -0700 Subject: [PATCH] new state machine in place --- setup.py | 4 +- sleekxmpp/clientxmpp.py | 9 + sleekxmpp/thirdparty/__init__.py | 0 sleekxmpp/thirdparty/statemachine.py | 278 +++++++++++++++++++++++++++ sleekxmpp/xmlstream/scheduler.py | 68 ++++--- sleekxmpp/xmlstream/statemachine.py | 59 ------ sleekxmpp/xmlstream/statemanager.py | 139 -------------- sleekxmpp/xmlstream/xmlstream.py | 261 +++++++++++++------------ 8 files changed, 462 insertions(+), 356 deletions(-) create mode 100644 sleekxmpp/thirdparty/__init__.py create mode 100644 sleekxmpp/thirdparty/statemachine.py delete mode 100644 sleekxmpp/xmlstream/statemachine.py delete mode 100644 sleekxmpp/xmlstream/statemanager.py diff --git a/setup.py b/setup.py index 4c90a1c..f6eace7 100644 --- a/setup.py +++ b/setup.py @@ -44,7 +44,9 @@ packages = [ 'sleekxmpp', 'sleekxmpp/xmlstream', 'sleekxmpp/xmlstream/matcher', 'sleekxmpp/xmlstream/handler', - 'sleekxmpp/xmlstream/tostring'] + 'sleekxmpp/xmlstream/tostring', + 'sleekxmpp/thirdparty', + ] setup( name = "sleekxmpp", diff --git a/sleekxmpp/clientxmpp.py b/sleekxmpp/clientxmpp.py index 13de39d..1aab4a7 100644 --- a/sleekxmpp/clientxmpp.py +++ b/sleekxmpp/clientxmpp.py @@ -92,6 +92,7 @@ class ClientXMPP(BaseXMPP): self.sessionstarted = False self.bound = False self.bindfail = False + self.add_event_handler('connected', self.handle_connected) self.register_handler( Callback('Stream Features', @@ -116,6 +117,14 @@ class ClientXMPP(BaseXMPP): self.register_feature( "", self._handle_start_session) + + def handle_connected(self, event=None): + #TODO: Use stream state here + self.authenticated = False + self.sessionstarted = False + self.bound = False + self.bindfail = False + def connect(self, address=tuple()): """ diff --git a/sleekxmpp/thirdparty/__init__.py b/sleekxmpp/thirdparty/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sleekxmpp/thirdparty/statemachine.py b/sleekxmpp/thirdparty/statemachine.py new file mode 100644 index 0000000..140c081 --- /dev/null +++ b/sleekxmpp/thirdparty/statemachine.py @@ -0,0 +1,278 @@ +""" + SleekXMPP: The Sleek XMPP Library + Copyright (C) 2010 Nathanael C. Fritz + This file is part of SleekXMPP. + + See the file LICENSE for copying permission. +""" +import threading +import time +import logging + +log = logging.getLogger(__name__) + + +class StateMachine(object): + + def __init__(self, states=[]): + self.lock = threading.Lock() + self.notifier = threading.Event() + self.__states= [] + self.addStates(states) + self.__default_state = self.__states[0] + self.__current_state = self.__default_state + + def addStates(self, states): + self.lock.acquire() + try: + for state in states: + if state in self.__states: + raise IndexError("The state '%s' is already in the StateMachine." % state) + self.__states.append( state ) + finally: self.lock.release() + + + def transition(self, from_state, to_state, wait=0.0, func=None, args=[], kwargs={} ): + ''' + Transition from the given `from_state` to the given `to_state`. + This method will return `True` if the state machine is now in `to_state`. It + will return `False` if a timeout occurred the transition did not occur. + If `wait` is 0 (the default,) this method returns immediately if the state machine + is not in `from_state`. + + If you want the thread to block and transition once the state machine to enters + `from_state`, set `wait` to a non-negative value. Note there is no 'block + indefinitely' flag since this leads to deadlock. If you want to wait indefinitely, + choose a reasonable value for `wait` (e.g. 20 seconds) and do so in a while loop like so: + + :: + + while not thread_should_exit and not state_machine.transition('disconnected', 'connecting', wait=20 ): + pass # timeout will occur every 20s unless transition occurs + if thread_should_exit: return + # perform actions here after successful transition + + This allows the thread to be responsive by setting `thread_should_exit=True`. + + The optional `func` argument allows the user to pass a callable operation which occurs + within the context of the state transition (e.g. while the state machine is locked.) + If `func` returns a True value, the transition will occur. If `func` returns a non- + True value or if an exception is thrown, the transition will not occur. Any thrown + exception is not caught by the state machine and is the caller's responsibility to handle. + If `func` completes normally, this method will return the value returned by `func.` If + values for `args` and `kwargs` are provided, they are expanded and passed like so: + `func( *args, **kwargs )`. + ''' + + return self.transition_any( (from_state,), to_state, wait=wait, + func=func, args=args, kwargs=kwargs ) + + + def transition_any(self, from_states, to_state, wait=0.0, func=None, args=[], kwargs={} ): + ''' + Transition from any of the given `from_states` to the given `to_state`. + ''' + + if not (isinstance(from_states,tuple) or isinstance(from_states,list)): + raise ValueError( "from_states should be a list or tuple" ) + + for state in from_states: + if not state in self.__states: + raise ValueError( "StateMachine does not contain from_state %s." % state ) + if not to_state in self.__states: + raise ValueError( "StateMachine does not contain to_state %s." % to_state ) + + + start = time.time() + while not self.__current_state in from_states or not self.lock.acquire(False): + # detect timeout: + remainder = start + wait - time.time() + if remainder > 0: self.notifier.wait(remainder) + else: return False + + try: # lock is acquired; all other threads will return false or wait until notify/timeout + if self.__current_state in from_states: # should always be True due to lock + + # Note that func might throw an exception, but that's OK, it aborts the transition + return_val = func(*args,**kwargs) if func is not None else True + + # some 'false' value returned from func, + # indicating that transition should not occur: + if not return_val: return return_val + + log.debug(' ==== TRANSITION %s -> %s', self.__current_state, to_state) + self._set_state( to_state ) + return return_val # some 'true' value returned by func or True if func was None + else: + log.error( "StateMachine bug!! The lock should ensure this doesn't happen!" ) + return False + finally: + self.notifier.set() # notify any waiting threads that the state has changed. + self.notifier.clear() + self.lock.release() + + + def transition_ctx(self, from_state, to_state, wait=0.0): + ''' + Use the state machine as a context manager. The transition occurs on /exit/ from + the `with` context, so long as no exception is thrown. For example: + + :: + + with state_machine.transition_ctx('one','two', wait=5) as locked: + if locked: + # the state machine is currently locked in state 'one', and will + # transition to 'two' when the 'with' statement ends, so long as + # no exception is thrown. + print 'Currently locked in state one: %s' % state_machine['one'] + + else: + # The 'wait' timed out, and no lock has been acquired + print 'Timed out before entering state "one"' + + print 'Since no exception was thrown, we are now in state "two": %s' % state_machine['two'] + + + The other main difference between this method and `transition()` is that the + state machine is locked for the duration of the `with` statement. Normally, + after a `transition()` occurs, the state machine is immediately unlocked and + available to another thread to call `transition()` again. + ''' + + if not from_state in self.__states: + raise ValueError( "StateMachine does not contain from_state %s." % from_state ) + if not to_state in self.__states: + raise ValueError( "StateMachine does not contain to_state %s." % to_state ) + + return _StateCtx(self, from_state, to_state, wait) + + + def ensure(self, state, wait=0.0, block_on_transition=False ): + ''' + Ensure the state machine is currently in `state`, or wait until it enters `state`. + ''' + return self.ensure_any( (state,), wait=wait, block_on_transition=block_on_transition ) + + + def ensure_any(self, states, wait=0.0, block_on_transition=False): + ''' + Ensure we are currently in one of the given `states` or wait until + we enter one of those states. + + Note that due to the nature of the function, you cannot guarantee that + the entirety of some operation completes while you remain in a given + state. That would require acquiring and holding a lock, which + would mean no other threads could do the same. (You'd essentially + be serializing all of the threads that are 'ensuring' their tasks + occurred in some state. + ''' + if not (isinstance(states,tuple) or isinstance(states,list)): + raise ValueError('states arg should be a tuple or list') + + for state in states: + if not state in self.__states: + raise ValueError( "StateMachine does not contain state '%s'" % state ) + + # if we're in the middle of a transition, determine whether we should + # 'fall back' to the 'current' state, or wait for the new state, in order to + # avoid an operation occurring in the wrong state. + # TODO another option would be an ensure_ctx that uses a semaphore to allow + # threads to indicate they want to remain in a particular state. + + # will return immediately if no transition is in process. + if block_on_transition: + # we're not in the middle of a transition; don't hold the lock + if self.lock.acquire(False): self.lock.release() + # wait for the transition to complete + else: self.notifier.wait() + + start = time.time() + while not self.__current_state in states: + # detect timeout: + remainder = start + wait - time.time() + if remainder > 0: self.notifier.wait(remainder) + else: return False + return True + + + def reset(self): + # TODO need to lock before calling this? + self.transition(self.__current_state, self.__default_state) + + + def _set_state(self, state): #unsynchronized, only call internally after lock is acquired + self.__current_state = state + return state + + + def current_state(self): + ''' + Return the current state name. + ''' + return self.__current_state + + + def __getitem__(self, state): + ''' + Non-blocking, non-synchronized test to determine if we are in the given state. + Use `StateMachine.ensure(state)` to wait until the machine enters a certain state. + ''' + return self.__current_state == state + + def __str__(self): + return "".join(( "StateMachine(", ','.join(self.__states), "): ", self.__current_state )) + + + +class _StateCtx: + + def __init__( self, state_machine, from_state, to_state, wait ): + self.state_machine = state_machine + self.from_state = from_state + self.to_state = to_state + self.wait = wait + self._locked = False + + def __enter__(self): + start = time.time() + while not self.state_machine[ self.from_state ] or not self.state_machine.lock.acquire(False): + # detect timeout: + remainder = start + self.wait - time.time() + if remainder > 0: self.state_machine.notifier.wait(remainder) + else: + log.debug('StateMachine timeout while waiting for state: %s', self.from_state ) + return False + + self._locked = True # lock has been acquired at this point + self.state_machine.notifier.clear() + log.debug('StateMachine entered context in state: %s', + self.state_machine.current_state() ) + return True + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_val is not None: + log.exception( "StateMachine exception in context, remaining in state: %s\n%s:%s", + self.state_machine.current_state(), exc_type.__name__, exc_val ) + + if self._locked: + if exc_val is None: + log.debug(' ==== TRANSITION %s -> %s', + self.state_machine.current_state(), self.to_state) + self.state_machine._set_state( self.to_state ) + + self.state_machine.notifier.set() + self.state_machine.lock.release() + + return False # re-raise any exception + +if __name__ == '__main__': + + def callback(s, s2): + print 1, s.transition('on', 'off', wait=0.0, func=callback, args=[s,s2]) + print 2, s2.transition('off', 'on', func=callback, args=[s,s2]) + return True + + s = StateMachine(('off', 'on')) + s2 = StateMachine(('off', 'on')) + print 3, s.transition('off', 'on', wait=0.0, func=callback, args=[s,s2]), + print s.current_state(), s2.current_state() diff --git a/sleekxmpp/xmlstream/scheduler.py b/sleekxmpp/xmlstream/scheduler.py index 0e3d30b..0af1f50 100644 --- a/sleekxmpp/xmlstream/scheduler.py +++ b/sleekxmpp/xmlstream/scheduler.py @@ -104,7 +104,7 @@ class Scheduler(object): quit -- Stop the scheduler. """ - def __init__(self, parentqueue=None): + def __init__(self, parentqueue=None, parentstop=None): """ Create a new scheduler. @@ -116,6 +116,7 @@ class Scheduler(object): self.thread = None self.run = False self.parentqueue = parentqueue + self.parentstop = parentstop def process(self, threaded=True): """ @@ -135,38 +136,41 @@ class Scheduler(object): def _process(self): """Process scheduled tasks.""" self.run = True - while self.run: - try: - wait = 1 - updated = False - if self.schedule: - wait = self.schedule[0].next - time.time() - try: - if wait <= 0.0: - newtask = self.addq.get(False) - else: - newtask = self.addq.get(True, wait) - except queue.Empty: - cleanup = [] - for task in self.schedule: - if time.time() >= task.next: - updated = True - if not task.run(): - cleanup.append(task) + try: + while self.run: + wait = 1 + updated = False + if self.schedule: + wait = self.schedule[0].next - time.time() + try: + if wait <= 0.0: + newtask = self.addq.get(False) else: - break - for task in cleanup: - x = self.schedule.pop(self.schedule.index(task)) - else: - updated = True - self.schedule.append(newtask) - finally: - if updated: - self.schedule = sorted(self.schedule, - key=lambda task: task.next) - except KeyboardInterrupt: - self.run = False - + newtask = self.addq.get(True, wait) + except queue.Empty: + cleanup = [] + for task in self.schedule: + if time.time() >= task.next: + updated = True + if not task.run(): + cleanup.append(task) + else: + break + for task in cleanup: + x = self.schedule.pop(self.schedule.index(task)) + else: + updated = True + self.schedule.append(newtask) + finally: + if updated: + self.schedule = sorted(self.schedule, + key=lambda task: task.next) + except KeyboardInterrupt: + self.run = False + if self.parentstop is not None: self.parentstop.set() + except SystemExit: + self.run = False + if self.parentstop is not None: self.parentstop.set() logging.debug("Qutting Scheduler thread") if self.parentqueue is not None: self.parentqueue.put(('quit', None, None)) diff --git a/sleekxmpp/xmlstream/statemachine.py b/sleekxmpp/xmlstream/statemachine.py deleted file mode 100644 index 8a1aa22..0000000 --- a/sleekxmpp/xmlstream/statemachine.py +++ /dev/null @@ -1,59 +0,0 @@ -""" - SleekXMPP: The Sleek XMPP Library - Copyright (C) 2010 Nathanael C. Fritz - This file is part of SleekXMPP. - - See the file LICENSE for copying permission. -""" -from __future__ import with_statement -import threading - -class StateMachine(object): - - def __init__(self, states=[], groups=[]): - self.lock = threading.Lock() - self.__state = {} - self.__default_state = {} - self.__group = {} - self.addStates(states) - self.addGroups(groups) - - def addStates(self, states): - with self.lock: - for state in states: - if state in self.__state or state in self.__group: - raise IndexError("The state or group '%s' is already in the StateMachine." % state) - self.__state[state] = states[state] - self.__default_state[state] = states[state] - - def addGroups(self, groups): - with self.lock: - for gstate in groups: - if gstate in self.__state or gstate in self.__group: - raise IndexError("The key or group '%s' is already in the StateMachine." % gstate) - for state in groups[gstate]: - if state in self.__state: - raise IndexError("The group %s contains a key %s which is not set in the StateMachine." % (gstate, state)) - self.__group[gstate] = groups[gstate] - - def set(self, state, status): - with self.lock: - if state in self.__state: - self.__state[state] = bool(status) - else: - raise KeyError("StateMachine does not contain state %s." % state) - - def __getitem__(self, key): - if key in self.__group: - for state in self.__group[key]: - if not self.__state[state]: - return False - return True - return self.__state[key] - - def __getattr__(self, attr): - return self.__getitem__(attr) - - def reset(self): - self.__state = self.__default_state - diff --git a/sleekxmpp/xmlstream/statemanager.py b/sleekxmpp/xmlstream/statemanager.py deleted file mode 100644 index c7f76e7..0000000 --- a/sleekxmpp/xmlstream/statemanager.py +++ /dev/null @@ -1,139 +0,0 @@ -""" - SleekXMPP: The Sleek XMPP Library - Copyright (C) 2010 Nathanael C. Fritz, Lance J.T. Stout - This file is part of SleekXMPP. - - See the file LICENSE for copying permission. -""" - -from __future__ import with_statement -import threading - - -class StateError(Exception): - """Raised whenever a state transition was attempted but failed.""" - - -class StateManager(object): - """ - At the very core of SleekXMPP there is a need to track various - library configuration settings, XML stream features, and the - network connection status. The state manager is responsible for - tracking this information in a thread-safe manner. - - State 'variables' store the current state of these items as simple - string values or booleans. Changing those values must be done - according to transitions defined when creating the state variable. - - If a state variable is given a value that is not allowed according - to the transition definitions, a StateError is raised. When a - valid value is assigned an event is raised named: - - _state_changed_nameofthestatevariable - - The event carries a dictionary containing the previous and the new - state values. - """ - - def __init__(self, event_func=None): - """ - Initialize the state manager. The parameter event_func should be - the event() method of a SleekXMPP object in order to enable - _state_changed_* events. - """ - self.main_lock = threading.Lock() - self.locks = {} - self.state_variables = {} - - if event_func is not None: - self.event = event_func - else: - self.event = lambda name, data: None - - def add(self, name, default=False, values=None, transitions=None): - """ - Create a new state variable. - - When transitions is specified, only those defined state change - transitions will be allowed. - - When values is specified (and not transitions), any state changes - between those values are allowed. - - If neither values nor transitions are defined, then the state variable - will be a binary switch between True and False. - """ - if name in self.state_variables: - raise IndexError("State variable %s already exists" % name) - - self.locks[name] = threading.Lock() - with self.locks[name]: - var = {'value': default, - 'default': default, - 'transitions': {}} - - if transitions is not None: - for start in transitions: - var['transitions'][start] = set(transitions[start]) - elif values is not None: - values = set(values) - for value in values: - var['transitions'][value] = values - elif values is None: - var['transitions'] = {True: [False], - False: [True]} - - self.state_variables[name] = var - - def addStates(self, var_defs): - """ - Create multiple state variables at once. - """ - for var, data in var_defs: - self.add(var, - default=data.get('default', False), - values=data.get('values', None), - transitions=data.get('transitions', None)) - - def force_set(self, name, val): - """ - Force setting a state variable's value by overriding transition checks. - """ - with self.locks[name]: - self.state_variables[name]['value'] = val - - def reset(self, name): - """ - Reset a state variable to its default value. - """ - with self.locks[name]: - default = self.state_variables[name]['default'] - self.state_variables[name]['value'] = default - - def __getitem__(self, name): - """ - Get the value of a state variable if it exists. - """ - with self.locks[name]: - if name not in self.state_variables: - raise IndexError("State variable %s does not exist" % name) - return self.state_variables[name]['value'] - - def __setitem__(self, name, val): - """ - Attempt to set the value of a state variable, but raise StateError - if the transition is undefined. - - A _state_changed_* event is triggered after a successful transition. - """ - with self.locks[name]: - if name not in self.state_variables: - raise IndexError("State variable %s does not exist" % name) - current = self.state_variables[name]['value'] - if current == val: - return - if val in self.state_variables[name]['transitions'][current]: - self.state_variables[name]['value'] = val - self.event('_state_changed_%s' % name, {'from': current, 'to': val}) - else: - raise StateError("Can not transition from '%s' to '%s'" % (str(current), str(val))) diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py index 3152ec9..23aef6b 100644 --- a/sleekxmpp/xmlstream/xmlstream.py +++ b/sleekxmpp/xmlstream/xmlstream.py @@ -21,7 +21,8 @@ try: except ImportError: import Queue as queue -from sleekxmpp.xmlstream import StateMachine, Scheduler, tostring +from sleekxmpp.thirdparty.statemachine import StateMachine +from sleekxmpp.xmlstream import Scheduler, tostring from sleekxmpp.xmlstream.stanzabase import StanzaBase, ET # In Python 2.x, file socket objects are broken. A patched socket @@ -92,6 +93,8 @@ class XMLStream(object): stream_header -- The closing tag of the stream's root element. use_ssl -- Flag indicating if SSL should be used. use_tls -- Flag indicating if TLS should be used. + stop -- threading Event used to stop all threads. + auto_reconnect-- Flag to determine whether we auto reconnect. Methods: add_event_handler -- Add a handler for a custom event. @@ -152,15 +155,8 @@ class XMLStream(object): self.ssl_support = SSL_SUPPORT - # TODO: Integrate the new state machine. - self.state = StateMachine() - self.state.addStates({'connected': False, - 'is client': False, - 'ssl': False, - 'tls': False, - 'reconnect': True, - 'processing': False, - 'disconnecting': False}) + self.state = StateMachine(('disconnected', 'connected')) + self.state._set_state('disconnected') self.address = (host, int(port)) self.filesocket = None @@ -178,9 +174,10 @@ class XMLStream(object): self.stream_header = "" self.stream_footer = "" + self.stop = threading.Event() self.event_queue = queue.Queue() self.send_queue = queue.Queue() - self.scheduler = Scheduler(self.event_queue) + self.scheduler = Scheduler(self.event_queue, self.stop) self.namespace_map = {} @@ -193,7 +190,8 @@ class XMLStream(object): self._id = 0 self._id_lock = threading.Lock() - self.run = True + self.auto_reconnect = True + self.is_client = False def new_id(self): """ @@ -232,17 +230,23 @@ class XMLStream(object): if host and port: self.address = (host, int(port)) + self.is_client = True # Respect previous SSL and TLS usage directives. if use_ssl is not None: self.use_ssl = use_ssl if use_tls is not None: self.use_tls = use_tls - self.state.set('is client', True) - # Repeatedly attempt to connect until a successful connection # is established. - while reattempt and not self.state['connected']: + connected = self.state.transition('disconnected', 'connected', func=self._connect) + while reattempt and not connected: + connected = self.state.transition('disconnected', 'connected', func=self._connect) + return connected + + + def _connect(self): + self.stop.clear() self.socket = self.socket_class(socket.AF_INET, socket.SOCK_STREAM) self.socket.settimeout(None) if self.use_ssl and self.ssl_support: @@ -257,13 +261,15 @@ class XMLStream(object): try: self.socket.connect(self.address) - self.set_socket(self.socket) - self.state.set('connected', True) + self.set_socket(self.socket, ignore=True) + #this event is where you should set your application state + self.event("connected", direct=True) return True except socket.error as serr: error_msg = "Could not connect. Socket Error #%s: %s" logging.error(error_msg % (serr.errno, serr.strerror)) time.sleep(1) + return False def disconnect(self, reconnect=False): """ @@ -277,40 +283,38 @@ class XMLStream(object): and processing should be restarted. Defaults to False. """ - self.event("disconnected") - self.state.set('reconnect', reconnect) - if self.state['disconnecting']: - return - if not self.state['reconnect']: - logging.debug("Disconnecting...") - self.state.set('disconnecting', True) - self.run = False - self.scheduler.run = False - if self.state['connected']: - # Send the end of stream marker. - self.send_raw(self.stream_footer) - # Wait for confirmation that the stream was - # closed in the other direction. - time.sleep(1) + self.state.transition('connected', 'disconnected', wait=0.0, func=self._disconnect, args=(reconnect,)) + + def _disconnect(self, reconnect=False): + # Send the end of stream marker. + self.send_raw(self.stream_footer) + # Wait for confirmation that the stream was + # closed in the other direction. + time.sleep(1) + if not reconnect: + self.auto_reconnect = False + self.stop.set() try: self.socket.close() self.filesocket.close() self.socket.shutdown(socket.SHUT_RDWR) except socket.error as serr: pass + finally: + #clear your application state + self.event("disconnected", direct=True) + return True + def reconnect(self): """ Reset the stream's state and reconnect to the server. """ - logging.info("Reconnecting") - self.event("disconnected") - self.state.set('tls', False) - self.state.set('ssl', False) - time.sleep(1) - self.connect() + logging.debug("reconnecting...") + self.state.transition('connected', 'disconnected', wait=0.0, func=self._disconnect, args=(True,)) + return self.state.transition('disconnected', 'connected', wait=0.0, func=self._connect) - def set_socket(self, socket): + def set_socket(self, socket, ignore=False): """ Set the socket to use for the stream. @@ -318,6 +322,7 @@ class XMLStream(object): Arguments: socket -- The new socket to use. + ignore -- don't set the state """ self.socket = socket if socket is not None: @@ -331,7 +336,8 @@ class XMLStream(object): self.filesocket = FileSocket(self.socket) else: self.filesocket = self.socket.makefile('rb', 0) - self.state.set('connected', True) + if not ignore: + self.state._set_state('connected') def start_tls(self): """ @@ -490,17 +496,21 @@ class XMLStream(object): self.__event_handlers[name] = filter(filter_pointers, self.__event_handlers[name]) - def event(self, name, data={}): + def event(self, name, data={}, direct=False): """ Manually trigger a custom event. Arguments: - name -- The name of the event to trigger. - data -- Data that will be passed to each event handler. - Defaults to an empty dictionary. + name -- The name of the event to trigger. + data -- Data that will be passed to each event handler. + Defaults to an empty dictionary. + direct -- Runs the event directly if True. """ for handler in self.__event_handlers.get(name, []): - self.event_queue.put(('event', handler, copy.copy(data))) + if direct: + handler[0](copy.copy(data)) + else: + self.event_queue.put(('event', handler, copy.copy(data))) if handler[2]: # If the handler is disposable, we will go ahead and # remove it now instead of waiting for it to be @@ -640,49 +650,36 @@ class XMLStream(object): # The body of this loop will only execute once per connection. # Additional passes will be made only if an error occurs and # reconnecting is permitted. - while self.run and (firstrun or self.state['reconnect']): - self.state.set('processing', True) + while not self.stop.isSet() and firstrun or self.auto_reconnect: firstrun = False try: - if self.state['is client']: + if self.is_client: self.send_raw(self.stream_header) # The call to self.__read_xml will block and prevent # the body of the loop from running until a disconnect # occurs. After any reconnection, the stream header will # be resent and processing will resume. - while self.run and self.__read_xml(): + while not self.stop.isSet() and self.__read_xml(): # Ensure the stream header is sent for any # new connections. - if self.state['is client']: + if self.is_client: self.send_raw(self.stream_header) except KeyboardInterrupt: - logging.debug("Keyboard Escape Detected") - self.state.set('processing', False) - self.state.set('reconnect', False) - self.disconnect() - self.run = False - self.scheduler.run = False - self.event_queue.put(('quit', None, None)) - return + logging.debug("Keyboard Escape Detected in _process") + self.stop.set() except SystemExit: - self.event_queue.put(('quit', None, None)) - return + logging.debug("SystemExit in _process") + self.stop.set() except socket.error: - if not self.state.reconnect: - return - self.state.set('processing', False) logging.exception('Socket Error') - self.disconnect(reconnect=True) except: - if not self.state.reconnect: - return - self.state.set('processing', False) logging.exception('Connection error. Reconnecting.') - self.disconnect(reconnect=True) - if self.state['reconnect']: + if not self.stop.isSet() and self.auto_reconnect: self.reconnect() - self.state.set('processing', False) - self.event_queue.put(('quit', None, None)) + else: + self.disconnect() + self.event_queue.put(('quit', None, None)) + self.scheduler.run = False def __read_xml(self): """ @@ -705,7 +702,6 @@ class XMLStream(object): if depth == 0: # The stream's root element has closed, # terminating the stream. - self.disconnect(reconnect=self.state['reconnect']) logging.debug("Ending read XML loop") return False elif depth == 1: @@ -754,8 +750,11 @@ class XMLStream(object): stanza_copy = stanza_type(self, copy.deepcopy(xml)) handler.prerun(stanza_copy) self.event_queue.put(('stanza', handler, stanza_copy)) - if handler.checkDelete(): - self.__handlers.pop(self.__handlers.index(handler)) + try: + if handler.checkDelete(): + self.__handlers.pop(self.__handlers.index(handler)) + except: + pass #not thread safe unhandled = False # Some stanzas require responses, such as Iq queries. A default @@ -773,61 +772,73 @@ class XMLStream(object): handlers may be spawned in individual threads. """ logging.debug("Loading event runner") - while self.run: - try: - event = self.event_queue.get(True, timeout=5) - except queue.Empty: - event = None - except KeyboardInterrupt: - self.run = False - self.scheduler.run = False - if event is None: - continue + try: + while not self.stop.isSet(): + try: + event = self.event_queue.get(True, timeout=5) + except queue.Empty: + event = None + if event is None: + continue - etype, handler = event[0:2] - args = event[2:] + etype, handler = event[0:2] + args = event[2:] - if etype == 'stanza': - try: - handler.run(args[0]) - except Exception as e: - error_msg = 'Error processing stream handler: %s' - logging.exception(error_msg % handler.name) - args[0].exception(e) - elif etype == 'schedule': - try: - logging.debug(args) - handler(*args[0]) - except: - logging.exception('Error processing scheduled task') - elif etype == 'event': - func, threaded, disposable = handler - try: - if threaded: - x = threading.Thread(name="Event_%s" % str(func), - target=func, - args=args) - x.start() - else: - func(*args) - except: - logging.exception('Error processing event handler: %s') - elif etype == 'quit': - logging.debug("Quitting event runner thread") - return False + if etype == 'stanza': + try: + handler.run(args[0]) + except Exception as e: + error_msg = 'Error processing stream handler: %s' + logging.exception(error_msg % handler.name) + args[0].exception(e) + elif etype == 'schedule': + try: + logging.debug(args) + handler(*args[0]) + except: + logging.exception('Error processing scheduled task') + elif etype == 'event': + func, threaded, disposable = handler + try: + if threaded: + x = threading.Thread(name="Event_%s" % str(func), + target=func, + args=args) + x.start() + else: + func(*args) + except: + logging.exception('Error processing event handler: %s') + elif etype == 'quit': + logging.debug("Quitting event runner thread") + return False + except KeyboardInterrupt: + logging.debug("Keyboard Escape Detected in _event_runner") + self.disconnect() + return + except SystemExit: + self.disconnect() + self.event_queue.put(('quit', None, None)) + return def _send_thread(self): """ Extract stanzas from the send queue and send them on the stream. """ - while self.run: - data = self.send_queue.get(True) - logging.debug("SEND: %s" % data) - try: - self.socket.send(data.encode('utf-8')) - except: - logging.warning("Failed to send %s" % data) - self.state.set('connected', False) - if self.state.reconnect: - logging.exception("Disconnected. Socket Error.") - self.disconnect(reconnect=True) + try: + while not self.stop.isSet(): + data = self.send_queue.get(True) + logging.debug("SEND: %s" % data) + try: + self.socket.send(data.encode('utf-8')) + except: + logging.warning("Failed to send %s" % data) + self.disconnect(self.auto_reconnect) + except KeyboardInterrupt: + logging.debug("Keyboard Escape Detected in _send_thread") + self.disconnect() + return + except SystemExit: + self.disconnect() + self.event_queue.put(('quit', None, None)) + return