Await in memory database start when creating

neo4j · Mar 22, 2024 · 6116e01 · 6116e01
1 parent 42ea42f
commit 6116e01
Showing 1 changed file with 45 additions and 1 deletion.
diff --git a/...age-engine-adapter/src/main/java/org/neo4j/gds/storageengine/InMemoryDatabaseCreator.java b/...age-engine-adapter/src/main/java/org/neo4j/gds/storageengine/InMemoryDatabaseCreator.java
@@ -27,6 +27,10 @@
 import org.neo4j.gds.core.loading.CatalogRequest;
 import org.neo4j.gds.core.loading.GraphStoreCatalog;
 import org.neo4j.graphdb.GraphDatabaseService;
+import org.neo4j.graphdb.event.DatabaseEventContext;
+import org.neo4j.graphdb.event.DatabaseEventListener;
+
+import java.util.concurrent.CountDownLatch;
 
 import static org.neo4j.gds.core.cypher.CypherGraphStoreCatalogHelper.setWrappedGraphStore;
 
@@ -55,11 +59,51 @@ public static void createDatabase(
             var graphStoreWithConfig = GraphStoreCatalog.get(CatalogRequest.of(username, databaseService.databaseName()), graphName);
             var cypherGraphStore = new CypherGraphStore(graphStoreWithConfig.graphStore());
             setWrappedGraphStore(graphStoreWithConfig.config(), cypherGraphStore);
-            StorageEngineProxy.createInMemoryDatabase(dbms, dbName, graphName, config);
+            createAndAwaitDatabase(dbms, dbName, graphName, config);
         } catch (Exception e) {
             InMemoryDatabaseCreationCatalog.removeDatabaseEntry(dbName);
             throw e;
         }
     }
 
+    private static void createAndAwaitDatabase(DatabaseManagementService dbms, String dbName, String graphName, Config config) {
+        var databaseCreationLatch = new CountDownLatch(1);
+        dbms.registerDatabaseEventListener(new DatabaseEventListener() {
+            @Override
+            public void databaseStart(DatabaseEventContext eventContext) {
+                if (eventContext.getDatabaseName().equals(dbName)) {
+                    databaseCreationLatch.countDown();
+                }
+            }
+
+            @Override
+            public void databaseShutdown(DatabaseEventContext eventContext) {
+
+            }
+
+            @Override
+            public void databasePanic(DatabaseEventContext eventContext) {
+                if (eventContext.getDatabaseName().equals(dbName)) {
+                    databaseCreationLatch.countDown();
+                }
+            }
+
+            @Override
+            public void databaseCreate(DatabaseEventContext eventContext) {
+
+            }
+
+            @Override
+            public void databaseDrop(DatabaseEventContext eventContext) {
+
+            }
+        });
+
+        StorageEngineProxy.createInMemoryDatabase(dbms, dbName, graphName, config);
+        try {
+            databaseCreationLatch.await();
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }