Make DefaultTabletMap initialize tabletMap later

core/src/main/java/com/datastax/oss/driver/api/core/metadata/Metadata.java

@@ -123,7 +123,7 @@ default Optional<KeyspaceMetadata> getKeyspace(@NonNull String keyspaceName) {
    * <p>Starts as an empty map that will gradually receive updates on each query of a yet unknown
    * tablet.
    */
-  TabletMap getTabletMap();
+  Optional<TabletMap> getTabletMap();
 
   /**
    * The cluster name to which this session is connected. The Optional returned should contain the

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);

core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java

@@ -292,14 +292,21 @@ 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 {}.",
           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();

core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java

@@ -301,7 +301,7 @@ protected Set<Node> getReplicas(@Nullable Request request, @Nullable Session ses
     }
 
     Optional<TokenMap> maybeTokenMap = context.getMetadataManager().getMetadata().getTokenMap();
-    TabletMap tabletMap = context.getMetadataManager().getMetadata().getTabletMap();
+    Optional<TabletMap> 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<Node> getReplicas(@Nullable Request request, @Nullable Session ses
 
     Optional<KeyspaceMetadata> 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();
         }

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<UUID, Node> nodes;
   protected final Map<CqlIdentifier, KeyspaceMetadata> keyspaces;
@@ -98,8 +97,8 @@ public Optional<TokenMap> getTokenMap() {
   }
 
   @Override
-  public TabletMap getTabletMap() {
-    return tabletMap;
+  public Optional<TabletMap> 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);
   }
 
   /**

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());
   }
 }

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<KeyspaceTableNamePair, ConcurrentSkipListSet<Tablet>> tabletMapping =
-        session.getMetadata().getTabletMap().getMapping();
+        session.getMetadata().getTabletMap().get().getMapping();
     KeyspaceTableNamePair ktPair =
         new KeyspaceTableNamePair(
             CqlIdentifier.fromCql(KEYSPACE_NAME), CqlIdentifier.fromCql(TABLE_NAME));

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);
   }
 }