Skip to content

Commit df4ebed

Browse files
soerenreichardtsoerenreichardt
committed
Use caches instead of hash maps
Also test that cache eviction works correctly
1 parent bc0106b commit df4ebed

File tree

2 files changed

+140
-21
lines changed

2 files changed

+140
-21
lines changed

‎core/src/main/java/org/neo4j/gds/api/EphemeralResultStore.java‎

Lines changed: 44 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,19 @@
1919
*/
2020
package org.neo4j.gds.api;
2121

22+
import com.github.benmanes.caffeine.cache.Cache;
23+
import com.github.benmanes.caffeine.cache.Caffeine;
24+
import com.github.benmanes.caffeine.cache.Scheduler;
25+
import com.github.benmanes.caffeine.cache.Ticker;
2226
import org.neo4j.gds.api.nodeproperties.ValueType;
2327
import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
28+
import org.neo4j.gds.core.concurrency.ExecutorServiceUtil;
29+
import org.neo4j.gds.core.utils.ClockService;
2430

2531
import java.util.Collection;
26-
import java.util.HashMap;
2732
import java.util.List;
28-
import java.util.Map;
33+
import java.util.concurrent.ScheduledExecutorService;
34+
import java.util.concurrent.TimeUnit;
2935
import java.util.function.LongUnaryOperator;
3036
import java.util.stream.Stream;
3137

