Auto reconnection when channel state changed to idle (#1846)

Signed-off-by: shaoyue.chen <[email protected]>

pymilvus/client/grpc_handler.py

@@ -4,7 +4,7 @@
 import socket
 import time
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 from urllib import parse
 
 import grpc
@@ -87,6 +87,16 @@ def __init__(
  self._set_authorization(**kwargs)
  self._setup_db_interceptor(kwargs.get("db_name", None))
  self._setup_grpc_channel()
+ self.callbacks = []
+
+ def register_state_change_callback(self, callback: Callable):
+ self.callbacks.append(callback)
+ self._channel.subscribe(callback, try_to_connect=True)
+
+ def deregister_state_change_callbacks(self):
+ for callback in self.callbacks:
+ self._channel.unsubscribe(callback)
+ self.callbacks = []
 
  def __get_address(self, uri: str, host: str, port: str) -> str:
  if host != "" and port != "" and is_legal_host(host) and is_legal_port(port):
@@ -141,6 +151,7 @@ def _wait_for_channel_ready(self, timeout: Union[float] = 10):
  raise e from e
 
  def close(self):
+ self.deregister_state_change_callbacks()
  self._channel.close()
 
  def reset_db_name(self, db_name: str):

pymilvus/orm/connections.py

@@ -11,7 +11,9 @@
 # the License.
 
 import copy
+import logging
 import threading
+import time
 from typing import Callable, Tuple, Union
 from urllib import parse
 
@@ -24,6 +26,8 @@
 )
 from pymilvus.settings import Config
 
+logger = logging.getLogger(__name__)
+
 VIRTUAL_PORT = 443
 
 
@@ -58,6 +62,53 @@ def __new__(cls, *args, **kwargs):
  return super().__new__(cls, *args, **kwargs)
 
 
+class ReconnectHandler:
+ def __init__(self, conns: object, connection_name: str, kwargs: object) -> None:
+ self.connection_name = connection_name
+ self.conns = conns
+ self._kwargs = kwargs
+ self.is_idle_state = False
+ self.reconnect_lock = threading.Lock()
+
+ def check_state_and_reconnect_later(self):
+ check_after_seconds = 3
+ logger.debug(f"state is idle, schedule reconnect in {check_after_seconds} seconds")
+ time.sleep(check_after_seconds)
+ if not self.is_idle_state:
+ logger.debug("idle state changed, skip reconnect")
+ return
+ with self.reconnect_lock:
+ logger.info("reconnect on idle state")
+ self.is_idle_state = False
+ try:
+ logger.debug("try disconnecting old connection...")
+ self.conns.disconnect(self.connection_name)
+ except Exception:
+ logger.warning("disconnect failed: {e}")
+ finally:
+ reconnected = False
+ while not reconnected:
+ try:
+ logger.debug("try reconnecting...")
+ self.conns.connect(self.connection_name, **self._kwargs)
+ reconnected = True
+ except Exception as e:
+ logger.warning(
+ f"reconnect failed: {e}, try again after {check_after_seconds} seconds"
+ )
+ time.sleep(check_after_seconds)
+ logger.info("reconnected")
+
+ def reconnect_on_idle(self, state: object):
+ logger.debug(f"state change to: {state}")
+ with self.reconnect_lock:
+ if state.value[1] != "idle":
+ self.is_idle_state = False
+ return
+ self.is_idle_state = True
+ threading.Thread(target=self.check_state_and_reconnect_later).start()
+
+
 class Connections(metaclass=SingleInstanceMetaClass):
  """Class for managing all connections of milvus. Used as a singleton in this module."""
 
@@ -270,6 +321,8 @@ def connect(
  Optional. Serving as the key for identification and authentication purposes.
  Whenever a token is furnished, we shall supplement the corresponding header
  to each RPC call.
+ * *keep_alive* (``bool``) --
+ Optional. Default is false. If set to true, client will keep an alive connection.
  * *db_name* (``str``) --
  Optional. default database name of this connection
  * *client_key_path* (``str``) --
@@ -293,13 +346,24 @@ def connect(
  >>> connections.connect("test", host="localhost", port="19530")
  """
 
+ # kwargs_copy is used for auto reconnect
+ kwargs_copy = copy.deepcopy(kwargs)
+ kwargs_copy["user"] = user
+ kwargs_copy["password"] = password
+ kwargs_copy["db_name"] = db_name
+ kwargs_copy["token"] = token
+
  def connect_milvus(**kwargs):
  gh = GrpcHandler(**kwargs)
 
  t = kwargs.get("timeout")
  timeout = t if isinstance(t, (int, float)) else Config.MILVUS_CONN_TIMEOUT
 
  gh._wait_for_channel_ready(timeout=timeout)
+ if kwargs.get("keep_alive", False):
+ gh.register_state_change_callback(
+ ReconnectHandler(self, alias, kwargs_copy).reconnect_on_idle
+ )
  kwargs.pop("password")
  kwargs.pop("token", None)
  kwargs.pop("secure", None)