[tor-commits] [stem/master] Move send and recv implementation to BaseSocket

[atagar at torproject.org]

commit f0d78cf8e509531aad79da2f73ea3a6af3c9548e
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Dec 31 14:12:58 2017 -0800

    Move send and recv implementation to BaseSocket
    
    Boilerplate for sending and receiving is common for all sockets. Only thing
    that differs is message formatting and parsing.
---
 stem/socket.py | 114 ++++++++++++++++++++++++++++++++++-----------------------
 1 file changed, 68 insertions(+), 46 deletions(-)

diff --git a/stem/socket.py b/stem/socket.py
index 8ff316a7..6bf13a39 100644
--- a/stem/socket.py
+++ b/stem/socket.py
@@ -223,6 +223,72 @@ class BaseSocket(object):
       if is_change:
         self._close()
 
+  def _send(self, message, handler):
+    """
+    Send message in a thread safe manner. Handler is expected to be of the form...
+
+    ::
+
+      my_handler(socket_file, message)
+    """
+
+    with self._send_lock:
+      try:
+        if not self.is_alive():
+          raise stem.SocketClosed()
+
+        handler(self._socket_file, message)
+      except stem.SocketClosed:
+        # if send_message raises a SocketClosed then we should properly shut
+        # everything down
+
+        if self.is_alive():
+          self.close()
+
+        raise
+
+  def _recv(self, handler):
+    """
+    Receives a message in a thread safe manner. Handler is expected to be of the form...
+
+    ::
+
+      my_handler(socket_file)
+    """
+
+    with self._recv_lock:
+      try:
+        # makes a temporary reference to the _socket_file because connect()
+        # and close() may set or unset it
+
+        socket_file = self._socket_file
+
+        if not socket_file:
+          raise stem.SocketClosed()
+
+        return handler(socket_file)
+      except stem.SocketClosed:
+        # If recv_message raises a SocketClosed then we should properly shut
+        # everything down. However, there's a couple cases where this will
+        # cause deadlock...
+        #
+        # * This SocketClosed was *caused by* a close() call, which is joining
+        #   on our thread.
+        #
+        # * A send() call that's currently in flight is about to call close(),
+        #   also attempting to join on us.
+        #
+        # To resolve this we make a non-blocking call to acquire the send lock.
+        # If we get it then great, we can close safely. If not then one of the
+        # above are in progress and we leave the close to them.
+
+        if self.is_alive():
+          if self._send_lock.acquire(False):
+            self.close()
+            self._send_lock.release()
+
+        raise
+
   def _get_send_lock(self):
     """
     The send lock is useful to classes that interact with us at a deep level
@@ -300,20 +366,7 @@ class ControlSocket(BaseSocket):
       * :class:`stem.SocketClosed` if the socket is known to be shut down
     """
 
-    with self._send_lock:
-      try:
-        if not self.is_alive():
-          raise stem.SocketClosed()
-
-        send_message(self._socket_file, message)
-      except stem.SocketClosed:
-        # if send_message raises a SocketClosed then we should properly shut
-        # everything down
-
-        if self.is_alive():
-          self.close()
-
-        raise
+    self._send(message, send_message)
 
   def recv(self):
     """
@@ -327,38 +380,7 @@ class ControlSocket(BaseSocket):
       * :class:`stem.SocketClosed` if the socket closes before we receive a complete message
     """
 
-    with self._recv_lock:
-      try:
-        # makes a temporary reference to the _socket_file because connect()
-        # and close() may set or unset it
-
-        socket_file = self._socket_file
-
-        if not socket_file:
-          raise stem.SocketClosed()
-
-        return recv_message(socket_file)
-      except stem.SocketClosed:
-        # If recv_message raises a SocketClosed then we should properly shut
-        # everything down. However, there's a couple cases where this will
-        # cause deadlock...
-        #
-        # * This SocketClosed was *caused by* a close() call, which is joining
-        #   on our thread.
-        #
-        # * A send() call that's currently in flight is about to call close(),
-        #   also attempting to join on us.
-        #
-        # To resolve this we make a non-blocking call to acquire the send lock.
-        # If we get it then great, we can close safely. If not then one of the
-        # above are in progress and we leave the close to them.
-
-        if self.is_alive():
-          if self._send_lock.acquire(False):
-            self.close()
-            self._send_lock.release()
-
-        raise
+    return self._recv(recv_message)
 
 
 class ControlPort(ControlSocket):