[#596] Implementation of SCYLLA_USE_METADATA_ID protocol feature negotiation. Introducing ProtocolFeatureStore.

[#596] Initial implementation of SCYLLA_USE_METADATA_ID feature.

[#596] Fixing NPE when accessing sharding info from empty ProtocolFeatureStore.

[#596] Storing `ProtocolFeatureStore` in `ProtocolEncoder`/`ProtocolDecoder`. Getting read of sharing it via `Channel` attributes.

[#596] Addressing missing `ProtocolFeatureStore` in a `statement` param of `executeAsync` to fix ITs.

[#596] Adjusting the way feature store is handled in `DefaultResultSetFuture`. Removing feature store from its constructor.

[#596] Moving channel attribute management for feature store to utility methods of `ProtocolFeatureStore`.

 [#596] Setting `Host`'s feature store separately from `Channel`. Code clean up

Apply suggestion from @dkropachev

Co-authored-by: Dmitry Kropachev <[email protected]>

 [#596] `ProtocolEncoder` clean up.

Author: nikagra

driver-core/src/main/java/com/datastax/driver/core/ArrayBackedResultSet.java

@@ -68,7 +68,8 @@ static ArrayBackedResultSet fromMessage(
       SessionManager session,
       ProtocolVersion protocolVersion,
       ExecutionInfo info,
-      Statement statement) {
+      Statement statement,
+      ProtocolFeatureStore featureStore) {
 
     switch (msg.kind) {
       case ROWS:
@@ -94,7 +95,8 @@ static ArrayBackedResultSet fromMessage(
           // CASSANDRA-10786).
           MD5Digest newMetadataId = r.metadata.metadataId;
           assert !(actualStatement instanceof BoundStatement)
-              || ProtocolFeature.PREPARED_METADATA_CHANGES.isSupportedBy(protocolVersion)
+              || ProtocolFeatures.PREPARED_METADATA_CHANGES.isSupportedBy(
+                  protocolVersion, featureStore)
               || newMetadataId == null;
           if (newMetadataId != null) {
             BoundStatement bs = ((BoundStatement) actualStatement);

driver-core/src/main/java/com/datastax/driver/core/BatchStatement.java

@@ -228,10 +228,10 @@ public int requestSizeInBytes(ProtocolVersion protocolVersion, CodecRegistry cod
           // overestimate by a
           // few bytes.
           size += CBUtil.sizeOfConsistencyLevel(getSerialConsistencyLevel());
-          if (ProtocolFeature.CLIENT_TIMESTAMPS.isSupportedBy(protocolVersion)) {
+          if (ProtocolFeatures.CLIENT_TIMESTAMPS.isSupportedBy(protocolVersion)) {
             size += 8; // timestamp
           }
-          if (ProtocolFeature.CUSTOM_PAYLOADS.isSupportedBy(protocolVersion)
+          if (ProtocolFeatures.CUSTOM_PAYLOADS.isSupportedBy(protocolVersion)
               && getOutgoingPayload() != null) {
             size += CBUtil.sizeOfBytesMap(getOutgoingPayload());
           }

driver-core/src/main/java/com/datastax/driver/core/BoundStatement.java

@@ -325,7 +325,8 @@ public int requestSizeInBytes(ProtocolVersion protocolVersion, CodecRegistry cod
     try {
       size +=
           CBUtil.sizeOfShortBytes(preparedStatement().getPreparedId().boundValuesMetadata.id.bytes);
-      if (ProtocolFeature.PREPARED_METADATA_CHANGES.isSupportedBy(protocolVersion)) {
+      ProtocolFeatureStore featureStore = getHost().getProtocolFeatureStore();
+      if (ProtocolFeatures.PREPARED_METADATA_CHANGES.isSupportedBy(protocolVersion, featureStore)) {
         size +=
             CBUtil.sizeOfShortBytes(preparedStatement().getPreparedId().resultSetMetadata.id.bytes);
       }
@@ -353,10 +354,10 @@ public int requestSizeInBytes(ProtocolVersion protocolVersion, CodecRegistry cod
             size += CBUtil.sizeOfValue(getPagingState());
           }
           size += CBUtil.sizeOfConsistencyLevel(getSerialConsistencyLevel());
-          if (ProtocolFeature.CLIENT_TIMESTAMPS.isSupportedBy(protocolVersion)) {
+          if (ProtocolFeatures.CLIENT_TIMESTAMPS.isSupportedBy(protocolVersion)) {
             size += 8; // timestamp
           }
-          if (ProtocolFeature.CUSTOM_PAYLOADS.isSupportedBy(protocolVersion)
+          if (ProtocolFeatures.CUSTOM_PAYLOADS.isSupportedBy(protocolVersion)
               && getOutgoingPayload() != null) {
             size += CBUtil.sizeOfBytesMap(getOutgoingPayload());
           }

driver-core/src/main/java/com/datastax/driver/core/Connection.java

@@ -155,6 +155,7 @@ enum State {
 
   private final AtomicReference<Owner> ownerRef = new AtomicReference<Owner>();
   private final ApplicationInfo applicationInfo;
+  private ProtocolFeatureStore protocolFeatureStore;
 
   /**
    * Create a new connection to a Cassandra node and associate it with the given pool.
@@ -449,35 +450,31 @@ private AsyncFunction<Message.Response, Void> onOptionsResponse(
       final ProtocolVersion protocolVersion, final Executor initExecutor) {
     return new AsyncFunction<Message.Response, Void>() {
       @Override
-      public ListenableFuture<Void> apply(Message.Response response) throws Exception {
+      public ListenableFuture<Void> apply(Message.Response response) {
         switch (response.type) {
           case SUPPORTED:
-            Responses.Supported msg = (Supported) response;
-            ShardingInfo.ConnectionShardingInfo sharding =
-                ShardingInfo.parseShardingInfo(msg.supported);
-            if (sharding != null) {
-              getHost().setShardingInfo(sharding.shardingInfo);
-              Connection.this.shardId = sharding.shardId;
+            Supported supported = (Supported) response;
+            protocolFeatureStore = ProtocolFeatureStore.parseSupportedOptions(supported.supported);
+            protocolFeatureStore.storeInChannel(channel);
+            getHost().setProtocolFeatureStore(protocolFeatureStore);
+
+            ShardingInfo.ConnectionShardingInfo shardingInfo =
+                protocolFeatureStore.getConnectionShardingInfo();
+            if (protocolFeatureStore.getConnectionShardingInfo() != null) {
+              Connection.this.shardId = shardingInfo.shardId;
               if (Connection.this.requestedShardId != -1
-                  && Connection.this.requestedShardId != sharding.shardId) {
+                  && Connection.this.requestedShardId != shardingInfo.shardId) {
                 logger.warn(
                     "Advanced shard awareness: requested connection to shard {}, but connected to {}. Is there a NAT between client and server?",
                     Connection.this.requestedShardId,
-                    sharding.shardId);
+                    shardingInfo.shardId);
                 // Owner is a HostConnectionPool if we are using adv. shard awareness
                 ((HostConnectionPool) Connection.this.ownerRef.get())
                     .tempBlockAdvShardAwareness(ADV_SHARD_AWARENESS_BLOCK_ON_NAT);
               }
             } else {
-              getHost().setShardingInfo(null);
               Connection.this.shardId = 0;
             }
-            LwtInfo lwt = LwtInfo.parseLwtInfo(msg.supported);
-            if (lwt != null) {
-              getHost().setLwtInfo(lwt);
-            }
-            TabletInfo tabletInfo = TabletInfo.parseTabletInfo(msg.supported);
-            getHost().setTabletInfo(tabletInfo);
             return MoreFutures.VOID_SUCCESS;
           case ERROR:
             Responses.Error error = (Responses.Error) response;
@@ -506,20 +503,13 @@ private AsyncFunction<Void, Void> onOptionsReady(
       @Override
       public ListenableFuture<Void> apply(Void input) throws Exception {
         ProtocolOptions protocolOptions = factory.configuration.getProtocolOptions();
-        Map<String, String> extraOptions = new HashMap<String, String>();
+        Map<String, String> extraOptions = new HashMap<>();
         if (applicationInfo != null) {
           applicationInfo.addOption(extraOptions);
         }
-        LwtInfo lwtInfo = getHost().getLwtInfo();
-        if (lwtInfo != null) {
-          lwtInfo.addOption(extraOptions);
-        }
-        TabletInfo tabletInfo = getHost().getTabletInfo();
-        if (tabletInfo != null
-            && tabletInfo.isEnabled()
-            && ProtocolFeature.CUSTOM_PAYLOADS.isSupportedBy(protocolVersion)) {
-          logger.debug("Enabling tablet support in OPTIONS message");
-          TabletInfo.addOption(extraOptions);
+
+        if (protocolFeatureStore != null) {
+          protocolFeatureStore.populateStartupOptions(protocolVersion, extraOptions);
         }
 
         Future startupResponseFuture =
@@ -1065,6 +1055,10 @@ public int shardId() {
     return shardId == null ? 0 : shardId;
   }
 
+  public ProtocolFeatureStore getProtocolFeatureStore() {
+    return protocolFeatureStore;
+  }
+
   /**
    * If the connection is part of a pool, return it to the pool. The connection should generally not
    * be reused after that.
@@ -1955,21 +1949,6 @@ interface DefaultResponseHandler {
   }
 
   private static class Initializer extends ChannelInitializer<SocketChannel> {
-    // Stateless handlers
-    private static final Message.ProtocolDecoder messageDecoder = new Message.ProtocolDecoder();
-    private static final Message.ProtocolEncoder messageEncoderV1 =
-        new Message.ProtocolEncoder(ProtocolVersion.V1);
-    private static final Message.ProtocolEncoder messageEncoderV2 =
-        new Message.ProtocolEncoder(ProtocolVersion.V2);
-    private static final Message.ProtocolEncoder messageEncoderV3 =
-        new Message.ProtocolEncoder(ProtocolVersion.V3);
-    private static final Message.ProtocolEncoder messageEncoderV4 =
-        new Message.ProtocolEncoder(ProtocolVersion.V4);
-    private static final Message.ProtocolEncoder messageEncoderV5 =
-        new Message.ProtocolEncoder(ProtocolVersion.V5);
-    private static final Message.ProtocolEncoder messageEncoderV6 =
-        new Message.ProtocolEncoder(ProtocolVersion.V6);
-    private static final Frame.Encoder frameEncoder = new Frame.Encoder();
 
     private final ProtocolVersion protocolVersion;
     private final Connection connection;
@@ -2033,7 +2012,7 @@ protected void initChannel(SocketChannel channel) throws Exception {
       }
 
       pipeline.addLast("frameDecoder", new Frame.Decoder());
-      pipeline.addLast("frameEncoder", frameEncoder);
+      pipeline.addLast("frameEncoder", new Frame.Encoder());
 
       pipeline.addLast("framingFormatHandler", new FramingFormatHandler(connection.factory));
 
@@ -2046,7 +2025,7 @@ protected void initChannel(SocketChannel channel) throws Exception {
         pipeline.addLast("frameCompressor", new Frame.Compressor(compressor));
       }
 
-      pipeline.addLast("messageDecoder", messageDecoder);
+      pipeline.addLast("messageDecoder", new Message.ProtocolDecoder(null));
       pipeline.addLast("messageEncoder", messageEncoderFor(protocolVersion));
 
       pipeline.addLast("idleStateHandler", idleStateHandler);
@@ -2056,23 +2035,11 @@ protected void initChannel(SocketChannel channel) throws Exception {
       nettyOptions.afterChannelInitialized(channel);
     }
 
-    private Message.ProtocolEncoder messageEncoderFor(ProtocolVersion version) {
-      switch (version) {
-        case V1:
-          return messageEncoderV1;
-        case V2:
-          return messageEncoderV2;
-        case V3:
-          return messageEncoderV3;
-        case V4:
-          return messageEncoderV4;
-        case V5:
-          return messageEncoderV5;
-        case V6:
-          return messageEncoderV6;
-        default:
-          throw new DriverInternalError("Unsupported protocol version " + protocolVersion);
+    private static Message.ProtocolEncoder messageEncoderFor(ProtocolVersion version) {
+      if (version.toInt() > ProtocolVersion.V6.toInt()) {
+        throw new DriverInternalError("Unsupported protocol version " + version);
       }
+      return new Message.ProtocolEncoder(version, null);
     }
   }
 

driver-core/src/main/java/com/datastax/driver/core/DefaultResultSetFuture.java

@@ -87,15 +87,29 @@ public void onSet(
                     table,
                     rm.getCustomPayload().get(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY));
           }
+
           switch (rm.kind) {
             case SET_KEYSPACE:
               // propagate the keyspace change to other connections
               session.poolsState.setKeyspace(((Responses.Result.SetKeyspace) rm).keyspace);
-              set(ArrayBackedResultSet.fromMessage(rm, session, protocolVersion, info, statement));
+              set(
+                  ArrayBackedResultSet.fromMessage(
+                      rm,
+                      session,
+                      protocolVersion,
+                      info,
+                      statement,
+                      connection.getProtocolFeatureStore()));
               break;
             case SCHEMA_CHANGE:
               ResultSet rs =
-                  ArrayBackedResultSet.fromMessage(rm, session, protocolVersion, info, statement);
+                  ArrayBackedResultSet.fromMessage(
+                      rm,
+                      session,
+                      protocolVersion,
+                      info,
+                      statement,
+                      connection.getProtocolFeatureStore());
               final Cluster.Manager cluster = session.cluster.manager;
               if (!cluster.configuration.getQueryOptions().isMetadataEnabled()) {
                 cluster.waitForSchemaAgreementAndSignal(connection, this, rs);
@@ -224,7 +238,14 @@ public void run() {
               }
               break;
             default:
-              set(ArrayBackedResultSet.fromMessage(rm, session, protocolVersion, info, statement));
+              set(
+                  ArrayBackedResultSet.fromMessage(
+                      rm,
+                      session,
+                      protocolVersion,
+                      info,
+                      statement,
+                      connection.getProtocolFeatureStore()));
               break;
           }
           break;

driver-core/src/main/java/com/datastax/driver/core/FramingFormatHandler.java

@@ -16,6 +16,7 @@
 package com.datastax.driver.core;
 
 import com.datastax.driver.core.Message.Response.Type;
+import com.datastax.driver.core.exceptions.DriverInternalError;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPipeline;
 import io.netty.handler.codec.MessageToMessageDecoder;
@@ -43,8 +44,11 @@ protected void decode(ChannelHandlerContext ctx, Frame frame, List<Object> out)
       // By default, the pipeline is configured for legacy framing since this is the format used
       // by all protocol versions until handshake; after handshake however, we need to switch to
       // modern framing for protocol v5 and higher.
+      ProtocolFeatureStore featureStore = ProtocolFeatureStore.loadFromChannel(ctx.channel());
       if (frame.header.version.compareTo(ProtocolVersion.V5) >= 0) {
         switchToModernFraming(ctx);
+      } else if (featureStore != null && featureStore.isUseMetadataId()) {
+        switchToCQL4MetadataId(ctx, frame.header.version, featureStore);
       }
       // once the handshake is successful, the framing format cannot change anymore;
       // we can safely remove ourselves from the pipeline.
@@ -53,6 +57,16 @@ protected void decode(ChannelHandlerContext ctx, Frame frame, List<Object> out)
     out.add(frame);
   }
 
+  private void switchToCQL4MetadataId(
+      ChannelHandlerContext ctx,
+      ProtocolVersion protocolVersion,
+      ProtocolFeatureStore featureStore) {
+    ChannelPipeline pipeline = ctx.pipeline();
+    pipeline.replace("messageDecoder", "messageDecoder", new Message.ProtocolDecoder(featureStore));
+    pipeline.replace(
+        "messageEncoder", "messageEncoder", messageEncoderFor(protocolVersion, featureStore));
+  }
+
   private void switchToModernFraming(ChannelHandlerContext ctx) {
     ChannelPipeline pipeline = ctx.pipeline();
     SegmentCodec segmentCodec =
@@ -75,4 +89,12 @@ private void switchToModernFraming(ChannelHandlerContext ctx) {
     pipeline.addAfter(
         "bytesToSegmentDecoder", "segmentToFrameDecoder", new SegmentToFrameDecoder());
   }
+
+  private static Message.ProtocolEncoder messageEncoderFor(
+      ProtocolVersion version, ProtocolFeatureStore protocolFeatureStore) {
+    if (version.toInt() > ProtocolVersion.V6.toInt()) {
+      throw new DriverInternalError("Unsupported protocol version " + version);
+    }
+    return new Message.ProtocolEncoder(version, protocolFeatureStore);
+  }
 }