new state machine in place

8 files changed, 462 insertions, 356 deletions

diff --git a/setup.py b/setup.py
index 4c90a1c5..f6eace7f 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 13de39df..1aab4a78 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(
             "<session xmlns='urn:ietf:params:xml:ns:xmpp-session' />",
             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 00000000..e69de29b
--- /dev/null
+++ b/sleekxmpp/thirdparty/__init__.py
diff --git a/sleekxmpp/thirdparty/statemachine.py b/sleekxmpp/thirdparty/statemachine.py
new file mode 100644
index 00000000..140c0819
--- /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 0e3d30b3..0af1f508 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 8a1aa22d..00000000
--- 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 c7f76e75..00000000
--- 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 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