4.x: Adjust BasicLoadBalancingPolicy#getReplicas (#405)

* Fetch additional keyspace metadata from scylla_keyspaces

Extends schema related classes with methods for querying scylla_keyspaces
in order to determine whether keyspace is tablets-enabled.

Adjusts schema queries tests and several other integration tests
that mock responses from system tables.

Adds `isUsingTablets()` and `setUsingTablets()` to KeyspaceMetadata.

* Use keyspace metadata in `BasicLoadBalancingPolicy#getReplicas`

Modifies the getReplicas method to do the tablet lookup if and only if the
keyspace metadata indicates that it's a tablets-based keyspace. Otherwise refer
to the token map.
Previous behavior was to try tablet map lookup first regardless of the
keyspace configuration.

Author: Bouncheck

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

@@ -40,6 +40,12 @@ public interface KeyspaceMetadata extends Describable {
   /** Whether this keyspace is virtual */
   boolean isVirtual();
 
+  default boolean isUsingTablets() {
+    return false;
+  }
+
+  default void setUsingTablets(boolean predicate) {}
+
   /** The replication options defined for this keyspace. */
   @NonNull
   Map<String, String> getReplication();

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

@@ -37,6 +37,7 @@
 import com.datastax.oss.driver.api.core.metadata.Tablet;
 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.api.core.metadata.token.Partitioner;
 import com.datastax.oss.driver.api.core.metadata.token.Token;
 import com.datastax.oss.driver.api.core.session.Request;
@@ -328,22 +329,24 @@ protected Set<Node> getReplicas(@Nullable Request request, @Nullable Session ses
       return Collections.emptySet();
     }
 
-    if (table != null) {
-      if (token == null) {
-        if (partitioner != null) {
-          token = partitioner.hash(key);
-        }
+    if (token == null && partitioner != null) {
+      token = partitioner.hash(key);
+    }
+
+    Optional<KeyspaceMetadata> ksMetadata =
+        context.getMetadataManager().getMetadata().getKeyspace(keyspace);
+    if (ksMetadata.isPresent() && ksMetadata.get().isUsingTablets()) {
+      if (table == null) {
+        return Collections.emptySet();
       }
       if (token instanceof TokenLong64) {
         Tablet targetTablet =
             tabletMap.getTablet(keyspace, table, ((TokenLong64) token).getValue());
         if (targetTablet != null) {
-          Set<Node> replicas = targetTablet.getReplicaNodes();
-          if (!replicas.isEmpty()) {
-            return replicas;
-          }
+          return targetTablet.getReplicaNodes();
         }
       }
+      return Collections.emptySet();
     }
 
     if (!maybeTokenMap.isPresent()) {

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

@@ -45,6 +45,7 @@ public class DefaultKeyspaceMetadata implements KeyspaceMetadata, Serializable {
   @NonNull private final Map<CqlIdentifier, ViewMetadata> views;
   @NonNull private final Map<FunctionSignature, FunctionMetadata> functions;
   @NonNull private final Map<FunctionSignature, AggregateMetadata> aggregates;
+  private boolean usingTablets = false;
 
   public DefaultKeyspaceMetadata(
       @NonNull CqlIdentifier name,
@@ -119,6 +120,16 @@ public Map<FunctionSignature, AggregateMetadata> getAggregates() {
     return aggregates;
   }
 
+  @Override
+  public boolean isUsingTablets() {
+    return this.usingTablets;
+  }
+
+  @Override
+  public void setUsingTablets(boolean predicate) {
+    this.usingTablets = predicate;
+  }
+
   @Override
   public boolean equals(Object other) {
     if (other == this) {

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/parsing/CassandraSchemaParser.java

@@ -75,6 +75,12 @@ public SchemaRefresh parse() {
     ImmutableMap.Builder<CqlIdentifier, KeyspaceMetadata> keyspacesBuilder = ImmutableMap.builder();
     for (AdminRow row : rows.keyspaces()) {
       KeyspaceMetadata keyspace = parseKeyspace(row);
+      AdminRow scyllaRow = rows.scyllaKeyspaces().getOrDefault(keyspace.getName(), null);
+      if (scyllaRow != null
+          && scyllaRow.contains("initial_tablets")
+          && !scyllaRow.isNull("initial_tablets")) {
+        keyspace.setUsingTablets(true);
+      }
       keyspacesBuilder.put(keyspace.getName(), keyspace);
     }
     for (AdminRow row : rows.virtualKeyspaces()) {

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/Cassandra21SchemaQueries.java

@@ -94,4 +94,9 @@ protected Optional<String> selectEdgesQuery() {
   protected Optional<String> selectVerticiesQuery() {
     return Optional.empty();
   }
+
+  @Override
+  protected Optional<String> selectScyllaKeyspacesQuery() {
+    return Optional.empty();
+  }
 }

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/Cassandra22SchemaQueries.java

@@ -94,4 +94,9 @@ protected Optional<String> selectEdgesQuery() {
   protected Optional<String> selectVerticiesQuery() {
     return Optional.empty();
   }
+
+  @Override
+  protected Optional<String> selectScyllaKeyspacesQuery() {
+    return Optional.empty();
+  }
 }

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/Cassandra3SchemaQueries.java

@@ -94,4 +94,9 @@ protected Optional<String> selectEdgesQuery() {
   protected Optional<String> selectVerticiesQuery() {
     return Optional.empty();
   }
+
+  @Override
+  protected Optional<String> selectScyllaKeyspacesQuery() {
+    return Optional.of("SELECT * FROM system_schema.scylla_keyspaces");
+  }
 }

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/CassandraSchemaQueries.java

@@ -28,10 +28,12 @@
 import com.datastax.oss.driver.internal.core.adminrequest.AdminRequestHandler;
 import com.datastax.oss.driver.internal.core.adminrequest.AdminResult;
 import com.datastax.oss.driver.internal.core.adminrequest.AdminRow;
+import com.datastax.oss.driver.internal.core.adminrequest.UnexpectedResponseException;
 import com.datastax.oss.driver.internal.core.channel.DriverChannel;
 import com.datastax.oss.driver.internal.core.util.NanoTime;
 import com.datastax.oss.driver.internal.core.util.concurrent.RunOrSchedule;
 import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
+import com.datastax.oss.protocol.internal.ProtocolConstants;
 import io.netty.util.concurrent.EventExecutor;
 import java.time.Duration;
 import java.util.Collections;
@@ -111,6 +113,8 @@ protected CassandraSchemaQueries(
 
   protected abstract Optional<String> selectVerticiesQuery();
 
+  protected abstract Optional<String> selectScyllaKeyspacesQuery();
+
   @Override
   public CompletionStage<SchemaRows> execute() {
     RunOrSchedule.on(adminExecutor, this::executeOnAdminExecutor);
@@ -125,6 +129,12 @@ private void executeOnAdminExecutor() {
     String usingClause = shouldApplyUsingTimeout() ? usingTimeoutClause : "";
 
     query(selectKeyspacesQuery() + whereClause + usingClause, schemaRowsBuilder::withKeyspaces);
+    selectScyllaKeyspacesQuery()
+        .ifPresent(
+            select ->
+                queryIfAvailable(
+                    select + whereClause + usingClause, schemaRowsBuilder::withScyllaKeyspaces));
+
     query(selectTypesQuery() + whereClause + usingClause, schemaRowsBuilder::withTypes);
     query(selectTablesQuery() + whereClause + usingClause, schemaRowsBuilder::withTables);
     query(selectColumnsQuery() + whereClause + usingClause, schemaRowsBuilder::withColumns);
@@ -176,6 +186,17 @@ private void query(
             (result, error) -> handleResult(result, error, builderUpdater), adminExecutor);
   }
 
+  private void queryIfAvailable(
+      String queryString,
+      Function<Iterable<AdminRow>, CassandraSchemaRows.Builder> builderUpdater) {
+    assert adminExecutor.inEventLoop();
+
+    pendingQueries += 1;
+    query(queryString)
+        .whenCompleteAsync(
+            (result, error) -> handleResult(result, error, builderUpdater, true), adminExecutor);
+  }
+
   @VisibleForTesting
   protected CompletionStage<AdminResult> query(String query) {
     return AdminRequestHandler.query(channel, query, timeout, pageSize, logPrefix).start();
@@ -185,12 +206,37 @@ private void handleResult(
       AdminResult result,
       Throwable error,
       Function<Iterable<AdminRow>, CassandraSchemaRows.Builder> builderUpdater) {
+    handleResult(result, error, builderUpdater, false);
+  }
+
+  private void handleResult(
+      AdminResult result,
+      Throwable error,
+      Function<Iterable<AdminRow>, CassandraSchemaRows.Builder> builderUpdater,
+      boolean ignoreServerErrors) {
 
     // If another query already failed, we've already propagated the failure so just ignore this one
     if (schemaRowsFuture.isCompletedExceptionally()) {
       return;
     }
 
+    // Meant to allow through "(keyspace/table) does not exist" or "unconfigured" errors for
+    // specific, optional queries
+    if (ignoreServerErrors && error instanceof UnexpectedResponseException) {
+      UnexpectedResponseException castedError = (UnexpectedResponseException) error;
+      if (castedError.message.opcode == ProtocolConstants.ErrorCode.SERVER_ERROR) {
+        LOG.debug("Silencing error: ", error);
+        // Consider such query 'done', but ignore its result
+        pendingQueries -= 1;
+        if (pendingQueries == 0) {
+          LOG.debug(
+              "[{}] Schema queries took {}", logPrefix, NanoTime.formatTimeSince(startTimeNs));
+          schemaRowsFuture.complete(schemaRowsBuilder.build());
+        }
+        return;
+      }
+    }
+
     if (error != null) {
       schemaRowsFuture.completeExceptionally(error);
     } else {
@@ -200,7 +246,8 @@ private void handleResult(
         result
             .nextPage()
             .whenCompleteAsync(
-                (nextResult, nextError) -> handleResult(nextResult, nextError, builderUpdater),
+                (nextResult, nextError) ->
+                    handleResult(nextResult, nextError, builderUpdater, ignoreServerErrors),
                 adminExecutor);
       } else {
         pendingQueries -= 1;

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/CassandraSchemaRows.java

@@ -56,6 +56,7 @@ public class CassandraSchemaRows implements SchemaRows {
   private final Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> indexes;
   private final Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> vertices;
   private final Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> edges;
+  private final Map<CqlIdentifier, AdminRow> scyllaKeyspaces;
 
   private CassandraSchemaRows(
       Node node,
@@ -72,7 +73,8 @@ private CassandraSchemaRows(
       Multimap<CqlIdentifier, AdminRow> functions,
       Multimap<CqlIdentifier, AdminRow> aggregates,
       Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> vertices,
-      Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> edges) {
+      Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> edges,
+      Map<CqlIdentifier, AdminRow> scyllaKeyspaces) {
     this.node = node;
     this.dataTypeParser = dataTypeParser;
     this.keyspaces = keyspaces;
@@ -88,6 +90,7 @@ private CassandraSchemaRows(
     this.aggregates = aggregates;
     this.vertices = vertices;
     this.edges = edges;
+    this.scyllaKeyspaces = scyllaKeyspaces;
   }
 
   @NonNull
@@ -166,6 +169,11 @@ public Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> edges() {
     return edges;
   }
 
+  @Override
+  public Map<CqlIdentifier, AdminRow> scyllaKeyspaces() {
+    return scyllaKeyspaces;
+  }
+
   public static class Builder {
     private static final Logger LOG = LoggerFactory.getLogger(Builder.class);
 
@@ -198,6 +206,8 @@ public static class Builder {
         verticesBuilders = new LinkedHashMap<>();
     private final Map<CqlIdentifier, ImmutableMultimap.Builder<CqlIdentifier, AdminRow>>
         edgesBuilders = new LinkedHashMap<>();
+    private final ImmutableMap.Builder<CqlIdentifier, AdminRow> scyllaKeyspacesBuilder =
+        ImmutableMap.builder();
 
     public Builder(Node node, KeyspaceFilter keyspaceFilter, String logPrefix) {
       this.node = node;
@@ -323,6 +333,13 @@ public Builder withEdges(Iterable<AdminRow> rows) {
       return this;
     }
 
+    public Builder withScyllaKeyspaces(Iterable<AdminRow> rows) {
+      for (AdminRow row : rows) {
+        putByKeyspacePk(row, scyllaKeyspacesBuilder);
+      }
+      return this;
+    }
+
     private void put(ImmutableList.Builder<AdminRow> builder, AdminRow row) {
       String keyspace = row.getString("keyspace_name");
       if (keyspace == null) {
@@ -342,6 +359,16 @@ private void putByKeyspace(
       }
     }
 
+    private void putByKeyspacePk(
+        AdminRow row, ImmutableMap.Builder<CqlIdentifier, AdminRow> builder) {
+      String keyspace = row.getString("keyspace_name");
+      if (keyspace == null) {
+        LOG.warn("[{}] Skipping system row with missing keyspace name", logPrefix);
+      } else if (keyspaceFilter.includes(keyspace)) {
+        builder.put(CqlIdentifier.fromInternal(keyspace), row);
+      }
+    }
+
     private void putByKeyspaceAndTable(
         AdminRow row,
         Map<CqlIdentifier, ImmutableMultimap.Builder<CqlIdentifier, AdminRow>> builders) {
@@ -375,7 +402,8 @@ public CassandraSchemaRows build() {
           functionsBuilder.build(),
           aggregatesBuilder.build(),
           build(verticesBuilders),
-          build(edgesBuilders));
+          build(edgesBuilders),
+          scyllaKeyspacesBuilder.build());
     }
 
     private static <K1, K2, V> Map<K1, Multimap<K2, V>> build(

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/SchemaRows.java

@@ -23,6 +23,7 @@
 import com.datastax.oss.driver.internal.core.metadata.schema.parsing.DataTypeParser;
 import com.datastax.oss.driver.shaded.guava.common.collect.Multimap;
 import edu.umd.cs.findbugs.annotations.NonNull;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -70,4 +71,8 @@ default Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> vertices() {
   default Map<CqlIdentifier, Multimap<CqlIdentifier, AdminRow>> edges() {
     return new LinkedHashMap<>();
   }
+
+  default Map<CqlIdentifier, AdminRow> scyllaKeyspaces() {
+    return new HashMap<>();
+  }
 }