Pull ShardingInfo into ProtocolFeatureStore

Author: dkropachev

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

@@ -62,8 +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");
 
   @SuppressWarnings("RedundantStringConstructorCall")
   static final Object GRACEFUL_CLOSE_MESSAGE = new String("GRACEFUL_CLOSE_MESSAGE");
@@ -149,13 +147,13 @@ 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() {

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

@@ -42,8 +42,6 @@
 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;
@@ -227,10 +225,6 @@ void onResponse(Message response) {
           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);
-          }
           tabletInfo = TabletInfo.parseTabletInfo(res.options);
           featureStore.storeInChannel(channel);
           step = Step.STARTUP;

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

@@ -11,18 +11,25 @@ public class ProtocolFeatureStore {
       AttributeKey.valueOf("protocol_feature_store");
 
   private final LwtInfo lwtInfo;
+  private final ShardingInfo.ConnectionShardingInfo shardingInfo;
 
-  ProtocolFeatureStore(LwtInfo lwtInfo) {
+  ProtocolFeatureStore(LwtInfo lwtInfo, ShardingInfo.ConnectionShardingInfo shardingInfo) {
     this.lwtInfo = lwtInfo;
+    this.shardingInfo = shardingInfo;
   }
 
   public LwtInfo getLwtFeatureInfo() {
     return lwtInfo;
   }
 
+  public ShardingInfo.ConnectionShardingInfo getShardingInfo() {
+    return shardingInfo;
+  }
+
   public static ProtocolFeatureStore parseSupportedOptions(Map<String, List<String>> options) {
     LwtInfo lwtInfo = LwtInfo.loadFromSupportedOptions(options);
-    return new ProtocolFeatureStore(lwtInfo);
+    ShardingInfo.ConnectionShardingInfo shardingInfo = ShardingInfo.parseShardingInfo(options);
+    return new ProtocolFeatureStore(lwtInfo, shardingInfo);
   }
 
   public void populateStartupOptions(Map<String, String> options) {