optimize close behavior

Author: Time-Coder

.gitignore

@@ -2,4 +2,4 @@ __pycache__
 *.egg-info
 dist
 build
-test.py
+test*

Connector/NodeBasicMethods.py

@@ -42,6 +42,8 @@ def __init__(self, connection, server=None):
     self._queues = Queues(self)
     self._result_queues_for_future = Queues(self, private=True)
     self._recved_binary_queue = CloseableQueue()
+    self._is_closing = False
+    self._is_internal_closing = False
 
     if self._parent is None:
         self._local_vars = {}
@@ -99,20 +101,21 @@ def hold_on(self):
     threading.Event().wait()
 
 
-def close(self):
-    if self._parent is None:
-        try:
-            self._request(session_id=self._get_session_id())
-        except BaseException:
-            pass
-
-        try:
-            self._connection.close()
-        except BaseException:
-            pass
+def _close(self):
+    if self._is_internal_closing or self.is_closed:
+        return
+    
+    self._is_internal_closing = True
+    
+    self_address = self.address
+    try:
+        self._connection.close()
+    except BaseException:
+        pass
 
-        self._connection = None
+    self._connection = None
 
+    if self._parent is None:
         try:
             self._queue.close()
         except BaseException:
@@ -147,25 +150,18 @@ def close(self):
             self._actual_pipes.clear()
         except BaseException:
             pass
-
     else:
         try:
-            self._connection.close()
-        except BaseException:
-            pass
-
-        self._connection = None
-
-        try:
-            if self.address is not None:
+            if self_address is not None:
                 result = self._parent._callback_threads_pool.submit(
-                    self._parent._on_disconnected, self.address)
-                self._parent.disconnect_result[self.address] = result
+                    self._parent._on_disconnected, self_address
+                )
+                self._parent.disconnect_result[self_address] = result
         except BaseException:
             pass
 
         try:
-            self._parent._connected_nodes.pop(self.address)
+            self._parent._connected_nodes.pop(self_address)
         except BaseException:
             pass
 
@@ -193,6 +189,28 @@ def close(self):
     except BaseException:
         pass
 
+    self._is_internal_closing = False
+
+
+def close(self):
+    if self._is_closing or self.is_closed:
+        return
+    
+    self._is_closing = True
+
+    try:
+        self._request(session_id=self._get_session_id())
+    except BaseException:
+        pass
+
+    self._close()
+
+    self._is_closing = False
+
+
+def _process_close(self, session_id, request):
+    self._close()
+
 
 def _get_session_id(self):
     return uuid.uuid4().bytes
@@ -228,10 +246,6 @@ def _traceback(self, remote=True):
     return exception_str
 
 
-def _process_close(self, request):
-    self.close()
-
-
 def _send(self, session_id, type, value):
     binary = dill.dumps(value, 3)
     type_bytes = type.to_bytes(1, byteorder='little', signed=False)
@@ -392,10 +406,11 @@ def _recving_loop(self):
     while True:
         try:
             recved_binary = self._connection.recv(self.recv_buffer)
-            self._recved_binary_queue.put(recved_binary)
         except BaseException:
-            self.close()
+            self._close()
             return
+        
+        self._recved_binary_queue.put(recved_binary)
 
 
 def _decoding_loop(self):
@@ -449,12 +464,13 @@ def _decoding_loop(self):
 def init_basic_methods(cls):
     cls.__init__ = __init__
     cls.__del__ = __del__
+    cls._close = _close
     cls.close = close
+    cls._process_close = _process_close
     cls.hold_on = hold_on
     cls._start = _start
     cls._get_session_id = _get_session_id
     cls._traceback = _traceback
-    cls._process_close = _process_close
     cls._send = _send
     cls._send_signal = _send_signal
     cls._request = _request

Connector/NodeFileTransferMethods.py

@@ -10,7 +10,7 @@ def _write_to_file(self, data, file_name):
     self._request(self._get_session_id(), data=data, file_name=file_name)
 
 
