temporary

Author: Bouncheck

core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java

@@ -641,6 +641,13 @@ public enum DefaultDriverOption implements DriverOption {
    */
   METADATA_TOKEN_MAP_ENABLED("advanced.metadata.token-map.enabled"),
 
+  /**
+   * Whether tablet metadata is enabled.
+   *
+   * <p>Value-type: boolean
+   */
+  METADATA_TABLET_MAP_ENABLED("advanced.metadata.tablet-map.enabled"),
+
   /**
    * How long the driver waits for responses to control queries.
    *

core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java

@@ -354,6 +354,7 @@ protected static void fillWithDriverDefaults(OptionsMap map) {
     map.put(TypedDriverOption.METADATA_SCHEMA_WINDOW, Duration.ofSeconds(1));
     map.put(TypedDriverOption.METADATA_SCHEMA_MAX_EVENTS, 20);
     map.put(TypedDriverOption.METADATA_TOKEN_MAP_ENABLED, true);
+    map.put(TypedDriverOption.METADATA_TABLET_MAP_ENABLED, false);
     map.put(TypedDriverOption.CONTROL_CONNECTION_TIMEOUT, initQueryTimeout);
     map.put(TypedDriverOption.CONTROL_CONNECTION_AGREEMENT_INTERVAL, Duration.ofMillis(200));
     map.put(TypedDriverOption.CONTROL_CONNECTION_AGREEMENT_TIMEOUT, Duration.ofSeconds(10));

core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java

@@ -527,6 +527,9 @@ public String toString() {
   /** Whether token metadata is enabled. */
   public static final TypedDriverOption<Boolean> METADATA_TOKEN_MAP_ENABLED =
       new TypedDriverOption<>(DefaultDriverOption.METADATA_TOKEN_MAP_ENABLED, GenericType.BOOLEAN);
+  /** Whether tablet metadata is enabled. */
+  public static final TypedDriverOption<Boolean> METADATA_TABLET_MAP_ENABLED =
+      new TypedDriverOption<>(DefaultDriverOption.METADATA_TABLET_MAP_ENABLED, GenericType.BOOLEAN);
   /** How long the driver waits for responses to control queries. */
   public static final TypedDriverOption<Duration> CONTROL_CONNECTION_TIMEOUT =
       new TypedDriverOption<>(DefaultDriverOption.CONTROL_CONNECTION_TIMEOUT, GenericType.DURATION);

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

@@ -115,6 +115,17 @@ default Optional<KeyspaceMetadata> getKeyspace(@NonNull String keyspaceName) {
   @NonNull
   Optional<TokenMap> getTokenMap();
 
+  /**
+   * The tablet map for this cluster.
+   *
+   * <p>Note that this property might be absent if tablet metadata was disabled, or if there was a
+   * runtime error while computing the map (this would generate a warning log).
+   *
+   * @see DefaultDriverOption#NEW_OPTION
+   */
+  @NonNull
+  Optional<TabletMap> getTabletMap();
+
   /**
    * The cluster name to which this session is connected. The Optional returned should contain the
    * value from the server for <b>system.local.cluster_name</b>.

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

@@ -0,0 +1,12 @@
+package com.datastax.oss.driver.api.core.metadata;
+
+import com.datastax.oss.driver.internal.core.metadata.DefaultTabletMap;
+
+import java.util.Map;
+import java.util.NavigableSet;
+import java.util.Set;
+
+public interface TabletMap {
+  public Map<DefaultTabletMap.KeyspaceTableNamePair, NavigableSet<DefaultTabletMap.Tablet>> getRows();
+  public Set<String> getKeyspaces();
+}

core/src/main/java/com/datastax/oss/driver/internal/core/control/ControlConnection.java

@@ -499,6 +499,16 @@ private void onSuccessfulReconnect() {
                     }
                   }
                 });
+        context
+            .getMetadataManager()
+            .fullTabletMapRefresh()
+            .whenComplete(
+                (result, error) -> {
+                  if (error != null) {
+                    LOG.debug("[{}] Error while refreshing tablet map", logPrefix, error);
+                  }
+                }
+            );
       }
     }
 

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/DefaultMetadata.java

@@ -16,8 +16,10 @@
 package com.datastax.oss.driver.internal.core.metadata;
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
+import com.datastax.oss.driver.api.core.cql.Row;
 import com.datastax.oss.driver.api.core.metadata.Metadata;
 import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.metadata.TabletMap;
 import com.datastax.oss.driver.api.core.metadata.TokenMap;
 import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
@@ -30,6 +32,7 @@
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
@@ -46,21 +49,24 @@
 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);
+      new DefaultMetadata(Collections.emptyMap(), Collections.emptyMap(), null, null, null);
 
   protected final Map<UUID, Node> nodes;
   protected final Map<CqlIdentifier, KeyspaceMetadata> keyspaces;
   protected final TokenMap tokenMap;
+  protected final TabletMap tabletMap;
   protected final String clusterName;
 
   protected DefaultMetadata(
       Map<UUID, Node> nodes,
       Map<CqlIdentifier, KeyspaceMetadata> keyspaces,
       TokenMap tokenMap,
+      TabletMap tabletMap,
       String clusterName) {
     this.nodes = nodes;
     this.keyspaces = keyspaces;
     this.tokenMap = tokenMap;
+    this.tabletMap = tabletMap;
     this.clusterName = clusterName;
   }
 
@@ -82,6 +88,12 @@ public Optional<TokenMap> getTokenMap() {
     return Optional.ofNullable(tokenMap);
   }
 
+  @NonNull
+  @Override
+  public Optional<TabletMap> getTabletMap() {
+    return Optional.ofNullable(tabletMap);
+  }
+
   @NonNull
   @Override
   public Optional<String> getClusterName() {
@@ -114,6 +126,7 @@ public DefaultMetadata withNodes(
         this.keyspaces,
         rebuildTokenMap(
             newNodes, keyspaces, tokenMapEnabled, forceFullRebuild, tokenFactory, context),
+        tabletMap,
         context.getChannelFactory().getClusterName());
   }
 
@@ -125,9 +138,23 @@ public DefaultMetadata withSchema(
         this.nodes,
         ImmutableMap.copyOf(newKeyspaces),
         rebuildTokenMap(nodes, newKeyspaces, tokenMapEnabled, false, null, context),
+        tabletMap,
         context.getChannelFactory().getClusterName());
   }
 
+  public DefaultMetadata withTablets(
+      Iterable<DefaultTabletMap.Tablet> tablets,
+      boolean tabletMapEnabled,
+      InternalDriverContext context) {
+    return new DefaultMetadata(
+        this.nodes,
+        this.keyspaces,
+        this.tokenMap,
+        rebuildTabletMap(tablets, tabletMapEnabled, context),
+        context.getChannelFactory().getClusterName());
+  }
+
+
   @Nullable
   protected TokenMap rebuildTokenMap(
       Map<UUID, Node> newNodes,
@@ -188,4 +215,31 @@ protected TokenMap rebuildTokenMap(
       LOG.debug("[{}] Rebuilding token map took {}", logPrefix, NanoTime.formatTimeSince(start));
     }
   }
+
+  @Nullable
+  protected TabletMap rebuildTabletMap(
+      Iterable<DefaultTabletMap.Tablet> tablets,
+      boolean tabletMapEnabled,
+      InternalDriverContext context) {
+
+    String logPrefix = context.getSessionName();
+
+    if (!tabletMapEnabled) {
+      LOG.debug("[{}] Tablet map is disabled, skipping", logPrefix);
+      return this.tabletMap;
+    }
+    long start = System.nanoTime();
+    try {
+      return new DefaultTabletMap(tablets);
+    } catch (Throwable t) {
+      Loggers.warnWithException(
+          LOG,
+          "[{}] Unexpected error while refreshing tablet map, keeping previous version",
+          logPrefix,
+          t);
+      return this.tabletMap;
+    } finally {
+      LOG.debug("[{}] Rebuilding tablet map took {}", logPrefix, NanoTime.formatTimeSince(start));
+    }
+  }
 }