Skip to content

Commit 2c4e547

Browse files
soerenreichardtsoerenreichardt
committed
Refactor token holder code
1 parent dbdc9d7 commit 2c4e547

File tree

3 files changed

+176
-63
lines changed

3 files changed

+176
-63
lines changed

‎cypher/api/storage-engine-adapter/src/main/java/org/neo4j/gds/compat/AbstractInMemoryStorageEngine.java‎

Lines changed: 8 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
*/
2020
package org.neo4j.gds.compat;
2121

22-
import org.apache.commons.lang3.mutable.MutableInt;
2322
import org.neo4j.counts.CountsAccessor;
2423
import org.neo4j.counts.CountsStore;
2524
import org.neo4j.exceptions.KernelException;
@@ -51,10 +50,8 @@
5150
import org.neo4j.storageengine.api.TransactionApplicationMode;
5251
import org.neo4j.storageengine.api.txstate.ReadableTransactionState;
5352
import org.neo4j.token.TokenHolders;
54-
import org.neo4j.token.api.NamedToken;
5553

5654
import java.util.Collection;
57-
import java.util.HashSet;
5855
import java.util.List;
5956
import java.util.function.BiFunction;
6057
import java.util.function.Supplier;
@@ -65,10 +62,10 @@
6562
public abstract class AbstractInMemoryStorageEngine implements StorageEngine {
6663

6764
private final DatabaseLayout databaseLayout;
68-
private final TokenHolders tokenHolders;
6965
private final InMemoryTransactionStateVisitor txStateVisitor;
7066
private final Supplier<CommandCreationContext> commandCreationContextSupplier;
7167
private final TriFunction<CypherGraphStore, TokenHolders, CountsStore, StorageReader> storageReaderFn;
68+
private final TokenManager tokenManager;
7269
private final CountsStore countsStore;
7370
protected final MetadataProvider metadataProvider;
7471
protected final CypherGraphStore graphStore;
@@ -83,16 +80,11 @@ public AbstractInMemoryStorageEngine(
8380
TriFunction<CypherGraphStore, TokenHolders, CountsStore, StorageReader> storageReaderFn
8481
) {
8582
this.databaseLayout = databaseLayout;
86-
this.tokenHolders = tokenHolders;
87-
var graphName = InMemoryDatabaseCreationCatalog.getRegisteredDbCreationGraphName(databaseLayout.getDatabaseName());
88-
this.graphStore = getGraphStoreFromCatalog(graphName);
83+
this.graphStore = getGraphStoreFromCatalog(databaseLayout.getDatabaseName());
8984
this.txStateVisitor = txStateVisitorFn.apply(graphStore, tokenHolders);
9085
this.commandCreationContextSupplier = commandCreationContextSupplier;
9186
this.storageReaderFn = storageReaderFn;
92-
initializeTokenHolders();
93-
graphStore.initialize(tokenHolders);
94-
graphStore.registerPropertyRemovalCallback(txStateVisitor::removeNodeProperty);
95-
graphStore.registerPropertyAddedCallback(this::addProperty);
87+
this.tokenManager = new TokenManager(tokenHolders, txStateVisitor, graphStore);
9688
this.countsStore = countsStoreFn.apply(graphStore, tokenHolders);
9789
this.metadataProvider = metadataProvider;
9890
}
@@ -103,7 +95,7 @@ protected void createCommands(ReadableTransactionState txState) throws KernelExc
10395

10496
@Override
10597
public StorageReader newReader() {
106-
return storageReaderFn.apply(graphStore, tokenHolders, countsStore);
98+
return storageReaderFn.apply(graphStore, tokenManager.tokenHolders(), countsStore);
10799
}
108100

109101
@Override
@@ -116,7 +108,8 @@ public void apply(CommandsToApply batch, TransactionApplicationMode mode) {
116108
}
117109

118110
@Override
119-
public void init() { }
111+
public void init() {
112+
}
120113

121114
@Override
122115
public void start() {
@@ -184,41 +177,8 @@ public void lockRecoveryCommands(
184177

185178
}
186179

187-
private void initializeTokenHolders() {
188-
MutableInt labelCounter = new MutableInt(0);
189-
MutableInt typeCounter = new MutableInt(0);
190-
MutableInt propertyCounter = new MutableInt(0);
191-
192-
var propertyKeys = new HashSet<>(graphStore.nodePropertyKeys());
193-
propertyKeys.addAll(graphStore.relationshipPropertyKeys());
194-
propertyKeys.forEach(propertyKey ->
195-
tokenHolders
196-
.propertyKeyTokens()
197-
.addToken(new NamedToken(propertyKey, propertyCounter.getAndIncrement()))
198-
);
199-
200-
graphStore
201-
.nodeLabels()
202-
.forEach(nodeLabel -> tokenHolders
203-
.labelTokens()
204-
.addToken(new NamedToken(nodeLabel.name(), labelCounter.getAndIncrement())));
205-
206-
graphStore
207-
.relationshipTypes()
208-
.forEach(relType -> tokenHolders
209-
.relationshipTypeTokens()
210-
.addToken(new NamedToken(relType.name(), typeCounter.getAndIncrement())));
211-
}
212-
213-
private void addProperty(String propertyName) {
214-
try {
215-
tokenHolders.propertyKeyTokens().getOrCreateId(propertyName);
216-
} catch (KernelException e) {
217-
throw new RuntimeException(e);
218-
}
219-
}
220-
221-
private static CypherGraphStore getGraphStoreFromCatalog(String graphName) {
180+
private static CypherGraphStore getGraphStoreFromCatalog(String databaseName) {
181+
var graphName = InMemoryDatabaseCreationCatalog.getRegisteredDbCreationGraphName(databaseName);
222182
return (CypherGraphStore) GraphStoreCatalog.getAllGraphStores()
223183
.filter(graphStoreWithUserNameAndConfig -> graphStoreWithUserNameAndConfig
224184
.config()

‎cypher/api/storage-engine-adapter/src/main/java/org/neo4j/gds/compat/TokenManager.java‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
/*
2+
* Copyright (c) "Neo4j"
3+
* Neo4j Sweden AB [http://neo4j.com]
4+
*
5+
* This file is part of Neo4j.
6+
*
7+
* Neo4j is free software: you can redistribute it and/or modify
8+
* it under the terms of the GNU General Public License as published by
9+
* the Free Software Foundation, either version 3 of the License, or
10+
* (at your option) any later version.
11+
*
12+
* This program is distributed in the hope that it will be useful,
13+
* but WITHOUT ANY WARRANTY; without even the implied warranty of
14+
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15+
* GNU General Public License for more details.
16+
*
17+
* You should have received a copy of the GNU General Public License
18+
* along with this program. If not, see <http://www.gnu.org/licenses/>.
19+
*/
20+
package org.neo4j.gds.compat;
21+
22+
import org.apache.commons.lang3.mutable.MutableInt;
23+
import org.neo4j.exceptions.KernelException;
24+
import org.neo4j.gds.core.cypher.CypherGraphStore;
25+
import org.neo4j.gds.storageengine.InMemoryTransactionStateVisitor;
26+
import org.neo4j.kernel.lifecycle.LifecycleAdapter;
27+
import org.neo4j.token.TokenHolders;
28+
import org.neo4j.token.api.NamedToken;
29+
30+
import java.util.HashSet;
31+
32+
class TokenManager extends LifecycleAdapter implements CypherGraphStore.StateVisitor {
33+
34+
private final TokenHolders tokenHolders;
35+
private final InMemoryTransactionStateVisitor transactionStateVisitor;
36+
private final CypherGraphStore graphStore;
37+
38+
TokenManager(
39+
TokenHolders tokenHolders,
40+
InMemoryTransactionStateVisitor transactionStateVisitor,
41+
CypherGraphStore graphStore
42+
) {
43+
this.tokenHolders = tokenHolders;
44+
this.transactionStateVisitor = transactionStateVisitor;
45+
this.graphStore = graphStore;
46+
47+
init();
48+
}
49+
50+
@Override
51+
public void init() {
52+
initializeTokensFromGraphStore();
53+
graphStore.initialize(tokenHolders);
54+
graphStore.registerStateVisitor(this);
55+
}
56+
57+
@Override
58+
public void nodePropertyRemoved(String propertyKey) {
59+
this.transactionStateVisitor.removeNodeProperty(propertyKey);
60+
}
61+
62+
@Override
63+
public void nodePropertyAdded(String propertyKey) {
64+
addProperty(propertyKey);
65+
}
66+
67+
@Override
68+
public void nodeLabelAdded(String nodeLabel) {
69+
addNodeLabel(nodeLabel);
70+
}
71+
72+
@Override
73+
public void relationshipTypeAdded(String relationshipType) {
74+
addRelationshipType(relationshipType);
75+
}
76+
77+
public TokenHolders tokenHolders() {
78+
return this.tokenHolders;
79+
}
80+
81+
private void addProperty(String propertyName) {
82+
try {
83+
tokenHolders.propertyKeyTokens().getOrCreateId(propertyName);
84+
} catch (KernelException e) {
85+
throw new RuntimeException(e);
86+
}
87+
}
88+
89+
private void addNodeLabel(String nodeLabel) {
90+
try {
91+
tokenHolders.labelTokens().getOrCreateId(nodeLabel);
92+
} catch (KernelException e) {
93+
throw new RuntimeException(e);
94+
}
95+
}
96+
97+
private void addRelationshipType(String relationshipType) {
98+
try {
99+
tokenHolders.relationshipTypeTokens().getOrCreateId(relationshipType);
100+
} catch (KernelException e) {
101+
throw new RuntimeException(e);
102+
}
103+
}
104+
105+
private void initializeTokensFromGraphStore() {
106+
// When this method is called there is no kernel available
107+
// which is needed to use the `getOrCreate` method on
108+
// the `TokenHolders`
109+
var labelCounter = new MutableInt(0);
110+
var typeCounter = new MutableInt(0);
111+
var propertyCounter = new MutableInt(0);
112+
113+
var propertyKeys = new HashSet<>(graphStore.nodePropertyKeys());
114+
propertyKeys.addAll(graphStore.relationshipPropertyKeys());
115+
propertyKeys.forEach(propertyKey ->
116+
tokenHolders
117+
.propertyKeyTokens()
118+
.addToken(new NamedToken(propertyKey, propertyCounter.getAndIncrement()))
119+
);
120+
121+
graphStore
122+
.nodeLabels()
123+
.forEach(nodeLabel -> tokenHolders
124+
.labelTokens()
125+
.addToken(new NamedToken(nodeLabel.name(), labelCounter.getAndIncrement())));
126+
127+
graphStore
128+
.relationshipTypes()
129+
.forEach(relType -> tokenHolders
130+
.relationshipTypeTokens()
131+
.addToken(new NamedToken(relType.name(), typeCounter.getAndIncrement())));
132+
}
133+
134+
135+
}

‎cypher/cypher-core/src/main/java/org/neo4j/gds/core/cypher/CypherGraphStore.java‎

Lines changed: 33 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,38 +20,39 @@
2020
package org.neo4j.gds.core.cypher;
2121

2222
import org.neo4j.gds.NodeLabel;
23+
import org.neo4j.gds.RelationshipType;
2324
import org.neo4j.gds.api.GraphStore;
2425
import org.neo4j.gds.api.GraphStoreAdapter;
2526
import org.neo4j.gds.api.IdMap;
27+
import org.neo4j.gds.api.Relationships;
2628
import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
2729
import org.neo4j.token.TokenHolders;
30+
import org.neo4j.values.storable.NumberType;
2831

32+
import java.util.ArrayList;
33+
import java.util.List;
34+
import java.util.Optional;
2935
import java.util.Set;
30-
import java.util.function.Consumer;
3136

3237
public class CypherGraphStore extends GraphStoreAdapter implements NodeLabelUpdater {
3338

3439
private final CypherIdMap cypherIdMap;
40+
private final List<StateVisitor> stateVisitors;
3541

3642
private RelationshipIds relationshipIds;
37-
private Consumer<String> propertyRemovalCallback;
38-
private Consumer<String> propertyAddedCallback;
3943

4044
public CypherGraphStore(GraphStore graphStore) {
4145
super(graphStore);
4246
this.cypherIdMap = new CypherIdMap(super.nodes());
47+
this.stateVisitors = new ArrayList<>();
4348
}
4449

4550
public void initialize(TokenHolders tokenHolders) {
4651
this.relationshipIds = RelationshipIds.fromGraphStore(innerGraphStore(), tokenHolders);
4752
}
4853

49-
public void registerPropertyRemovalCallback(Consumer<String> propertyRemovalCallback) {
50-
this.propertyRemovalCallback = propertyRemovalCallback;
51-
}
52-
53-
public void registerPropertyAddedCallback(Consumer<String> propertyAddedCallback) {
54-
this.propertyAddedCallback = propertyAddedCallback;
54+
public void registerStateVisitor(StateVisitor stateVisitor) {
55+
this.stateVisitors.add(stateVisitor);
5556
}
5657

5758
@Override
@@ -72,22 +73,39 @@ public void addLabelToNode(long nodeId, NodeLabel nodeLabel) {
7273
@Override
7374
public void removeNodeProperty(String propertyKey) {
7475
super.removeNodeProperty(propertyKey);
75-
if (propertyRemovalCallback != null) {
76-
propertyRemovalCallback.accept(propertyKey);
77-
}
76+
stateVisitors.forEach(stateVisitor -> stateVisitor.nodePropertyRemoved(propertyKey));
7877
}
7978

8079
@Override
8180
public void addNodeProperty(
8281
Set<NodeLabel> nodeLabels, String propertyKey, NodePropertyValues propertyValues
8382
) {
8483
super.addNodeProperty(nodeLabels, propertyKey, propertyValues);
85-
if (propertyAddedCallback != null) {
86-
propertyAddedCallback.accept(propertyKey);
87-
}
84+
stateVisitors.forEach(stateVisitor -> stateVisitor.nodePropertyAdded(propertyKey));
85+
}
86+
87+
@Override
88+
public void addRelationshipType(
89+
RelationshipType relationshipType,
90+
Optional<String> relationshipPropertyKey,
91+
Optional<NumberType> relationshipPropertyType,
92+
Relationships relationships
93+
) {
94+
super.addRelationshipType(relationshipType, relationshipPropertyKey, relationshipPropertyType, relationships);
95+
stateVisitors.forEach(stateVisitor -> stateVisitor.relationshipTypeAdded(relationshipType.name()));
8896
}
8997

9098
public RelationshipIds relationshipIds() {
9199
return this.relationshipIds;
92100
}
101+
102+
public interface StateVisitor {
103+
void nodePropertyRemoved(String propertyKey);
104+
105+
void nodePropertyAdded(String propertyKey);
106+
107+
void nodeLabelAdded(String nodeLabel);
108+
109+
void relationshipTypeAdded(String relationshipType);
110+
}
93111
}

0 commit comments

Comments
 (0)