diff --git a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImpl.java b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImpl.java index e619aff4fb..7a61dd502a 100755 --- a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImpl.java +++ b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImpl.java @@ -74,6 +74,7 @@ public class HstoreSessionsImpl extends HstoreSessions { private static volatile Boolean initializedNode = Boolean.FALSE; private static volatile PDClient defaultPdClient; private static volatile HgStoreClient hgStoreClient; + private final GraphStoreClient graphStoreClient; private final HugeConfig config; private final HstoreSession session; private final Map tables; @@ -81,11 +82,23 @@ public class HstoreSessionsImpl extends HstoreSessions { private final String graphName; public HstoreSessionsImpl(HugeConfig config, String database, String store) { + this(config, database, store, null); + } + + HstoreSessionsImpl(HugeConfig config, String database, String store, + GraphStoreClient graphStoreClient) { super(config, database, store); this.config = config; this.graphName = database + "/" + store; - this.initStoreNode(config); - this.session = new HstoreSession(this.config, graphName); + if (graphStoreClient == null) { + this.initStoreNode(config); + this.graphStoreClient = new GraphStoreClient(); + } else { + this.graphStoreClient = graphStoreClient; + } + this.session = new HstoreSession( + this.config, this.graphName, + this.graphStoreClient.openSession(this.graphName)); this.tables = new ConcurrentHashMap<>(); this.refCount = new AtomicInteger(1); } @@ -130,10 +143,11 @@ public void open() throws Exception { E.checkArgument(partitionCount > -1, "The value of hstore.partition_count " + "cannot be less than 0."); - defaultPdClient.setGraph(Metapb.Graph.newBuilder() - .setGraphName(this.graphName) - .setPartitionCount(partitionCount) - .build()); + this.graphStoreClient.setGraph( + Metapb.Graph.newBuilder() + .setGraphName(this.graphName) + .setPartitionCount(partitionCount) + .build()); infoInitializedGraph.add(this.graphName); } } @@ -181,12 +195,18 @@ public void truncateTable(String table) { public void clear() { this.session.deleteGraph(); try { - hgStoreClient.getPdClient().delGraph(this.graphName); + this.graphStoreClient.delGraph(this.graphName); } catch (PDException ignored) { + } finally { + clearInitializedGraph(this.graphName); } } + private static void clearInitializedGraph(String graphName) { + infoInitializedGraph.remove(graphName); + } + @Override public final Session session() { return (Session) super.getOrNewSession(); @@ -194,7 +214,8 @@ public final Session session() { @Override protected final Session newSession() { - return new HstoreSession(this.config(), this.graphName); + return new HstoreSession(this.config(), this.graphName, + this.graphStoreClient.openSession(this.graphName)); } @Override @@ -216,6 +237,21 @@ protected synchronized void doClose() { private void checkValid() { } + static class GraphStoreClient { + + HgStoreSession openSession(String graphName) { + return hgStoreClient.openSession(graphName); + } + + void setGraph(Metapb.Graph graph) throws PDException { + defaultPdClient.setGraph(graph); + } + + void delGraph(String graphName) throws PDException { + hgStoreClient.getPdClient().delGraph(graphName); + } + } + private static class ColumnIterator implements BackendColumnIterator, Countable { @@ -418,10 +454,11 @@ private final class HstoreSession extends Session { private final HgStoreSession graph; int changedSize = 0; - public HstoreSession(HugeConfig conf, String graphName) { + public HstoreSession(HugeConfig conf, String graphName, + HgStoreSession graph) { setGraphName(graphName); setConf(conf); - this.graph = hgStoreClient.openSession(graphName); + this.graph = graph; } @Override diff --git a/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImplTest.java b/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImplTest.java index 69f2591e25..6476e37cc5 100644 --- a/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImplTest.java +++ b/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImplTest.java @@ -19,8 +19,15 @@ import java.io.IOException; import java.io.InputStream; +import java.lang.reflect.Proxy; import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.hugegraph.config.HugeConfig; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.store.HgStoreSession; import org.junit.Assert; import org.junit.Test; @@ -36,4 +43,61 @@ public void testProductionClassDoesNotReferenceTestAssert() throws IOException { "org/apache/hugegraph/testutil/Assert")); } } + + @Test + public void testClearAndRecreateGraphRepeatedly() throws Exception { + String graphName = "hugegraph/hstore-clear-test"; + AtomicInteger graphRegistrations = new AtomicInteger(); + AtomicInteger storeGraphDeletes = new AtomicInteger(); + AtomicInteger pdGraphDeletes = new AtomicInteger(); + HgStoreSession storeSession = (HgStoreSession) Proxy.newProxyInstance( + HgStoreSession.class.getClassLoader(), + new Class[]{HgStoreSession.class}, + (proxy, method, args) -> { + Assert.assertEquals("deleteGraph", method.getName()); + Assert.assertArrayEquals(new Object[]{graphName}, args); + storeGraphDeletes.incrementAndGet(); + return true; + }); + HstoreSessionsImpl.GraphStoreClient client = + new HstoreSessionsImpl.GraphStoreClient() { + @Override + HgStoreSession openSession(String name) { + Assert.assertEquals(graphName, name); + return storeSession; + } + + @Override + void setGraph(Metapb.Graph graph) { + Assert.assertEquals(graphName, graph.getGraphName()); + graphRegistrations.incrementAndGet(); + } + + @Override + void delGraph(String name) throws PDException { + Assert.assertEquals(graphName, name); + pdGraphDeletes.incrementAndGet(); + throw new PDException(0, "exercise clear() finally block"); + } + }; + HugeConfig config = new HugeConfig(Collections.emptyMap()); + HstoreSessionsImpl sessions = new HstoreSessionsImpl( + config, "hugegraph", "hstore-clear-test", client); + + try { + sessions.open(); + Assert.assertEquals(1, graphRegistrations.get()); + + for (int i = 0; i < 50; i++) { + sessions.clear(); + sessions.open(); + Assert.assertEquals(i + 2, graphRegistrations.get()); + } + + Assert.assertEquals(50, storeGraphDeletes.get()); + Assert.assertEquals(50, pdGraphDeletes.get()); + } finally { + sessions.clear(); + } + } }