Pull TabletInfo into ProtocolFeatureStore

Author: dkropachev

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

@@ -38,11 +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.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;
@@ -94,7 +92,6 @@ class ProtocolInitHandler extends ConnectInitHandler {
   private ChannelHandlerContext ctx;
   private final boolean querySupportedOptions;
   private ProtocolFeatureStore featureStore;
-  private TabletInfo tabletInfo;
 
   /**
    * @param querySupportedOptions whether to send OPTIONS as the first message, to request which
@@ -190,9 +187,6 @@ Message getRequest() {
         case STARTUP:
           Map<String, String> startupOptions = new HashMap<>(context.getStartupOptions());
           featureStore.populateStartupOptions(startupOptions);
-          if (tabletInfo != null && tabletInfo.isEnabled()) {
-            TabletInfo.addOption(startupOptions);
-          }
           return request = new Startup(startupOptions);
         case GET_CLUSTER_NAME:
           return request = CLUSTER_NAME_QUERY;
@@ -225,7 +219,6 @@ void onResponse(Message response) {
           channel.attr(DriverChannel.OPTIONS_KEY).set(((Supported) response).options);
           Supported res = (Supported) response;
           featureStore = ProtocolFeatureStore.parseSupportedOptions(res.options);
-          tabletInfo = TabletInfo.parseTabletInfo(res.options);
           featureStore.storeInChannel(channel);
           step = Step.STARTUP;
           send();

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

@@ -12,10 +12,13 @@ public class ProtocolFeatureStore {
 
   private final LwtInfo lwtInfo;
   private final ShardingInfo.ConnectionShardingInfo shardingInfo;
+  private final TabletInfo tabletInfo;
 
-  ProtocolFeatureStore(LwtInfo lwtInfo, ShardingInfo.ConnectionShardingInfo shardingInfo) {
+  ProtocolFeatureStore(
+      LwtInfo lwtInfo, ShardingInfo.ConnectionShardingInfo shardingInfo, TabletInfo tabletInfo) {
     this.lwtInfo = lwtInfo;
     this.shardingInfo = shardingInfo;
+    this.tabletInfo = tabletInfo;
   }
 
   public LwtInfo getLwtFeatureInfo() {
@@ -26,16 +29,24 @@ public ShardingInfo.ConnectionShardingInfo getShardingInfo() {
     return shardingInfo;
   }
 
+  public TabletInfo getTabletFeatureInfo() {
+    return tabletInfo;
+  }
+
   public static ProtocolFeatureStore parseSupportedOptions(Map<String, List<String>> options) {
     LwtInfo lwtInfo = LwtInfo.loadFromSupportedOptions(options);
     ShardingInfo.ConnectionShardingInfo shardingInfo = ShardingInfo.parseShardingInfo(options);
-    return new ProtocolFeatureStore(lwtInfo, shardingInfo);
+    TabletInfo tabletInfo = TabletInfo.loadFromSupportedOptions(options);
+    return new ProtocolFeatureStore(lwtInfo, shardingInfo, tabletInfo);
   }
 
   public void populateStartupOptions(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) {

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);
   }
 }