diff options
Diffstat (limited to 'sleekxmpp/xmlstream/xmlstream.py')
| -rw-r--r-- | sleekxmpp/xmlstream/xmlstream.py | 261 |
1 files changed, 136 insertions, 125 deletions
diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py index 3152ec94..23aef6b7 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 = "<stream>" self.stream_footer = "</stream>" + 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 |