[#596] Refactoring functions that update pipeline. Moving building of protocol features to ProtocolFeatureStore.

Author: nikagra

core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java

@@ -154,7 +154,7 @@ public ProtocolFeatureStore getSupportedFeatures() {
 
     ProtocolFeatureStore fromChannel = ProtocolFeatureStore.loadFromChannel(channel);
     if (fromChannel == null) {
-      return ProtocolFeatureStore.Empty;
+      return ProtocolFeatureStore.EMPTY;
     }
     // Features can't be renegotiated.
     // Once features is populated into channel it is enough to update cache and no need to

core/src/main/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandler.java

@@ -45,6 +45,7 @@
 import com.datastax.oss.driver.internal.core.protocol.SegmentToFrameDecoder;
 import com.datastax.oss.driver.internal.core.util.ProtocolUtils;
 import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions;
+import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
 import com.datastax.oss.protocol.internal.Message;
 import com.datastax.oss.protocol.internal.ProtocolConstants;
 import com.datastax.oss.protocol.internal.ProtocolConstants.ErrorCode;
@@ -225,16 +226,15 @@ void onResponse(Message response) {
           channel.attr(DriverChannel.OPTIONS_KEY).set(supported.options);
           featureStore = ProtocolFeatureStore.parseSupportedOptions(supported.options);
           featureStore.storeInChannel(channel);
-          maybeUpdatePipelineWithProtocolOptions(featureStore.isMetadataIdEnabled());
           step = Step.STARTUP;
           send();
         } else if (step == Step.STARTUP && response instanceof Ready) {
-          maybeSwitchToModernFraming();
+          maybeUpdatePipeline();
           context.getAuthProvider().ifPresent(provider -> provider.onMissingChallenge(endPoint));
           step = Step.GET_CLUSTER_NAME;
           send();
         } else if (step == Step.STARTUP && response instanceof Authenticate) {
-          maybeSwitchToModernFraming();
+          maybeUpdatePipeline();
           Authenticate authenticate = (Authenticate) response;
           authenticator = buildAuthenticator(endPoint, authenticate.authenticator);
           authenticator
@@ -400,11 +400,18 @@ public String toString() {
   }
 
   /**
-   * Rearranges the pipeline to deal with the new framing structure in protocol v5 and above. The
+   * Conditionally rebuilds pipeline.
+   *
+   * <p>Rearranges the pipeline to deal with the new framing structure in protocol v5 and above. The
    * first messages still use the legacy format, we only do this after a successful response to the
    * first STARTUP message.
+   *
+   * <p>If <code>SCYLLA_USE_METADATA_ID</code> feature was negotiated we need to replace {@link
+   * FrameEncoder} and {@link FrameDecoder} handlers with instances aware of a negotiated protocol
+   * feature.
    */
-  private void maybeSwitchToModernFraming() {
+  private void maybeUpdatePipeline() {
+    ProtocolFeatures protocolFeatures = featureStore.getProtocolFeatures();
     if (context
         .getProtocolVersionRegistry()
         .supports(initialProtocolVersion, DefaultProtocolFeature.MODERN_FRAMING)) {
@@ -432,20 +439,7 @@ private void maybeSwitchToModernFraming() {
           ChannelFactory.BYTES_TO_SEGMENT_DECODER_NAME,
           ChannelFactory.SEGMENT_TO_FRAME_DECODER_NAME,
           new SegmentToFrameDecoder(context.getFrameCodec(), logPrefix));
-    }
-  }
-
-  /**
-   * If <code>SCYLLA_USE_METADATA_ID</code> feature was negotiated we need to replace {@link
-   * FrameEncoder} and {@link FrameDecoder} handlers with instances aware of a negotiated protocol
-   * feature.
-   *
-   * @param metadataIdEnabled indicates if feature is successfully negotiated
-   */
-  private void maybeUpdatePipelineWithProtocolOptions(boolean metadataIdEnabled) {
-    if (metadataIdEnabled) {
-      ProtocolFeatures protocolFeatures =
-          new ProtocolFeatures.Builder().setScyllaUseMetadataId().build();
+    } else if (protocolFeatures.isScyllaUseMetadataId()) {
       int maxFrameLength =
           (int)
               context
@@ -457,7 +451,10 @@ private void maybeUpdatePipelineWithProtocolOptions(boolean metadataIdEnabled) {
       pipeline.replace(
           ChannelFactory.FRAME_TO_BYTES_ENCODER_NAME,
           ChannelFactory.FRAME_TO_BYTES_ENCODER_NAME,
-          new FrameEncoder(context.getFrameCodec(), protocolFeatures, maxFrameLength));
+          new FrameEncoder(
+              context.getFrameCodec(),
+              protocolFeatures, // Passing updated protocol features to alter codecs behaviors
+              maxFrameLength));
       pipeline.replace(
           ChannelFactory.BYTES_TO_FRAME_DECODER_NAME,
           ChannelFactory.BYTES_TO_FRAME_DECODER_NAME,
@@ -468,4 +465,9 @@ private void maybeUpdatePipelineWithProtocolOptions(boolean metadataIdEnabled) {
   private String getString(List<ByteBuffer> row, int i) {
     return TypeCodecs.TEXT.decode(row.get(i), DefaultProtocolVersion.DEFAULT);
   }
+
+  @VisibleForTesting
+  void setFeatureStore(ProtocolFeatureStore featureStore) {
+    this.featureStore = featureStore;
+  }
 }

core/src/main/java/com/datastax/oss/driver/internal/core/protocol/ProtocolFeatureStore.java

@@ -1,5 +1,6 @@
 package com.datastax.oss.driver.internal.core.protocol;
 
+import com.datastax.oss.protocol.internal.ProtocolFeatures;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import io.netty.channel.Channel;
 import io.netty.util.AttributeKey;
@@ -14,8 +15,9 @@ public class ProtocolFeatureStore {
   private final ShardingInfo.ConnectionShardingInfo shardingInfo;
   private final TabletInfo tabletInfo;
   private final boolean metadataIdEnabled;
+  private ProtocolFeatures protocolFeatures;
 
-  public static final ProtocolFeatureStore Empty =
+  public static final ProtocolFeatureStore EMPTY =
       new ProtocolFeatureStore(null, null, null, false);
 
   ProtocolFeatureStore(
@@ -73,4 +75,19 @@ public static ProtocolFeatureStore loadFromChannel(@NonNull Channel channel) {
   public void storeInChannel(@NonNull Channel channel) {
     channel.attr(ProtocolFeatureStore.CHANNEL_KEY).set(this);
   }
+
+  public ProtocolFeatures getProtocolFeatures() {
+    if (protocolFeatures == null) {
+      protocolFeatures = buildProtocolFeatures();
+    }
+    return protocolFeatures;
+  }
+
+  private ProtocolFeatures buildProtocolFeatures() {
+    if (metadataIdEnabled) {
+      return new ProtocolFeatures.Builder().setScyllaUseMetadataId().build();
+    } else {
+      return ProtocolFeatures.EMPTY;
+    }
+  }
 }

core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryTestBase.java

@@ -34,6 +34,7 @@
 import com.datastax.oss.driver.internal.core.context.NettyOptions;
 import com.datastax.oss.driver.internal.core.metrics.NodeMetricUpdater;
 import com.datastax.oss.driver.internal.core.protocol.ByteBufPrimitiveCodec;
+import com.datastax.oss.driver.internal.core.protocol.ProtocolFeatureStore;
 import com.datastax.oss.protocol.internal.Compressor;
 import com.datastax.oss.protocol.internal.Frame;
 import com.datastax.oss.protocol.internal.FrameCodec;
@@ -271,6 +272,7 @@ protected void initChannel(Channel channel) throws Exception {
                     options,
                     heartbeatHandler,
                     productType == null);
+            initHandler.setFeatureStore(ProtocolFeatureStore.EMPTY);
             channel
                 .pipeline()
                 .addLast(ChannelFactory.INFLIGHT_HANDLER_NAME, inFlightHandler)