fix connection

nebula3/gclient/net/Connection.py

@@ -44,6 +44,7 @@ def __init__(self):
         self._ssl_conf = None
         self.use_http2 = False
         self.http_headers = None
+        self._closed = True
 
     def open(self, ip, port, timeout, use_http2=False, http_headers=None):
         """open the connection
@@ -87,7 +88,9 @@ def open_SSL(
                 self._connection._iprot.trans.close()
                 raise ClientServerIncompatibleException(resp.error_msg)
         except Exception as e:
+            self.close()
             raise
+        self._closed = False
 
     def __get_protocol(self, timeout, ssl_config):
         try:
@@ -152,6 +155,7 @@ def _reopen(self):
                 self.use_http2,
                 self.http_headers,
             )
+            self._closed = False
         else:
             self.open(
                 self._ip, self._port, self._timeout, self.use_http2, self.http_headers
@@ -266,7 +270,9 @@ def close(self):
         :return: void
         """
         try:
-            self._connection._iprot.trans.close()
+            if not self._closed:
+                self._connection._iprot.trans.close()
+                self._closed = True
         except Exception as e:
             logger.error(
                 "Close connection to {}:{} failed:{}".format(self._ip, self._port, e)

nebula3/gclient/net/SessionPool.py

@@ -421,8 +421,19 @@ def _new_session(self):
                 session = Session(connection, auth_result, self, False)
 
                 # switch to the space specified in the configs
-                resp = session.execute("USE {}".format(self._space_name))
+                try:
+                    resp = session.execute("USE {}".format(self._space_name))
+                except Exception:
+                    session.release()
+                    connection.close()
+                    raise RuntimeError(
+                        "Failed to get session, execute `use {}` failed.".format(
+                            self._space_name
+                        )
+                    )
                 if not resp.is_succeeded():
+                    session.release()
+                    connection.close()
                     raise RuntimeError(
                         "Failed to get session, cannot set the session space to {} error: {} {}".format(
                             self._space_name, resp.error_code(), resp.error_msg()
@@ -440,8 +451,11 @@ def _new_session(self):
                         )
                     )
                     self.close()
+                else:
+                    connection.close()
                 raise e
             except Exception:
+                connection.close()
                 raise
 
         raise RuntimeError(