diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/metadata/Metadata.java b/core/src/main/java/com/datastax/oss/driver/api/core/metadata/Metadata.java index 68e7bfeb677..2f5330dacdc 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/metadata/Metadata.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/metadata/Metadata.java @@ -123,7 +123,7 @@ default Optional getKeyspace(@NonNull String keyspaceName) { *

Starts as an empty map that will gradually receive updates on each query of a yet unknown * tablet. */ - TabletMap getTabletMap(); + Optional getTabletMap(); /** * The cluster name to which this session is connected. The Optional returned should contain the diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java index dd59a76f7d8..cc725994d7c 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java @@ -661,8 +661,7 @@ protected SchemaChangeListener buildSchemaChangeListener( .ifPresent(listeners::add); } if (getMetadataManager().isSchemaEnabled()) { - listeners.add( - new TabletMapSchemaChangeListener(getMetadataManager().getMetadata().getTabletMap())); + listeners.add(new TabletMapSchemaChangeListener(getMetadataManager())); } if (listeners.isEmpty()) { return new NoopSchemaChangeListener(this); diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index 99297f4266a..c1eee844647 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java @@ -292,7 +292,6 @@ private Token getRoutingToken(Statement statement) { } public Integer getShardFromTabletMap(Statement statement, Node node, Token token) { - TabletMap tabletMap = context.getMetadataManager().getMetadata().getTabletMap(); if (!(token instanceof TokenLong64)) { LOG.trace( "Token ({}) is not a TokenLong64. Not performing tablet shard lookup for statement {}.", @@ -300,6 +299,14 @@ public Integer getShardFromTabletMap(Statement statement, Node node, Token token statement); return null; } + TabletMap tabletMap = context.getMetadataManager().getMetadata().getTabletMap().orElse(null); + if (tabletMap == null) { + LOG.trace( + "Tablet map is not initialized. Not performing tablet shard lookup for statement {}.", + statement); + return null; + } + CqlIdentifier statementKeyspace = statement.getKeyspace(); if (statementKeyspace == null) { statementKeyspace = statement.getRoutingKeyspace(); diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java index d097425026b..927fb94bff6 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java @@ -301,7 +301,7 @@ protected Set getReplicas(@Nullable Request request, @Nullable Session ses } Optional maybeTokenMap = context.getMetadataManager().getMetadata().getTokenMap(); - TabletMap tabletMap = context.getMetadataManager().getMetadata().getTabletMap(); + Optional maybeTabletMap = context.getMetadataManager().getMetadata().getTabletMap(); // Note: we're on the hot path and the getXxx methods are potentially more than simple getters, // so we only call each method when strictly necessary (which is why the code below looks a bit @@ -345,13 +345,13 @@ protected Set getReplicas(@Nullable Request request, @Nullable Session ses Optional ksMetadata = context.getMetadataManager().getMetadata().getKeyspace(keyspace); - if (ksMetadata.isPresent() && ksMetadata.get().isUsingTablets()) { + if (ksMetadata.isPresent() && ksMetadata.get().isUsingTablets() && maybeTabletMap.isPresent()) { if (table == null) { return Collections.emptySet(); } if (token instanceof TokenLong64) { Tablet targetTablet = - tabletMap.getTablet(keyspace, table, ((TokenLong64) token).getValue()); + maybeTabletMap.get().getTablet(keyspace, table, ((TokenLong64) token).getValue()); if (targetTablet != null) { return targetTablet.getReplicaNodes(); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/DefaultMetadata.java b/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/DefaultMetadata.java index 253ba279dde..52d602da71e 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/DefaultMetadata.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/DefaultMetadata.java @@ -49,8 +49,7 @@ public class DefaultMetadata implements Metadata { private static final Logger LOG = LoggerFactory.getLogger(DefaultMetadata.class); public static DefaultMetadata EMPTY = - new DefaultMetadata( - Collections.emptyMap(), Collections.emptyMap(), null, null, DefaultTabletMap.emptyMap()); + new DefaultMetadata(Collections.emptyMap(), Collections.emptyMap(), null, null, null); protected final Map nodes; protected final Map keyspaces; @@ -98,8 +97,8 @@ public Optional getTokenMap() { } @Override - public TabletMap getTabletMap() { - return tabletMap; + public Optional getTabletMap() { + return Optional.ofNullable(tabletMap); } @NonNull @@ -137,7 +136,7 @@ public DefaultMetadata withNodes( rebuildTokenMap( newNodes, keyspaces, tokenMapEnabled, forceFullRebuild, tokenFactory, context), context.getChannelFactory().getClusterName(), - this.tabletMap); + this.tabletMap == null ? DefaultTabletMap.emptyMap() : this.tabletMap); } /** diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/TabletMapSchemaChangeListener.java b/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/TabletMapSchemaChangeListener.java index 4358ed90a8b..b764f351e19 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/TabletMapSchemaChangeListener.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/TabletMapSchemaChangeListener.java @@ -1,36 +1,48 @@ package com.datastax.oss.driver.internal.core.metadata.schema; -import com.datastax.oss.driver.api.core.metadata.TabletMap; import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata; import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListenerBase; import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata; +import com.datastax.oss.driver.internal.core.metadata.MetadataManager; import edu.umd.cs.findbugs.annotations.NonNull; public class TabletMapSchemaChangeListener extends SchemaChangeListenerBase { - private final TabletMap tabletMap; + private final MetadataManager manager; - public TabletMapSchemaChangeListener(TabletMap tabletMap) { - this.tabletMap = tabletMap; + public TabletMapSchemaChangeListener(MetadataManager manager) { + this.manager = manager; } @Override public void onKeyspaceDropped(@NonNull KeyspaceMetadata keyspace) { - tabletMap.removeByKeyspace(keyspace.getName()); + if (!manager.getMetadata().getTabletMap().isPresent()) { + return; + } + manager.getMetadata().getTabletMap().get().removeByKeyspace(keyspace.getName()); } @Override public void onKeyspaceUpdated( @NonNull KeyspaceMetadata current, @NonNull KeyspaceMetadata previous) { - tabletMap.removeByKeyspace(previous.getName()); + if (!manager.getMetadata().getTabletMap().isPresent()) { + return; + } + manager.getMetadata().getTabletMap().get().removeByKeyspace(previous.getName()); } @Override public void onTableDropped(@NonNull TableMetadata table) { - tabletMap.removeByTable(table.getName()); + if (!manager.getMetadata().getTabletMap().isPresent()) { + return; + } + manager.getMetadata().getTabletMap().get().removeByTable(table.getName()); } @Override public void onTableUpdated(@NonNull TableMetadata current, @NonNull TableMetadata previous) { - tabletMap.removeByTable(previous.getName()); + if (!manager.getMetadata().getTabletMap().isPresent()) { + return; + } + manager.getMetadata().getTabletMap().get().removeByTable(previous.getName()); } } diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java index 32a5b2b29c2..1612caef687 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java @@ -120,8 +120,9 @@ public void should_receive_each_tablet_exactly_once() { // With enough queries we should hit a wrong node for each tablet exactly once. Assert.assertEquals(INITIAL_TABLETS, counter); + Assert.assertTrue(session.getMetadata().getTabletMap().isPresent()); ConcurrentMap> tabletMapping = - session.getMetadata().getTabletMap().getMapping(); + session.getMetadata().getTabletMap().get().getMapping(); KeyspaceTableNamePair ktPair = new KeyspaceTableNamePair( CqlIdentifier.fromCql(KEYSPACE_NAME), CqlIdentifier.fromCql(TABLE_NAME)); diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/TabletMapSchemaChangesIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/TabletMapSchemaChangesIT.java index f6660f0bd6e..3e802a86ff6 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/TabletMapSchemaChangesIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/TabletMapSchemaChangesIT.java @@ -21,6 +21,7 @@ import com.datastax.oss.driver.internal.core.metadata.schema.TabletMapSchemaChangeListener; import java.time.Duration; import java.util.concurrent.TimeUnit; +import org.junit.Assert; import org.junit.Before; import org.junit.ClassRule; import org.junit.Test; @@ -118,12 +119,14 @@ public void setup() { .atMost(30, TimeUnit.SECONDS) .until( () -> - SESSION_RULE - .session() - .getMetadata() - .getTabletMap() - .getMapping() - .containsKey(TABLET_MAP_KEY)); + SESSION_RULE.session().getMetadata().getTabletMap().isPresent() + && SESSION_RULE + .session() + .getMetadata() + .getTabletMap() + .get() + .getMapping() + .containsKey(TABLET_MAP_KEY)); // Reset invocations for the next test method Mockito.clearInvocations(listener); } @@ -137,7 +140,8 @@ public void should_remove_tablets_on_keyspace_update() { Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1)) .onKeyspaceUpdated(Mockito.any(), previous.capture()); assertThat(previous.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(KEYSPACE_NAME)); - assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet()) + Assert.assertTrue(SESSION_RULE.session().getMetadata().getTabletMap().isPresent()); + assertThat(SESSION_RULE.session().getMetadata().getTabletMap().get().getMapping().keySet()) .doesNotContain(TABLET_MAP_KEY); } @@ -148,7 +152,8 @@ public void should_remove_tablets_on_keyspace_drop() { Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1)) .onKeyspaceDropped(keyspace.capture()); assertThat(keyspace.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(KEYSPACE_NAME)); - assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet()) + Assert.assertTrue(SESSION_RULE.session().getMetadata().getTabletMap().isPresent()); + assertThat(SESSION_RULE.session().getMetadata().getTabletMap().get().getMapping().keySet()) .doesNotContain(TABLET_MAP_KEY); } @@ -161,7 +166,8 @@ public void should_remove_tablets_on_table_update() { Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1)) .onTableUpdated(Mockito.any(), previous.capture()); assertThat(previous.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(TABLE_NAME)); - assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet()) + Assert.assertTrue(SESSION_RULE.session().getMetadata().getTabletMap().isPresent()); + assertThat(SESSION_RULE.session().getMetadata().getTabletMap().get().getMapping().keySet()) .doesNotContain(TABLET_MAP_KEY); } @@ -172,7 +178,8 @@ public void should_remove_tablets_on_table_drop() { Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1)) .onTableDropped(table.capture()); assertThat(table.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(TABLE_NAME)); - assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet()) + Assert.assertTrue(SESSION_RULE.session().getMetadata().getTabletMap().isPresent()); + assertThat(SESSION_RULE.session().getMetadata().getTabletMap().get().getMapping().keySet()) .doesNotContain(TABLET_MAP_KEY); } }