Skip to content

Commit f6cc48d

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 tokenMap, 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 f6cc48d

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;
3+
import com.datastax.oss.driver.api.core.metadata.Metadata;
44
import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
55
import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListenerBase;
66
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
77
import edu.umd.cs.findbugs.annotations.NonNull;
88

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

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

1616
@Override
1717
public void onKeyspaceDropped(@NonNull KeyspaceMetadata keyspace) {
18-
tabletMap.removeByKeyspace(keyspace.getName());
18+
if (!metadata.getTabletMap().isPresent()) {
19+
return;
20+
}
21+
metadata.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 (!metadata.getTabletMap().isPresent()) {
28+
return;
29+
}
30+
metadata.getTabletMap().get().removeByKeyspace(previous.getName());
2531
}
2632

2733
@Override
2834
public void onTableDropped(@NonNull TableMetadata table) {
29-
tabletMap.removeByTable(table.getName());
35+
if (!metadata.getTabletMap().isPresent()) {
36+
return;
37+
}
38+
metadata.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 (!metadata.getTabletMap().isPresent()) {
44+
return;
45+
}
46+
metadata.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)