Implement ProtocolFeatureStore

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

@@ -32,6 +32,7 @@
 import com.datastax.oss.driver.internal.core.adminrequest.ThrottledAdminRequestHandler;
 import com.datastax.oss.driver.internal.core.pool.ChannelPool;
 import com.datastax.oss.driver.internal.core.protocol.LwtInfo;
+import com.datastax.oss.driver.internal.core.protocol.ProtocolFeatureStore;
 import com.datastax.oss.driver.internal.core.protocol.ShardingInfo;
 import com.datastax.oss.driver.internal.core.protocol.ShardingInfo.ConnectionShardingInfo;
 import com.datastax.oss.driver.internal.core.session.DefaultSession;
@@ -61,9 +62,6 @@ public class DriverChannel {
   static final AttributeKey<String> CLUSTER_NAME_KEY = AttributeKey.valueOf("cluster_name");
   static final AttributeKey<Map<String, List<String>>> OPTIONS_KEY =
       AttributeKey.valueOf("options");
-  static final AttributeKey<ConnectionShardingInfo> SHARDING_INFO_KEY =
-      AttributeKey.valueOf("sharding_info");
-  static final AttributeKey<LwtInfo> LWT_INFO_KEY = AttributeKey.valueOf("lwt_info");
 
   @SuppressWarnings("RedundantStringConstructorCall")
   static final Object GRACEFUL_CLOSE_MESSAGE = new String("GRACEFUL_CLOSE_MESSAGE");
@@ -149,17 +147,17 @@ public Map<String, List<String>> getOptions() {
   }
 
   public int getShardId() {
-    return channel.hasAttr(SHARDING_INFO_KEY) ? channel.attr(SHARDING_INFO_KEY).get().shardId : 0;
+    ConnectionShardingInfo info = ProtocolFeatureStore.loadFromChannel(channel).getShardingInfo();
+    return info != null ? info.shardId : 0;
   }
 
   public ShardingInfo getShardingInfo() {
-    return channel.hasAttr(SHARDING_INFO_KEY)
-        ? channel.attr(SHARDING_INFO_KEY).get().shardingInfo
-        : null;
+    ConnectionShardingInfo info = ProtocolFeatureStore.loadFromChannel(channel).getShardingInfo();
+    return info != null ? info.shardingInfo : null;
   }
 
   public LwtInfo getLwtInfo() {
-    return channel.attr(LWT_INFO_KEY).get();
+    return ProtocolFeatureStore.loadFromChannel(channel).getLwtFeatureInfo();
   }
 
   /**

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

@@ -23,8 +23,6 @@
  */
 package com.datastax.oss.driver.internal.core.channel;
 
-import static com.datastax.oss.driver.internal.core.channel.DriverChannel.LWT_INFO_KEY;
-
 import com.datastax.oss.driver.api.core.DefaultProtocolVersion;
 import com.datastax.oss.driver.api.core.InvalidKeyspaceException;
 import com.datastax.oss.driver.api.core.ProtocolVersion;
@@ -40,12 +38,9 @@
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.protocol.BytesToSegmentDecoder;
 import com.datastax.oss.driver.internal.core.protocol.FrameToSegmentEncoder;
-import com.datastax.oss.driver.internal.core.protocol.LwtInfo;
+import com.datastax.oss.driver.internal.core.protocol.ProtocolFeatureStore;
 import com.datastax.oss.driver.internal.core.protocol.SegmentToBytesEncoder;
 import com.datastax.oss.driver.internal.core.protocol.SegmentToFrameDecoder;
-import com.datastax.oss.driver.internal.core.protocol.ShardingInfo;
-import com.datastax.oss.driver.internal.core.protocol.ShardingInfo.ConnectionShardingInfo;
-import com.datastax.oss.driver.internal.core.protocol.TabletInfo;
 import com.datastax.oss.driver.internal.core.util.ProtocolUtils;
 import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions;
 import com.datastax.oss.protocol.internal.Message;
@@ -96,8 +91,7 @@ class ProtocolInitHandler extends ConnectInitHandler {
   private String logPrefix;
   private ChannelHandlerContext ctx;
   private final boolean querySupportedOptions;
-  private LwtInfo lwtInfo;
-  private TabletInfo tabletInfo;
+  private ProtocolFeatureStore featureStore;
 
   /**
    * @param querySupportedOptions whether to send OPTIONS as the first message, to request which
@@ -192,11 +186,8 @@ Message getRequest() {
           return request = Options.INSTANCE;
         case STARTUP:
           Map<String, String> startupOptions = new HashMap<>(context.getStartupOptions());
-          if (lwtInfo != null) {
-            lwtInfo.addOption(startupOptions);
-          }
-          if (tabletInfo != null && tabletInfo.isEnabled()) {
-            TabletInfo.addOption(startupOptions);
+          if (featureStore != null) {
+            featureStore.populateStartupOptions(startupOptions);
           }
           return request = new Startup(startupOptions);
         case GET_CLUSTER_NAME:
@@ -229,15 +220,8 @@ void onResponse(Message response) {
         if (step == Step.OPTIONS && response instanceof Supported) {
           channel.attr(DriverChannel.OPTIONS_KEY).set(((Supported) response).options);
           Supported res = (Supported) response;
-          ConnectionShardingInfo shardingInfo = ShardingInfo.parseShardingInfo(res.options);
-          if (shardingInfo != null) {
-            channel.attr(DriverChannel.SHARDING_INFO_KEY).set(shardingInfo);
-          }
-          lwtInfo = LwtInfo.parseLwtInfo(res.options);
-          if (lwtInfo != null) {
-            channel.attr(LWT_INFO_KEY).set(lwtInfo);
-          }
-          tabletInfo = TabletInfo.parseTabletInfo(res.options);
+          featureStore = ProtocolFeatureStore.parseSupportedOptions(res.options);
+          featureStore.storeInChannel(channel);
           step = Step.STARTUP;
           send();
         } else if (step == Step.STARTUP && response instanceof Ready) {

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

@@ -37,7 +37,7 @@ public boolean isLwt(int flags) {
     return (flags & mask) == mask;
   }
 
-  public static LwtInfo parseLwtInfo(Map<String, List<String>> supported) {
+  public static LwtInfo loadFromSupportedOptions(Map<String, List<String>> supported) {
     if (!supported.containsKey(SCYLLA_LWT_ADD_METADATA_MARK_KEY)) {
       return null;
     }
@@ -67,7 +67,7 @@ public static LwtInfo parseLwtInfo(Map<String, List<String>> supported) {
     return new LwtInfo((int) mask);
   }
 
-  public void addOption(Map<String, String> options) {
+  public void populateStartupOptions(Map<String, String> options) {
     options.put(SCYLLA_LWT_ADD_METADATA_MARK_KEY, Integer.toString(mask));
   }
 }

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

@@ -0,0 +1,60 @@
+package com.datastax.oss.driver.internal.core.protocol;
+
+import edu.umd.cs.findbugs.annotations.NonNull;
+import io.netty.channel.Channel;
+import io.netty.util.AttributeKey;
+import java.util.List;
+import java.util.Map;
+
+public class ProtocolFeatureStore {
+  private static final AttributeKey<ProtocolFeatureStore> CHANNEL_KEY =
+      AttributeKey.valueOf("protocol_feature_store");
+
+  private final LwtInfo lwtInfo;
+  private final ShardingInfo.ConnectionShardingInfo shardingInfo;
+  private final TabletInfo tabletInfo;
+
+  ProtocolFeatureStore(
+      LwtInfo lwtInfo, ShardingInfo.ConnectionShardingInfo shardingInfo, TabletInfo tabletInfo) {
+    this.lwtInfo = lwtInfo;
+    this.shardingInfo = shardingInfo;
+    this.tabletInfo = tabletInfo;
+  }
+
+  public LwtInfo getLwtFeatureInfo() {
+    return lwtInfo;
+  }
+
+  public ShardingInfo.ConnectionShardingInfo getShardingInfo() {
+    return shardingInfo;
+  }
+
+  public TabletInfo getTabletFeatureInfo() {
+    return tabletInfo;
+  }
+
+  public static ProtocolFeatureStore parseSupportedOptions(
+      @NonNull Map<String, List<String>> options) {
+    LwtInfo lwtInfo = LwtInfo.loadFromSupportedOptions(options);
+    ShardingInfo.ConnectionShardingInfo shardingInfo = ShardingInfo.parseShardingInfo(options);
+    TabletInfo tabletInfo = TabletInfo.loadFromSupportedOptions(options);
+    return new ProtocolFeatureStore(lwtInfo, shardingInfo, tabletInfo);
+  }
+
+  public void populateStartupOptions(@NonNull Map<String, String> options) {
+    if (lwtInfo != null) {
+      lwtInfo.populateStartupOptions(options);
+    }
+    if (tabletInfo != null && tabletInfo.isEnabled()) {
+      TabletInfo.populateStartupOptions(options);
+    }
+  }
+
+  public static ProtocolFeatureStore loadFromChannel(@NonNull Channel channel) {
+    return channel.attr(ProtocolFeatureStore.CHANNEL_KEY).get();
+  }
+
+  public void storeInChannel(@NonNull Channel channel) {
+    channel.attr(ProtocolFeatureStore.CHANNEL_KEY).set(this);
+  }
+}

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

@@ -19,15 +19,15 @@ public boolean isEnabled() {
     return enabled;
   }
 
-  public static TabletInfo parseTabletInfo(Map<String, List<String>> supported) {
+  public static TabletInfo loadFromSupportedOptions(Map<String, List<String>> supported) {
     List<String> values = supported.get(SCYLLA_TABLETS_STARTUP_OPTION_KEY);
     return new TabletInfo(
         values != null
             && values.size() == 1
             && values.get(0).equals(SCYLLA_TABLETS_STARTUP_OPTION_VALUE));
   }
 
-  public static void addOption(Map<String, String> options) {
+  public static void populateStartupOptions(Map<String, String> options) {
     options.put(SCYLLA_TABLETS_STARTUP_OPTION_KEY, SCYLLA_TABLETS_STARTUP_OPTION_VALUE);
   }
 }