Merge pull request #92 from camaze/nebula_optimize_cjn

【CodeFuse】optimize nebula startup time

Author: lightislost

.gitignore

@@ -19,7 +19,7 @@ zdatafront*
 *antgroup*
 *ipynb
 *log
-
+examples/mysql/db
 
 # frontend
 frontend/node_modules

muagent/db_handler/graph_db_handler/nebula_handler.py

@@ -63,13 +63,36 @@ def __init__(self,gb_config : GBConfig = None):
             self.nb_pw = '' or 'nebula'
             self.space_name = "client"
         else:
-            logger.info('NebulaGraph容器启动中，等待20秒')
-            time.sleep(20)
-            self.connection_pool.init([(gb_config.extra_kwargs.get("host"), gb_config.extra_kwargs.get("port"))], config)
+            logger.info('NebulaGraph容器启动中，请等待')
+            
+            if self.nebula_started(gb_config):
+                self.connection_pool.init([(gb_config.extra_kwargs.get("host"), gb_config.extra_kwargs.get("port"))], config)
+        
             self.username = gb_config.extra_kwargs.get("username")
             self.nb_pw = gb_config.extra_kwargs.get("password")
             self.space_name = gb_config.extra_kwargs.get("space")
 
+    def nebula_started(self, gb_config: GBConfig):
+        '''
+        Continuously ping the nebula server until it has started
+        If it does not start within 40 seconds, raise an exception
+        '''
+        # 80 attempts at 0.5 seconds each = 40 seconds
+        max_attempts = 80 
+
+        host = gb_config.extra_kwargs.get('host')
+        port = gb_config.extra_kwargs.get('port')
+
+        for attempt in range(max_attempts):
+            if self.connection_pool.ping([host, port]):
+                logger.info(f"nebula server is ok, total waiting time: {attempt * 0.5}s")
+                return True
+            
+            logger.info(f"ping nebula server.. waiting time: {attempt * 0.5}s")
+            time.sleep(0.5)
+
+        raise Exception('The nebula server did not start within 40 seconds.')
+
     def execute_cypher(self, cypher: str, space_name: str = '',ignore_log: bool = False, format_res: str = 'as_primitive', use_space_name: bool = True):
         '''
         @param space_name: space_name, if provided, will execute use space_name first
@@ -172,6 +195,11 @@ def show_space(self):
     def drop_space(self, space_name):
         cypher = f'DROP SPACE {space_name}'
         return self.execute_cypher(cypher)
+    
+    def show_hosts(self):
+        cypher = 'SHOW HOSTS'
+        resp = self.execute_cypher(cypher, use_space_name=False)
+        return resp
 
     def create_tag(self, tag_name: str, prop_dict: dict = {}):
         '''

muagent/service/ekg_construct/ekg_construct_base.py

@@ -180,20 +180,19 @@ def init_gb(self, do_init: bool=None):
             initialize_space = self.initialize_space  # True or False
             if initialize_space and self.gb_config.gb_type=="NebulaHandler":
                 self.gb.add_hosts('storaged0', 9779)
-                print('增加NebulaGraph Storage主机中，等待20秒')
-                time.sleep(20)
+                self.waiting_host_initialize()
 
                 if self.clear_history_data:
                     # 初始化space
                     self.gb.drop_space('client')
 
                 self.gb.create_space('client')
+                self.waiting_space_initialize()
 
                 # 创建node tags和edge types
                 self.create_gb_tags_and_edgetypes()
+                self.waiting_tags_edgetypes_initialize()
 
-                print('Node Tags和Edge Types初始化中，等待20秒......')
-                time.sleep(20)
         else:
             self.gb = None
 
@@ -223,7 +222,54 @@ def init_sls(self, do_init: bool=None):
                 self.embed_config, vb_config=self.vb_config)
         else:
             self.sls = None
+
+    def wait_for_condition(self, condition, timeout=40, wait_time=0.5, log_prefix=''):
+        max_attempts = int(timeout / wait_time)
+
+        total_time = 0
+
+        for _ in range(max_attempts):
+            if condition():
+                logger.info(f'[{log_prefix}] initialized successfully, total waiting time: {total_time}s')
+                return
+            
+            logger.info(f'[{log_prefix}] have waited for {total_time}s...')
+            total_time += wait_time
+            time.sleep(wait_time)
+
+        raise Exception(f'[{log_prefix}] not initialized within {timeout} seconds')
 
+    def waiting_host_initialize(self):
+        self.wait_for_condition(
+            condition=lambda: any(
+                host.get('Host') == 'storaged0' and host.get('Status') == 'ONLINE'
+                for host in self.gb.show_hosts()
+            ),
+            log_prefix='host_initializing'
+        )
+
+    def waiting_tags_edgetypes_initialize(self):
+        '''
+        Make sure that tags and edgetypes are created properly
+        Refer to create_gb_tags_and_edgetypes for more details
+        '''
+        self.wait_for_condition(
+            condition=lambda: (
+                len(self.gb.show_tags()) == len(TYPE2SCHEMA.items()) - 1 and
+                len(self.gb.show_edge_type()) == (len(TYPE2SCHEMA.items()) - 1) ** 2 * 3
+            ),
+            log_prefix='tags_edgetypes_initializing'
+        )
+
+    def waiting_space_initialize(self):
+        self.wait_for_condition(
+            condition=lambda:any(
+                space.get('Name') == 'client'
+                for space in self.gb.show_space()
+            ),
+            log_prefix='space_initializing'
+        )
+
     def _get_local_graph(
             self, nodes: List[GNode], edges: List[GEdge], rootid
         ) -> Tuple[List[str], Graph]: