From 540e14c7f1a82dffcb8b36432a646da1649e8596 Mon Sep 17 00:00:00 2001
From: Shubham Tagra <stagra@qubole.com>
Date: Thu, 30 Jan 2020 14:39:07 +0530
Subject: [PATCH] Rubix integration

---
 pom.xml                                       |   7 +
 presto-hive/pom.xml                           |   5 +
 .../hive/InternalHiveConnectorFactory.java    |  10 ++
 .../plugin/hive/rubix/RubixConfig.java        |  86 ++++++++++++
 .../rubix/RubixConfigurationInitializer.java  | 130 ++++++++++++++++++
 .../plugin/hive/rubix/RubixInitializer.java   | 110 +++++++++++++++
 .../plugin/hive/rubix/RubixModule.java        |  40 ++++++
 .../plugin/hive/util/HiveWriteUtils.java      |   3 +-
 .../plugin/hive/rubix/TestRubixConfig.java    |  59 ++++++++
 9 files changed, 449 insertions(+), 1 deletion(-)
 create mode 100644 presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixConfig.java
 create mode 100644 presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixConfigurationInitializer.java
 create mode 100644 presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixInitializer.java
 create mode 100644 presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixModule.java
 create mode 100644 presto-hive/src/test/java/io/prestosql/plugin/hive/rubix/TestRubixConfig.java

diff --git a/pom.xml b/pom.xml
index bfc1380248d20..4927c7b8a3e0c 100644
--- a/pom.xml
+++ b/pom.xml
@@ -53,6 +53,7 @@
         <dep.tempto.version>173</dep.tempto.version>
         <dep.gcs.version>2.0.0</dep.gcs.version>
         <dep.errorprone.version>2.3.3</dep.errorprone.version>
+        <rubix.version>0.3.6</rubix.version>
 
         <!--
           America/Bahia_Banderas has:
@@ -1215,6 +1216,12 @@
                 <type>pom</type>
                 <scope>import</scope>
             </dependency>
+
+            <dependency>
+                <groupId>com.qubole.rubix</groupId>
+                <artifactId>rubix-presto-shaded</artifactId>
+                <version>${rubix.version}</version>
+            </dependency>
         </dependencies>
     </dependencyManagement>
 
diff --git a/presto-hive/pom.xml b/presto-hive/pom.xml
index d77b6c8e1cbbd..0bb54a2e850d6 100644
--- a/presto-hive/pom.xml
+++ b/presto-hive/pom.xml
@@ -281,6 +281,11 @@
             <scope>provided</scope>
         </dependency>
 
+        <dependency>
+            <groupId>com.qubole.rubix</groupId>
+            <artifactId>rubix-presto-shaded</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>org.testng</groupId>
             <artifactId>testng</artifactId>
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/InternalHiveConnectorFactory.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/InternalHiveConnectorFactory.java
index 53e05e9aa8f5a..c02e84b8b8638 100644
--- a/presto-hive/src/main/java/io/prestosql/plugin/hive/InternalHiveConnectorFactory.java
+++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/InternalHiveConnectorFactory.java
@@ -34,6 +34,9 @@
 import io.prestosql.plugin.hive.metastore.HiveMetastore;
 import io.prestosql.plugin.hive.metastore.HiveMetastoreModule;
 import io.prestosql.plugin.hive.procedure.HiveProcedureModule;
+import io.prestosql.plugin.hive.rubix.RubixConfig;
+import io.prestosql.plugin.hive.rubix.RubixInitializer;
+import io.prestosql.plugin.hive.rubix.RubixModule;
 import io.prestosql.plugin.hive.s3.HiveS3Module;
 import io.prestosql.plugin.hive.security.HiveSecurityModule;
 import io.prestosql.plugin.hive.security.SystemTableAwareAccessControl;
