Implement ProtocolFeatureStore

Implement class to store information about all features negotiated.

Author: dkropachev

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;
@@ -63,7 +64,6 @@ public class DriverChannel {
       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");
@@ -159,7 +159,7 @@ public ShardingInfo getShardingInfo() {
   }
 
   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;
@@ -41,6 +39,7 @@
 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;
@@ -96,7 +95,7 @@ class ProtocolInitHandler extends ConnectInitHandler {
   private String logPrefix;
   private ChannelHandlerContext ctx;
   private final boolean querySupportedOptions;
-  private LwtInfo lwtInfo;
+  private ProtocolFeatureStore featureStore;
   private TabletInfo tabletInfo;
 
   /**
@@ -192,9 +191,7 @@ Message getRequest() {
           return request = Options.INSTANCE;
         case STARTUP:
           Map<String, String> startupOptions = new HashMap<>(context.getStartupOptions());
-          if (lwtInfo != null) {
-            lwtInfo.addOption(startupOptions);
-          }
+          featureStore.populateStartupOptions(startupOptions);
           if (tabletInfo != null && tabletInfo.isEnabled()) {
             TabletInfo.addOption(startupOptions);
           }
@@ -229,15 +226,13 @@ void onResponse(Message response) {
         if (step == Step.OPTIONS && response instanceof Supported) {
           channel.attr(DriverChannel.OPTIONS_KEY).set(((Supported) response).options);
           Supported res = (Supported) response;
+          featureStore = ProtocolFeatureStore.parseSupportedOptions(res.options);
           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.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,41 @@
+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;
+
+  ProtocolFeatureStore(LwtInfo lwtInfo) {
+    this.lwtInfo = lwtInfo;
+  }
+
+  public LwtInfo getLwtFeatureInfo() {
+    return lwtInfo;
+  }
+
+  public static ProtocolFeatureStore parseSupportedOptions(Map<String, List<String>> options) {
+    LwtInfo lwtInfo = LwtInfo.loadFromSupportedOptions(options);
+    return new ProtocolFeatureStore(lwtInfo);
+  }
+
+  public void populateStartupOptions(Map<String, String> options) {
+    if (lwtInfo != null) {
+      lwtInfo.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);
+  }
+}