Remove stream feature handlers on session_start.

Based on profiling, using around 35 stream handlers quarters the number
of basic message stanzas that can be processed in a second, in
comparison to only using the bare minimum of four handlers.

To help, we can drop handlers for stream features once the session
has started. So that we can re-enable these handlers when a stream
must restart, the 'stream_start' event has been added which fires
whenever a stream header is received.

The 'stream_start' event is a more generic replacement for the
existing start_stream_handler() method.

2 files changed, 30 insertions, 9 deletions

diff --git a/sleekxmpp/features/feature_mechanisms/mechanisms.py b/sleekxmpp/features/feature_mechanisms/mechanisms.py
index deff5d30..86dbb361 100644
--- a/sleekxmpp/features/feature_mechanisms/mechanisms.py
+++ b/sleekxmpp/features/feature_mechanisms/mechanisms.py
@@ -63,6 +63,17 @@ class feature_mechanisms(base_plugin):
         self.xmpp.register_stanza(stanza.Challenge)
         self.xmpp.register_stanza(stanza.Response)
 
+        self.xmpp.add_event_handler('stream_start',
+                self._handle_stream_start)
+        self.xmpp.add_event_handler('session_start',
+                self._handle_session_start)
+
+        self.xmpp.register_feature('mechanisms',
+                self._handle_sasl_auth,
+                restart=True,
+                order=self.config.get('order', 100))
+
+    def _handle_stream_start(self, root):
         self.xmpp.register_handler(
                 Callback('SASL Success',
                          MatchXPath(stanza.Success.tag_name()),
@@ -80,10 +91,10 @@ class feature_mechanisms(base_plugin):
                          MatchXPath(stanza.Challenge.tag_name()),
                          self._handle_challenge))
 
-        self.xmpp.register_feature('mechanisms',
-                self._handle_sasl_auth,
-                restart=True,
-                order=self.config.get('order', 100))
+    def _handle_session_start(self, e):
+        self.xmpp.remove_handler('SASL Success')
+        self.xmpp.remove_handler('SASL Failure')
+        self.xmpp.remove_handler('SASL Challenge')
 
     def _handle_sasl_auth(self, features):
         """
diff --git a/sleekxmpp/features/feature_starttls/starttls.py b/sleekxmpp/features/feature_starttls/starttls.py
index 4e2b6621..a7f048b1 100644
--- a/sleekxmpp/features/feature_starttls/starttls.py
+++ b/sleekxmpp/features/feature_starttls/starttls.py
@@ -27,20 +27,30 @@ class feature_starttls(base_plugin):
         self.description = "STARTTLS Stream Feature"
         self.stanza = stanza
 
-        self.xmpp.register_handler(
-                Callback('STARTTLS Proceed',
-                        MatchXPath(stanza.Proceed.tag_name()),
-                        self._handle_starttls_proceed,
-                        instream=True))
         self.xmpp.register_feature('starttls',
                 self._handle_starttls,
                 restart=True,
                 order=self.config.get('order', 0))
 
+        self.xmpp.add_event_handler('stream_start',
+                self._handle_stream_start)
+        self.xmpp.add_event_handler('session_start',
+                self._handle_session_start)
+
         self.xmpp.register_stanza(stanza.Proceed)
         self.xmpp.register_stanza(stanza.Failure)
         register_stanza_plugin(StreamFeatures, stanza.STARTTLS)
 
+    def _handle_stream_start(self, root):
+        self.xmpp.register_handler(
+                Callback('STARTTLS Proceed',
+                        MatchXPath(stanza.Proceed.tag_name()),
+                        self._handle_starttls_proceed,
+                        instream=True))
+
+    def _handle_session_start(self, e):
+        self.xmpp.remove_handler('STARTTLS Proceed')
+
     def _handle_starttls(self, features):
         """
         Handle notification that the server supports TLS.