-def _process__write_to_file(self, request):
+def _process__write_to_file(self, session_id, request):
     try:
         file_name = os.path.abspath(request["data"]["file_name"])
         if self._out_file is not None and self._out_file.name != file_name:
@@ -34,7 +34,7 @@ def _close_file(self):
     self._request(self._get_session_id())
 
 
-def _process__close_file(self, request):
+def _process__close_file(self, session_id, request):
     try:
         if self._out_file is not None:
             self._out_file.close()

Connector/NodeInternalClasses.py

@@ -31,6 +31,18 @@ def __delitem__(self, name):
 
         collections.OrderedDict.__delitem__(self, name)
 
+    def __repr__(self):
+        result = "[\n"
+        i = 0
+        for value in self.values():
+            address = value.address
+            result += f"    Client('{address[0]}', {address[1]})"
+            if i != len(self)-1:
+                result += ","
+            result += "\n"
+        result += "]"
+        return result
+
 
 class QueueDict(dict):
 

Connector/Server.py

@@ -23,6 +23,7 @@ def __init__(self, ip=None, port=None):
         self._queue = CloseableQueue()
         self._queues = QueueDict()
 
+        self._is_closing = False
         self._continue_accept = False
         self._accept_thread = None
 
@@ -39,6 +40,23 @@ def start(self, ip=None, port=None):
             return
 
         self._connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        if hasattr(self, "_send_buffer"):
+            self._connection.setsockopt(
+                socket.SOL_SOCKET, socket.SO_SNDBUF, self._send_buffer
+            )
+        else:
+            self._send_buffer = self._connection.getsockopt(
+                socket.SOL_SOCKET, socket.SO_SNDBUF
+            )
+
+        if hasattr(self, "_recv_buffer"):
+            self._connection.setsockopt(
+                socket.SOL_SOCKET, socket.SO_RCVBUF, self._recv_buffer
+            )
+        else:
+            self._recv_buffer = self._connection.getsockopt(
+                socket.SOL_SOCKET, socket.SO_RCVBUF
+            )
 
         if ip is None:
             ip = get_ip()
@@ -56,6 +74,10 @@ def start(self, ip=None, port=None):
             self._accept_thread.start()
 
     def close(self):
+        if self._is_closing:
+            return
+
+        self._is_closing = True
         self._continue_accept = False
 
         try:
@@ -92,6 +114,8 @@ def close(self):
             self._accept_thread.join()
         self._accept_thread = None
 
+        self._is_closing = False
+
     @property
     def address(self):
         try:
@@ -234,6 +258,7 @@ def _acceptting_loop(self):
                 self.connect_results[address] = result
                 node._respond_ok(session_id=b'\x00'*16)
             except BaseException:
+                self.close()
                 break
 
     @property

Connector/__init__.py

@@ -1,3 +1,17 @@
 from .Client import Client
 from .Server import Server
 from .Config import Config
+
+import gc
+import atexit
+
+
+@atexit.register
+def close_all():
+    for obj in gc.get_objects():
+        if isinstance(obj, Client):
+            obj.close()
+
+    for obj in gc.get_objects():
+        if isinstance(obj, Server):
+            obj.close()

static_check.py

@@ -1,18 +1,20 @@
 import os
+import subprocess
+import sys
 
 
 def check_syntax(folder_name):
     for home, dirs, files in os.walk(folder_name):
         for file_name in files:
             if file_name.endswith(".py") and file_name != "__init__.py":
-                os.system("pyflakes " + home + "/" + file_name)
+                subprocess.call([sys.executable, "-m", "pyflakes", home + "/" + file_name])
 
 
 def check_style(folder_name):
     for home, dirs, files in os.walk(folder_name):
         for file_name in files:
             if file_name.endswith(".py"):
-                os.system("pycodestyle " + home + "/" + file_name)
+                subprocess.call([sys.executable, "-m", "pycodestyle", home + "/" + file_name])
 
 
 check_syntax("Connector")