From 512e2b31c46708401c3cba3f3f65c0cc092feef6 Mon Sep 17 00:00:00 2001
From: "peng.xu" <peng.xu@zilliz.com>
Date: Thu, 19 Sep 2019 10:17:00 +0800
Subject: [PATCH] add pre run handlers

---
 mishards/main.py   | 10 ----------
 mishards/server.py | 28 +++++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/mishards/main.py b/mishards/main.py
index e9c47f9edf..7fac55dfa2 100644
--- a/mishards/main.py
+++ b/mishards/main.py
@@ -1,21 +1,11 @@
 import os, sys
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from urllib.parse import urlparse
-import socket
-
 from mishards import (
         settings,
-        db, connect_mgr,
-        discover,
         grpc_server as server)
 
 def main():
-    discover.start()
-    woserver = settings.WOSERVER if not settings.TESTING else settings.TESTING_WOSERVER
-    url = urlparse(woserver)
-    connect_mgr.register('WOSERVER',
-            '{}://{}:{}'.format(url.scheme, socket.gethostbyname(url.hostname), url.port))
     server.run(port=settings.SERVER_PORT)
     return 0
 
diff --git a/mishards/server.py b/mishards/server.py
index 185ed3c957..19cca2c18a 100644
--- a/mishards/server.py
+++ b/mishards/server.py
@@ -1,17 +1,21 @@
 import logging
 import grpc
 import time
+import socket
+from urllib.parse import urlparse
+from functools import wraps
 from concurrent import futures
 from grpc._cython import cygrpc
 from milvus.grpc_gen.milvus_pb2_grpc import add_MilvusServiceServicer_to_server
 from mishards.service_handler import ServiceHandler
-import mishards.settings
+from mishards import settings, discover
 
 logger = logging.getLogger(__name__)
 
 
 class Server:
     def __init__(self, conn_mgr, port=19530, max_workers=10, **kwargs):
+        self.pre_run_handlers = set()
         self.exit_flag = False
         self.port = int(port)
         self.conn_mgr = conn_mgr
@@ -21,6 +25,27 @@ class Server:
                 (cygrpc.ChannelArgKey.max_receive_message_length, -1)]
         )
 
+        self.register_pre_run_handler(self.pre_run_handler)
+
+    def pre_run_handler(self):
+        woserver = settings.WOSERVER if not settings.TESTING else settings.TESTING_WOSERVER
+        url = urlparse(woserver)
+        ip = socket.gethostbyname(url.hostname)
+        logger.error(ip)
+        socket.inet_pton(socket.AF_INET, ip)
+        self.conn_mgr.register('WOSERVER',
+                '{}://{}:{}'.format(url.scheme, ip, url.port))
+
+    def register_pre_run_handler(self, func):
+        logger.info('Regiterring {} into server pre_run_handlers'.format(func))
+        self.pre_run_handlers.add(func)
+        return func
+
+    def on_pre_run(self):
+        for handler in self.pre_run_handlers:
+            handler()
+        discover.start()
+
     def start(self, port=None):
         add_MilvusServiceServicer_to_server(ServiceHandler(conn_mgr=self.conn_mgr), self.server_impl)
         self.server_impl.add_insecure_port("[::]:{}".format(str(port or self._port)))
@@ -29,6 +54,7 @@ class Server:
     def run(self, port):
         logger.info('Milvus server start ......')
         port = port or self.port
+        self.on_pre_run()
 
         self.start(port)
         logger.info('Successfully')