@@ -84,6 +87,7 @@ public static Connector createConnector(String catalogName, Map<String, String>
                     new HiveS3Module(),
                     new HiveGcsModule(),
                     new HiveAzureModule(),
+                    new RubixModule(),
                     new HiveMetastoreModule(metastore),
                     new HiveSecurityModule(catalogName),
                     new HiveAuthenticationModule(),
@@ -106,6 +110,12 @@ public static Connector createConnector(String catalogName, Map<String, String>
                     .setRequiredConfigurationProperties(config)
                     .initialize();
 
+            if (injector.getInstance(RubixConfig.class).isCacheEnabled()) {
+                // RubixInitializer needs ConfigurationInitializers, hence kept outside RubixModule
+                RubixInitializer rubixInitializer = injector.getInstance(RubixInitializer.class);
+                rubixInitializer.initializeRubix(context.getNodeManager());
+            }
+
             LifeCycleManager lifeCycleManager = injector.getInstance(LifeCycleManager.class);
             HiveMetadataFactory metadataFactory = injector.getInstance(HiveMetadataFactory.class);
             HiveTransactionManager transactionManager = injector.getInstance(HiveTransactionManager.class);
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixConfig.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixConfig.java
new file mode 100644
index 0000000000000..85699752bc065
--- /dev/null
+++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixConfig.java
@@ -0,0 +1,86 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.plugin.hive.rubix;
+
+import com.qubole.rubix.spi.CacheConfig;
+import io.airlift.configuration.Config;
+
+public class RubixConfig
+{
+    private boolean cacheEnabled;
+    private boolean parallelWarmupEnabled = true;
+    private String cacheLocation = "/tmp";
+    private int bookKeeperServerPort = CacheConfig.DEFAULT_BOOKKEEPER_SERVER_PORT;
+    private int dataTransferServerPort = CacheConfig.DEFAULT_DATA_TRANSFER_SERVER_PORT;
+
+    public boolean isCacheEnabled()
+    {
+        return cacheEnabled;
+    }
+
+    @Config("hive.cache.enabled")
+    public RubixConfig setCacheEnabled(boolean value)
+    {
+        this.cacheEnabled = value;
+        return this;
+    }
+
+    public boolean isParallelWarmupEnabled()
+    {
+        return parallelWarmupEnabled;
+    }
+
+    @Config("hive.cache.parallel-warmup-enabled")
+    public RubixConfig setParallelWarmupEnabled(boolean value)
+    {
+        this.parallelWarmupEnabled = value;
+        return this;
+    }
+
+    public String getCacheLocation()
+    {
+        return cacheLocation;
+    }
+
+    @Config("hive.cache.location")
+    public RubixConfig setCacheLocation(String location)
+    {
+        this.cacheLocation = location;
+        return this;
+    }
+
+    public int getBookKeeperServerPort()
+    {
+        return bookKeeperServerPort;
+    }
+
+    @Config("hive.cache.rubix-bookkeeper-port")
+    public RubixConfig setBookKeeperServerPort(int port)
+    {
+        this.bookKeeperServerPort = port;
+        return this;
+    }
+
+    public int getDataTransferServerPort()
+    {
+        return dataTransferServerPort;
+    }
+
+    @Config("hive.cache.rubix-data-transfer-port")
+    public RubixConfig setDataTransferServerPort(int port)
+    {
+        this.dataTransferServerPort = port;
+        return this;
+    }
+}
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixConfigurationInitializer.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixConfigurationInitializer.java
new file mode 100644
index 0000000000000..1b8e77178c1cb
--- /dev/null
+++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixConfigurationInitializer.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.plugin.hive.rubix;
+
+import com.qubole.rubix.prestosql.CachingPrestoGoogleHadoopFileSystem;
+import com.qubole.rubix.prestosql.CachingPrestoNativeAzureFileSystem;
+import com.qubole.rubix.prestosql.CachingPrestoS3FileSystem;
+import com.qubole.rubix.prestosql.PrestoClusterManager;
+import io.prestosql.plugin.hive.ConfigurationInitializer;
+import io.prestosql.spi.HostAddress;
+import org.apache.hadoop.conf.Configuration;
+
+import javax.inject.Inject;
+
+import static com.google.common.base.Preconditions.checkState;
+import static com.qubole.rubix.spi.CacheConfig.enableHeartbeat;
+import static com.qubole.rubix.spi.CacheConfig.setBookKeeperServerPort;
+import static com.qubole.rubix.spi.CacheConfig.setCacheDataDirPrefix;
+import static com.qubole.rubix.spi.CacheConfig.setCacheDataEnabled;
+import static com.qubole.rubix.spi.CacheConfig.setClusterNodeRefreshTime;
+import static com.qubole.rubix.spi.CacheConfig.setClusterNodesFetchRetryCount;
+import static com.qubole.rubix.spi.CacheConfig.setCoordinatorHostName;
+import static com.qubole.rubix.spi.CacheConfig.setCurrentNodeHostName;
+import static com.qubole.rubix.spi.CacheConfig.setDataTransferServerPort;
+import static com.qubole.rubix.spi.CacheConfig.setEmbeddedMode;
+import static com.qubole.rubix.spi.CacheConfig.setIsParallelWarmupEnabled;
+import static com.qubole.rubix.spi.CacheConfig.setOnMaster;
+import static com.qubole.rubix.spi.CacheConfig.setRubixClusterType;
+import static com.qubole.rubix.spi.CacheConfig.setWorkerNodeInfoExpiryPeriod;
+import static com.qubole.rubix.spi.ClusterType.PRESTOSQL_CLUSTER_MANAGER;
+
+public class RubixConfigurationInitializer
+        implements ConfigurationInitializer
+{
+    private static final String RUBIX_S3_FS_CLASS_NAME = CachingPrestoS3FileSystem.class.getName();
+    private static final String RUBIX_AZURE_FS_CLASS_NAME = CachingPrestoNativeAzureFileSystem.class.getName();
+    private static final String RUBIX_GS_FS_CLASS_NAME = CachingPrestoGoogleHadoopFileSystem.class.getName();
+
+    private final boolean parallelWarmupEnabled;
+    private final String cacheLocation;
+    private final int bookKeeperServerPort;
+    private final int dataTransferServerPort;
+
+    // Configs below are dependent on node joining the cluster
+    private boolean cacheNotReady = true;
+    private boolean isMaster;
+    private HostAddress masterAddress;
+    private String nodeAddress;
+
+    @Inject
+    public RubixConfigurationInitializer(RubixConfig config)
+    {
+        this.parallelWarmupEnabled = config.isParallelWarmupEnabled();
+        this.cacheLocation = config.getCacheLocation();
+        this.bookKeeperServerPort = config.getBookKeeperServerPort();
+        this.dataTransferServerPort = config.getDataTransferServerPort();
+    }
+
+    @Override
+    public void initializeConfiguration(Configuration config)
+    {
+        if (cacheNotReady) {
+            setCacheDataEnabled(config, false);
+            return;
+        }
+
+        updateConfiguration(config);
+    }
+
+    public Configuration updateConfiguration(Configuration config)
+    {
+        checkState(masterAddress != null, "masterAddress is not set");
+        setCacheDataEnabled(config, true);
+        setOnMaster(config, isMaster);
+        setCoordinatorHostName(config, masterAddress.getHostText());
+        PrestoClusterManager.setPrestoServerPort(config, masterAddress.getPort());
+        setCurrentNodeHostName(config, nodeAddress);
+
+        setIsParallelWarmupEnabled(config, parallelWarmupEnabled);
+        setCacheDataDirPrefix(config, cacheLocation);
+        setBookKeeperServerPort(config, bookKeeperServerPort);
+        setDataTransferServerPort(config, dataTransferServerPort);
+
+        setEmbeddedMode(config, true);
+        setRubixClusterType(config, PRESTOSQL_CLUSTER_MANAGER);
+        enableHeartbeat(config, false);
+        setClusterNodeRefreshTime(config, 10);
+        setClusterNodesFetchRetryCount(config, Integer.MAX_VALUE);
+        setWorkerNodeInfoExpiryPeriod(config, 1);
+
+        config.set("fs.s3.impl", RUBIX_S3_FS_CLASS_NAME);
+        config.set("fs.s3a.impl", RUBIX_S3_FS_CLASS_NAME);
+        config.set("fs.s3n.impl", RUBIX_S3_FS_CLASS_NAME);
+        config.set("fs.wasb.impl", RUBIX_AZURE_FS_CLASS_NAME);
+        config.set("fs.gs.impl", RUBIX_GS_FS_CLASS_NAME);
+        return config;
+    }
+
+    public void setMaster(boolean master)
+    {
+        isMaster = master;
+    }
+
+    public void setMasterAddress(HostAddress masterAddress)
+    {
+        this.masterAddress = masterAddress;
+    }
+
+    public void setCurrentNodeAddress(String nodeAddress)
+    {
+        this.nodeAddress = nodeAddress;
+    }
+
+    public void initializationDone()
+    {
+        checkState(masterAddress != null, "masterAddress is not set");
+        cacheNotReady = false;
+    }
+}
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixInitializer.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixInitializer.java
new file mode 100644
index 0000000000000..f55479f317739
--- /dev/null
+++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixInitializer.java
@@ -0,0 +1,110 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.plugin.hive.rubix;
+
+import com.codahale.metrics.MetricRegistry;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.qubole.rubix.bookkeeper.BookKeeper;
+import com.qubole.rubix.bookkeeper.BookKeeperServer;
+import com.qubole.rubix.bookkeeper.LocalDataTransferServer;
+import com.qubole.rubix.core.CachingFileSystem;
+import io.airlift.log.Logger;
+import io.prestosql.plugin.hive.ConfigurationInitializer;
+import io.prestosql.plugin.hive.HiveCatalogName;
+import io.prestosql.spi.Node;
+import io.prestosql.spi.NodeManager;
+import org.apache.hadoop.conf.Configuration;
+
+import javax.inject.Inject;
+
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+/*
+ * Responsibilities of this initializer:
+ * 1. Lazily setup RubixConfigurationInitializer with information about master when it is available
+ * 2. Start Rubix Servers.
+ * 3. Inject BookKeeper object into CachingFileSystem class
+ */
+public class RubixInitializer
+{
+    private static final Logger log = Logger.get(RubixInitializer.class);
+
+    private final HiveCatalogName hiveCatalogName;
+    private final RubixConfigurationInitializer rubixConfigurationInitializer;
+    private final Set<ConfigurationInitializer> configurationInitializers;
+
+    @Inject
+    public RubixInitializer(HiveCatalogName hiveCatalogName, RubixConfigurationInitializer rubixConfigurationInitializer, Set<ConfigurationInitializer> configurationInitializers)
+    {
+        this.hiveCatalogName = hiveCatalogName;
+        this.rubixConfigurationInitializer = rubixConfigurationInitializer;
+        this.configurationInitializers = configurationInitializers;
+    }
+
+    public void initializeRubix(NodeManager nodeManager)
+    {
+        ExecutorService initializerService = Executors.newSingleThreadExecutor();
+        ListenableFuture<Boolean> nodeJoinFuture = MoreExecutors.listeningDecorator(initializerService).submit(() ->
+        {
+            while (!(nodeManager.getAllNodes().contains(nodeManager.getCurrentNode()) &&
+                nodeManager.getAllNodes().stream().anyMatch(Node::isCoordinator))) {
+                try {
+                    Thread.sleep(100);
+                }
+                catch (InterruptedException e) {
+                    return false;
+                }
+            }
+            return true;
+        });
+
+        Futures.transform(nodeJoinFuture,
+                nodeJoined ->
+                {
+                    if (nodeJoined) {
+                        Node master = nodeManager.getAllNodes().stream().filter(node -> node.isCoordinator()).findFirst().get();
+                        boolean isMaster = nodeManager.getCurrentNode().isCoordinator();
+
+                        rubixConfigurationInitializer.setMaster(isMaster);
+                        rubixConfigurationInitializer.setMasterAddress(master.getHostAndPort());
+                        rubixConfigurationInitializer.setCurrentNodeAddress(nodeManager.getCurrentNode().getHost());
+
+                        Configuration configuration = new Configuration(false);
+                        for (ConfigurationInitializer configurationInitializer : configurationInitializers) {
+                            configurationInitializer.initializeConfiguration(configuration);
+                        }
+
+                        // RubixConfigurationInitializer.initializeConfiguration will not update configurations yet as it has not been fully initialized
+                        // Apply configurations from it by skipping init check
+                        rubixConfigurationInitializer.updateConfiguration(configuration);
+                        MetricRegistry metricRegistry = new MetricRegistry();
+                        BookKeeperServer bookKeeperServer = new BookKeeperServer();
+                        BookKeeper bookKeeper = bookKeeperServer.startServer(configuration, metricRegistry);
+                        LocalDataTransferServer.startServer(configuration, metricRegistry, bookKeeper);
+
+                        CachingFileSystem.setLocalBookKeeper(bookKeeper, "catalog=" + hiveCatalogName);
+                        log.info("Rubix initialized successfully");
+                        rubixConfigurationInitializer.initializationDone();
+                    }
+
+                    // In case of node join failing, let the Rubix cache be in default disabled state
+                    return null;
+                },
+                initializerService);
+    }
+}
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixModule.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixModule.java
new file mode 100644
index 0000000000000..e4fb330a7cdc6
--- /dev/null
+++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/rubix/RubixModule.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.plugin.hive.rubix;
+
+import com.google.inject.Binder;
+import com.google.inject.Scopes;
+import io.airlift.configuration.AbstractConfigurationAwareModule;
+import io.prestosql.plugin.hive.ConfigurationInitializer;
+
+import static com.google.inject.multibindings.Multibinder.newSetBinder;
+import static io.airlift.configuration.ConfigBinder.configBinder;
+
+public class RubixModule
+        extends AbstractConfigurationAwareModule
+{
+    @Override
+    protected void setup(Binder binder)
+    {
+        RubixConfig config = buildConfigObject(RubixConfig.class);
+        if (!config.isCacheEnabled()) {
+            return;
+        }
+
+        configBinder(binder).bindConfig(RubixConfig.class);
+        binder.bind(RubixConfigurationInitializer.class).in(Scopes.SINGLETON);
+        newSetBinder(binder, ConfigurationInitializer.class).addBinding().to(RubixConfigurationInitializer.class).in(Scopes.SINGLETON);
+        binder.bind(RubixInitializer.class).in(Scopes.SINGLETON);
+    }
+}
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/util/HiveWriteUtils.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/util/HiveWriteUtils.java
index a2e70439a86f3..e95efb63c0c76 100644
--- a/presto-hive/src/main/java/io/prestosql/plugin/hive/util/HiveWriteUtils.java
+++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/util/HiveWriteUtils.java
@@ -17,6 +17,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Shorts;
 import com.google.common.primitives.SignedBytes;
+import com.qubole.rubix.prestosql.CachingPrestoS3FileSystem;
 import io.prestosql.plugin.hive.HdfsEnvironment;
 import io.prestosql.plugin.hive.HdfsEnvironment.HdfsContext;
 import io.prestosql.plugin.hive.HiveReadOnlyException;
@@ -489,7 +490,7 @@ public static boolean isS3FileSystem(HdfsContext context, HdfsEnvironment hdfsEn
     {
         try {
             FileSystem fileSystem = getRawFileSystem(hdfsEnvironment.getFileSystem(context, path));
-            return fileSystem instanceof PrestoS3FileSystem || fileSystem.getClass().getName().equals(EMR_FS_CLASS_NAME);
+            return fileSystem instanceof PrestoS3FileSystem || fileSystem.getClass().getName().equals(EMR_FS_CLASS_NAME) || fileSystem instanceof CachingPrestoS3FileSystem;
         }
         catch (IOException e) {
             throw new PrestoException(HIVE_FILESYSTEM_ERROR, "Failed checking path: " + path, e);
diff --git a/presto-hive/src/test/java/io/prestosql/plugin/hive/rubix/TestRubixConfig.java b/presto-hive/src/test/java/io/prestosql/plugin/hive/rubix/TestRubixConfig.java
new file mode 100644
index 0000000000000..eec4ee67deec7
--- /dev/null
+++ b/presto-hive/src/test/java/io/prestosql/plugin/hive/rubix/TestRubixConfig.java
@@ -0,0 +1,59 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.plugin.hive.rubix;
+
+import com.google.common.collect.ImmutableMap;
+import com.qubole.rubix.spi.CacheConfig;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+
+import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+
+public class TestRubixConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(RubixConfig.class)
+                .setBookKeeperServerPort(CacheConfig.DEFAULT_BOOKKEEPER_SERVER_PORT)
+                .setDataTransferServerPort(CacheConfig.DEFAULT_DATA_TRANSFER_SERVER_PORT)
+                .setCacheEnabled(false)
+                .setCacheLocation("/tmp")
+                .setParallelWarmupEnabled(true));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+    {
+        Map<String, String> properties = new ImmutableMap.Builder<String, String>()
+                .put("hive.cache.enabled", "true")
+                .put("hive.cache.parallel-warmup-enabled", "false")
+                .put("hive.cache.location", "/etc")
+                .put("hive.cache.rubix-bookkeeper-port", "1234")
+                .put("hive.cache.rubix-data-transfer-port", "1235")
+                .build();
+
+        RubixConfig expected = new RubixConfig()
+                .setCacheEnabled(true)
+                .setParallelWarmupEnabled(false)
+                .setCacheLocation("/etc")
+                .setBookKeeperServerPort(1234)
+                .setDataTransferServerPort(1235);
+
+        assertFullMapping(properties, expected);
+    }
+}