@@ -34,16 +40,17 @@ public class EphemeralResultStore implements ResultStore {
3440
private static final String NO_PROPERTY_KEY = "";
3541
private static final List<String> NO_PROPERTIES_LIST = List.of(NO_PROPERTY_KEY);
3642

37-
private final Map<NodeKey, NodePropertyValues> nodeProperties;
38-
private final Map<String, NodeLabelEntry> nodeIdsByLabel;
39-
private final Map<RelationshipKey, RelationshipEntry> relationships;
40-
private final Map<RelationshipKey, RelationshipStreamEntry> relationshipStreams;
43+
private final Cache<NodeKey, NodePropertyValues> nodeProperties;
44+
private final Cache<String, NodeLabelEntry> nodeIdsByLabel;
45+
private final Cache<RelationshipKey, RelationshipEntry> relationships;
46+
private final Cache<RelationshipKey, RelationshipStreamEntry> relationshipStreams;
4147

4248
public EphemeralResultStore() {
43-
this.nodeProperties = new HashMap<>();
44-
this.nodeIdsByLabel = new HashMap<>();
45-
this.relationships = new HashMap<>();
46-
this.relationshipStreams = new HashMap<>();
49+
var singleThreadScheduler = ExecutorServiceUtil.createSingleThreadScheduler("GDS-ResultStore");
50+
this.nodeProperties = createCache(singleThreadScheduler);
51+
this.nodeIdsByLabel = createCache(singleThreadScheduler);
52+
this.relationships = createCache(singleThreadScheduler);
53+
this.relationshipStreams = createCache(singleThreadScheduler);
4754
}
4855

4956
@Override
@@ -53,12 +60,12 @@ public void addNodePropertyValues(List<String> nodeLabels, String propertyKey, N
5360

5461
@Override
5562
public NodePropertyValues getNodePropertyValues(List<String> nodeLabels, String propertyKey) {
56-
return this.nodeProperties.get(new NodeKey(nodeLabels, propertyKey));
63+
return this.nodeProperties.getIfPresent(new NodeKey(nodeLabels, propertyKey));
5764
}
5865

5966
@Override
6067
public void removeNodePropertyValues(List<String> nodeLabels, String propertyKey) {
61-
this.nodeProperties.remove(new NodeKey(nodeLabels, propertyKey));
68+
this.nodeProperties.invalidate(new NodeKey(nodeLabels, propertyKey));
6269
}
6370

6471
@Override
@@ -68,17 +75,17 @@ public void addNodeLabel(String nodeLabel, long nodeCount, LongUnaryOperator toO
6875

6976
@Override
7077
public boolean hasNodeLabel(String nodeLabel) {
71-
return this.nodeIdsByLabel.containsKey(nodeLabel);
78+
return this.nodeIdsByLabel.getIfPresent(nodeLabel) != null;
7279
}
7380

7481
@Override
7582
public NodeLabelEntry getNodeIdsByLabel(String nodeLabel) {
76-
return this.nodeIdsByLabel.get(nodeLabel);
83+
return this.nodeIdsByLabel.getIfPresent(nodeLabel);
7784
}
7885

7986
@Override
8087
public void removeNodeLabel(String nodeLabel) {
81-
this.nodeIdsByLabel.remove(nodeLabel);
88+
this.nodeIdsByLabel.invalidate(nodeLabel);
8289
}
8390

8491
@Override
@@ -112,12 +119,12 @@ public void addRelationshipStream(
112119

113120
@Override
114121
public RelationshipStreamEntry getRelationshipStream(String relationshipType, List<String> propertyKeys) {
115-
return this.relationshipStreams.get(new RelationshipKey(relationshipType, propertyKeys));
122+
return this.relationshipStreams.getIfPresent(new RelationshipKey(relationshipType, propertyKeys));
116123
}
117124

118125
@Override
119126
public void removeRelationshipStream(String relationshipType, List<String> propertyKeys) {
120-
this.relationshipStreams.remove(new RelationshipKey(relationshipType, propertyKeys));
127+
this.relationshipStreams.invalidate(new RelationshipKey(relationshipType, propertyKeys));
121128
}
122129

123130
@Override
@@ -127,7 +134,7 @@ public RelationshipEntry getRelationship(String relationshipType) {
127134

128135
@Override
129136
public RelationshipEntry getRelationship(String relationshipType, String propertyKey) {
130-
return this.relationships.get(new RelationshipKey(relationshipType, List.of(propertyKey)));
137+
return this.relationships.getIfPresent(new RelationshipKey(relationshipType, List.of(propertyKey)));
131138
}
132139

133140
@Override
@@ -137,7 +144,7 @@ public void removeRelationship(String relationshipType) {
137144

138145
@Override
139146
public void removeRelationship(String relationshipType, String propertyKey) {
140-
this.relationships.remove(new RelationshipKey(relationshipType, List.of(propertyKey)));
147+
this.relationships.invalidate(new RelationshipKey(relationshipType, List.of(propertyKey)));
141148
}
142149

143150
@Override
@@ -147,12 +154,28 @@ public boolean hasRelationship(String relationshipType) {
147154

148155
@Override
149156
public boolean hasRelationship(String relationshipType, List<String> propertyKeys) {
150-
return this.relationships.containsKey(new RelationshipKey(relationshipType, propertyKeys));
157+
return this.relationships.getIfPresent(new RelationshipKey(relationshipType, propertyKeys)) != null;
151158
}
152159

153160
@Override
154161
public boolean hasRelationshipStream(String relationshipType, List<String> propertyKeys) {
155-
return this.relationshipStreams.containsKey(new RelationshipKey(relationshipType, propertyKeys));
162+
return this.relationshipStreams.getIfPresent(new RelationshipKey(relationshipType, propertyKeys)) != null;
163+
}
164+
165+
private static <K, V> Cache<K, V> createCache(ScheduledExecutorService singleThreadScheduler) {
166+
return Caffeine.newBuilder()
167+
.expireAfterAccess(10, TimeUnit.MINUTES)
168+
.ticker(new ClockServiceWrappingTicker())
169+
.executor(singleThreadScheduler)
170+
.scheduler(Scheduler.forScheduledExecutorService(singleThreadScheduler))
171+
.build();
172+
}
173+
174+
private static class ClockServiceWrappingTicker implements Ticker {
175+
@Override
176+
public long read() {
177+
return ClockService.clock().millis() * 1000000;
178+
}
156179
}
157180

158181
private record NodeKey(List<String> nodeLabels, String propertyKey) {}

‎core/src/test/java/org/neo4j/gds/api/EphemeralResultStoreTest.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,24 @@
2222
import org.junit.jupiter.api.Test;
2323
import org.neo4j.gds.api.nodeproperties.ValueType;
2424
import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
25+
import org.neo4j.gds.extension.FakeClockExtension;
26+
import org.neo4j.gds.extension.Inject;
27+
import org.neo4j.time.FakeClock;
2528

2629
import java.util.List;
30+
import java.util.concurrent.TimeUnit;
2731
import java.util.function.LongUnaryOperator;
2832
import java.util.stream.Stream;
2933

3034
import static org.assertj.core.api.Assertions.assertThat;
3135
import static org.mockito.Mockito.mock;
3236

37+
@FakeClockExtension
3338
class EphemeralResultStoreTest {
3439

40+
@Inject
41+
FakeClock clock;
42+
3543
@Test
3644
void shouldStoreNodeProperty() {
3745
var resultStore = new EphemeralResultStore();
@@ -66,6 +74,22 @@ void shouldRemoveNodeProperty() {
6674
assertThat(resultStore.getNodePropertyValues(List.of("A", "B"), "foo")).isNull();
6775
}
6876

77+
@Test
78+
void shouldEvictNodePropertyEntryAfter10Minutes() throws InterruptedException {
79+
var resultStore = new EphemeralResultStore();
80+
81+
var propertyValues = mock(NodePropertyValues.class);
82+
resultStore.addNodePropertyValues(List.of("A"), "prop", propertyValues);
83+
84+
assertThat(resultStore.getNodePropertyValues(List.of("A"), "prop")).isNotNull();
85+
86+
clock.forward(11, TimeUnit.MINUTES);
87+
// make some room for the cache eviction thread to trigger a cleanup
88+
Thread.sleep(100);
89+
90+
assertThat(resultStore.getNodePropertyValues(List.of("A"), "prop")).isNull();
91+
}
92+
6993
@Test
7094
void shouldStoreNodeLabel() {
7195
var resultStore = new EphemeralResultStore();
@@ -96,6 +120,23 @@ void shouldRemoveNodeLabel() {
96120
assertThat(resultStore.hasNodeLabel("Label")).isFalse();
97121
}
98122

123+
@Test
124+
void shouldEvictNodeLabelEntryAfter10Minutes() throws InterruptedException {
125+
var resultStore = new EphemeralResultStore();
126+
127+
var nodeCount = 1337L;
128+
var toOriginalId = mock(LongUnaryOperator.class);
129+
resultStore.addNodeLabel("Label", nodeCount, toOriginalId);
130+
131+
assertThat(resultStore.hasNodeLabel("Label")).isTrue();
132+
133+
clock.forward(11, TimeUnit.MINUTES);
134+
// make some room for the cache eviction thread to trigger a cleanup
135+
Thread.sleep(100);
136+
137+
assertThat(resultStore.hasNodeLabel("Label")).isFalse();
138+
}
139+
99140
@Test
100141
void shouldStoreGraphBasedRelationshipsWithoutProperty() {
101142
var resultStore = new EphemeralResultStore();
@@ -128,6 +169,24 @@ void shouldRemoveGraphBasedRelationshipsWithoutProperty() {
128169
assertThat(resultStore.hasRelationship("Type")).isFalse();
129170
}
130171

172+
@Test
173+
void shouldEvictGraphBasedRelationshipsWithoutPropertyAfter10Minutes() throws InterruptedException {
174+
var resultStore = new EphemeralResultStore();
175+
176+
var graph = mock(Graph.class);
177+
var toOriginalId = mock(LongUnaryOperator.class);
178+
179+
resultStore.addRelationship("Type", graph, toOriginalId);
180+
181+
assertThat(resultStore.hasRelationship("Type")).isTrue();
182+
183+
clock.forward(11, TimeUnit.MINUTES);
184+
// make some room for the cache eviction thread to trigger a cleanup
185+
Thread.sleep(100);
186+
187+
assertThat(resultStore.hasRelationship("Type")).isFalse();
188+
}
189+
131190
@Test
132191
void shouldStoreGraphBasedRelationshipsWithProperty() {
133192
var resultStore = new EphemeralResultStore();
@@ -160,6 +219,25 @@ void shouldRemoveGraphBasedRelationshipsWithProperty() {
160219
assertThat(resultStore.hasRelationship("Type", List.of("prop"))).isFalse();
161220
}
162221

222+
@Test
223+
void shouldEvictGraphBasedRelationshipsWithPropertyAfter10Minutes() throws InterruptedException {
224+
var resultStore = new EphemeralResultStore();
225+
226+
var graph = mock(Graph.class);
227+
var toOriginalId = mock(LongUnaryOperator.class);
228+
229+
resultStore.addRelationship("Type", "prop", graph, toOriginalId);
230+
231+
assertThat(resultStore.hasRelationship("Type", List.of("prop"))).isTrue();
232+
233+
234+
clock.forward(11, TimeUnit.MINUTES);
235+
// make some room for the cache eviction thread to trigger a cleanup
236+
Thread.sleep(100);
237+
238+
assertThat(resultStore.hasRelationship("Type", List.of("prop"))).isFalse();
239+
}
240+
163241
@Test
164242
void shouldStoreStreamBasedRelationships() {
165243
var resultStore = new EphemeralResultStore();
@@ -192,4 +270,22 @@ void shouldRemoveStreamBasedRelationships() {
192270

193271
assertThat(resultStore.hasRelationshipStream("Type", List.of("foo"))).isFalse();
194272
}
273+
274+
@Test
275+
void shouldEvictStreamBasedRelationshipsAfter10Minutes() throws InterruptedException {
276+
var resultStore = new EphemeralResultStore();
277+
278+
var relationshipStream = mock(Stream.class);
279+
var toOriginalId = mock(LongUnaryOperator.class);
280+
281+
resultStore.addRelationshipStream("Type", List.of("foo"), List.of(ValueType.DOUBLE), relationshipStream, toOriginalId);
282+
283+
assertThat(resultStore.hasRelationshipStream("Type", List.of("foo"))).isTrue();
284+
285+
clock.forward(11, TimeUnit.MINUTES);
286+
// make some room for the cache eviction thread to trigger a cleanup
287+
Thread.sleep(100);
288+
289+
assertThat(resultStore.hasRelationshipStream("Type", List.of("foo"))).isFalse();
290+
}
195291
}

0 commit comments

Comments
 (0)