Skip to content

Commit ed0eab6

Browse files
dkropachevdkropachev
committed
Make DefaultTabletMap initialize tabletMap later
Currently `tabletMap` is initialized once at `DefaultMetadata.EMPTY` As result all the session instances have same `tabletMap`, which can lead to all sort of problems. This commit postpone tabletMap initialization and changes Metadata API to reflect these changes.
1 parent 2a62541 commit ed0eab6

File tree

8 files changed

+56
-31
lines changed

8 files changed

+56
-31
lines changed

Diff for: ‎core/src/main/java/com/datastax/oss/driver/api/core/metadata/Metadata.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ default Optional<KeyspaceMetadata> getKeyspace(@NonNull String keyspaceName) {
123123
* <p>Starts as an empty map that will gradually receive updates on each query of a yet unknown
124124
* tablet.
125125
*/
126-
TabletMap getTabletMap();
126+
Optional<TabletMap> getTabletMap();
127127

128128
/**
129129
* The cluster name to which this session is connected. The Optional returned should contain the

Diff for: ‎core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java

+1-2
Original file line numberDiff line numberDiff line change
@@ -661,8 +661,7 @@ protected SchemaChangeListener buildSchemaChangeListener(
661661
.ifPresent(listeners::add);
662662
}
663663
if (getMetadataManager().isSchemaEnabled()) {
664-
listeners.add(
665-
new TabletMapSchemaChangeListener(getMetadataManager().getMetadata().getTabletMap()));
664+
listeners.add(new TabletMapSchemaChangeListener(getMetadataManager().getMetadata()));
666665
}
667666
if (listeners.isEmpty()) {
668667
return new NoopSchemaChangeListener(this);

Diff for: ‎core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java

+8-1
Original file line numberDiff line numberDiff line change
@@ -292,14 +292,21 @@ private Token getRoutingToken(Statement statement) {
292292
}
293293

294294
public Integer getShardFromTabletMap(Statement statement, Node node, Token token) {
295-
TabletMap tabletMap = context.getMetadataManager().getMetadata().getTabletMap();
296295
if (!(token instanceof TokenLong64)) {
297296
LOG.trace(
298297
"Token ({}) is not a TokenLong64. Not performing tablet shard lookup for statement {}.",
299298
token,
300299
statement);
301300
return null;
302301
}
302+
TabletMap tabletMap = context.getMetadataManager().getMetadata().getTabletMap().orElse(null);
303+
if (tabletMap == null) {
304+
LOG.trace(
305+
"Tablet map is not initialized. Not performing tablet shard lookup for statement {}.",
306+
statement);
307+
return null;
308+
}
309+
303310
CqlIdentifier statementKeyspace = statement.getKeyspace();
304311
if (statementKeyspace == null) {
305312
statementKeyspace = statement.getRoutingKeyspace();

Diff for: ‎core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -301,7 +301,7 @@ protected Set<Node> getReplicas(@Nullable Request request, @Nullable Session ses
301301
}
302302

303303
Optional<TokenMap> maybeTokenMap = context.getMetadataManager().getMetadata().getTokenMap();
304-
TabletMap tabletMap = context.getMetadataManager().getMetadata().getTabletMap();
304+
Optional<TabletMap> maybeTabletMap = context.getMetadataManager().getMetadata().getTabletMap();
305305

306306
// Note: we're on the hot path and the getXxx methods are potentially more than simple getters,
307307
// so we only call each method when strictly necessary (which is why the code below looks a bit
@@ -345,13 +345,13 @@ protected Set<Node> getReplicas(@Nullable Request request, @Nullable Session ses
345345

346346
Optional<KeyspaceMetadata> ksMetadata =
347347
context.getMetadataManager().getMetadata().getKeyspace(keyspace);
348-
if (ksMetadata.isPresent() && ksMetadata.get().isUsingTablets()) {
348+
if (ksMetadata.isPresent() && ksMetadata.get().isUsingTablets() && maybeTabletMap.isPresent()) {
349349
if (table == null) {
350350
return Collections.emptySet();
351351
}
352352
if (token instanceof TokenLong64) {
353353
Tablet targetTablet =
354-
tabletMap.getTablet(keyspace, table, ((TokenLong64) token).getValue());
354+
maybeTabletMap.get().getTablet(keyspace, table, ((TokenLong64) token).getValue());
355355
if (targetTablet != null) {
356356
return targetTablet.getReplicaNodes();
357357
}

Diff for: ‎core/src/main/java/com/datastax/oss/driver/internal/core/metadata/DefaultMetadata.java

+4-5
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,7 @@
4949
public class DefaultMetadata implements Metadata {
5050
private static final Logger LOG = LoggerFactory.getLogger(DefaultMetadata.class);
5151
public static DefaultMetadata EMPTY =
52-
new DefaultMetadata(
53-
Collections.emptyMap(), Collections.emptyMap(), null, null, DefaultTabletMap.emptyMap());
52+
new DefaultMetadata(Collections.emptyMap(), Collections.emptyMap(), null, null, null);
5453

5554
protected final Map<UUID, Node> nodes;
5655
protected final Map<CqlIdentifier, KeyspaceMetadata> keyspaces;
@@ -98,8 +97,8 @@ public Optional<TokenMap> getTokenMap() {
9897
}
9998

10099
@Override
101-
public TabletMap getTabletMap() {
102-
return tabletMap;
100+
public Optional<TabletMap> getTabletMap() {
101+
return Optional.ofNullable(tabletMap);
103102
}
104103

105104
@NonNull
@@ -137,7 +136,7 @@ public DefaultMetadata withNodes(
137136
rebuildTokenMap(
138137
newNodes, keyspaces, tokenMapEnabled, forceFullRebuild, tokenFactory, context),
139138
context.getChannelFactory().getClusterName(),
140-
this.tabletMap);
139+
this.tabletMap == null ? DefaultTabletMap.emptyMap() : this.tabletMap);
141140
}
142141

143142
/**

Diff for: ‎core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/TabletMapSchemaChangeListener.java

+20-8
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,48 @@
11
package com.datastax.oss.driver.internal.core.metadata.schema;
22

3-
import com.datastax.oss.driver.api.core.metadata.TabletMap;
43
import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
54
import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListenerBase;
65
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
6+
import com.datastax.oss.driver.internal.core.metadata.MetadataManager;
77
import edu.umd.cs.findbugs.annotations.NonNull;
88

99
public class TabletMapSchemaChangeListener extends SchemaChangeListenerBase {
10-
private final TabletMap tabletMap;
10+
private final MetadataManager manager;
1111

12-
public TabletMapSchemaChangeListener(TabletMap tabletMap) {
13-
this.tabletMap = tabletMap;
12+
public TabletMapSchemaChangeListener(MetadataManager manager) {
13+
this.manager = manager;
1414
}
1515

1616
@Override
1717
public void onKeyspaceDropped(@NonNull KeyspaceMetadata keyspace) {
18-
tabletMap.removeByKeyspace(keyspace.getName());
18+
if (!manager.getMetadata().getTabletMap().isPresent()) {
19+
return;
20+
}
21+
manager.getMetadata().getTabletMap().get().removeByKeyspace(keyspace.getName());
1922
}
2023

2124
@Override
2225
public void onKeyspaceUpdated(
2326
@NonNull KeyspaceMetadata current, @NonNull KeyspaceMetadata previous) {
24-
tabletMap.removeByKeyspace(previous.getName());
27+
if (!manager.getMetadata().getTabletMap().isPresent()) {
28+
return;
29+
}
30+
manager.getMetadata().getTabletMap().get().removeByKeyspace(previous.getName());
2531
}
2632

2733
@Override
2834
public void onTableDropped(@NonNull TableMetadata table) {
29-
tabletMap.removeByTable(table.getName());
35+
if (!manager.getMetadata().getTabletMap().isPresent()) {
36+
return;
37+
}
38+
manager.getMetadata().getTabletMap().get().removeByTable(table.getName());
3039
}
3140

3241
@Override
3342
public void onTableUpdated(@NonNull TableMetadata current, @NonNull TableMetadata previous) {
34-
tabletMap.removeByTable(previous.getName());
43+
if (!manager.getMetadata().getTabletMap().isPresent()) {
44+
return;
45+
}
46+
manager.getMetadata().getTabletMap().get().removeByTable(previous.getName());
3547
}
3648
}

Diff for: ‎integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,9 @@ public void should_receive_each_tablet_exactly_once() {
120120
// With enough queries we should hit a wrong node for each tablet exactly once.
121121
Assert.assertEquals(INITIAL_TABLETS, counter);
122122

123+
Assert.assertTrue(session.getMetadata().getTabletMap().isPresent());
123124
ConcurrentMap<KeyspaceTableNamePair, ConcurrentSkipListSet<Tablet>> tabletMapping =
124-
session.getMetadata().getTabletMap().getMapping();
125+
session.getMetadata().getTabletMap().get().getMapping();
125126
KeyspaceTableNamePair ktPair =
126127
new KeyspaceTableNamePair(
127128
CqlIdentifier.fromCql(KEYSPACE_NAME), CqlIdentifier.fromCql(TABLE_NAME));

Diff for: ‎integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/TabletMapSchemaChangesIT.java

+17-10
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import com.datastax.oss.driver.internal.core.metadata.schema.TabletMapSchemaChangeListener;
2222
import java.time.Duration;
2323
import java.util.concurrent.TimeUnit;
24+
import org.junit.Assert;
2425
import org.junit.Before;
2526
import org.junit.ClassRule;
2627
import org.junit.Test;
@@ -118,12 +119,14 @@ public void setup() {
118119
.atMost(30, TimeUnit.SECONDS)
119120
.until(
120121
() ->
121-
SESSION_RULE
122-
.session()
123-
.getMetadata()
124-
.getTabletMap()
125-
.getMapping()
126-
.containsKey(TABLET_MAP_KEY));
122+
SESSION_RULE.session().getMetadata().getTabletMap().isPresent()
123+
&& SESSION_RULE
124+
.session()
125+
.getMetadata()
126+
.getTabletMap()
127+
.get()
128+
.getMapping()
129+
.containsKey(TABLET_MAP_KEY));
127130
// Reset invocations for the next test method
128131
Mockito.clearInvocations(listener);
129132
}
@@ -137,7 +140,8 @@ public void should_remove_tablets_on_keyspace_update() {
137140
Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1))
138141
.onKeyspaceUpdated(Mockito.any(), previous.capture());
139142
assertThat(previous.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(KEYSPACE_NAME));
140-
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet())
143+
Assert.assertTrue(SESSION_RULE.session().getMetadata().getTabletMap().isPresent());
144+
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().get().getMapping().keySet())
141145
.doesNotContain(TABLET_MAP_KEY);
142146
}
143147

@@ -148,7 +152,8 @@ public void should_remove_tablets_on_keyspace_drop() {
148152
Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1))
149153
.onKeyspaceDropped(keyspace.capture());
150154
assertThat(keyspace.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(KEYSPACE_NAME));
151-
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet())
155+
Assert.assertTrue(SESSION_RULE.session().getMetadata().getTabletMap().isPresent());
156+
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().get().getMapping().keySet())
152157
.doesNotContain(TABLET_MAP_KEY);
153158
}
154159

@@ -161,7 +166,8 @@ public void should_remove_tablets_on_table_update() {
161166
Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1))
162167
.onTableUpdated(Mockito.any(), previous.capture());
163168
assertThat(previous.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(TABLE_NAME));
164-
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet())
169+
Assert.assertTrue(SESSION_RULE.session().getMetadata().getTabletMap().isPresent());
170+
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().get().getMapping().keySet())
165171
.doesNotContain(TABLET_MAP_KEY);
166172
}
167173

@@ -172,7 +178,8 @@ public void should_remove_tablets_on_table_drop() {
172178
Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1))
173179
.onTableDropped(table.capture());
174180
assertThat(table.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(TABLE_NAME));
175-
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet())
181+
Assert.assertTrue(SESSION_RULE.session().getMetadata().getTabletMap().isPresent());
182+
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().get().getMapping().keySet())
176183
.doesNotContain(TABLET_MAP_KEY);
177184
}
178185
}

0 commit comments

Comments
 (0)