[QoS] Qos ete test

atlasdb-dropwizard-bundle/src/main/java/com/palantir/atlasdb/dropwizard/AtlasDbConfigurationProvider.java

@@ -15,8 +15,13 @@
  */
 package com.palantir.atlasdb.dropwizard;
 
+import java.util.Optional;
+
 import com.palantir.atlasdb.config.AtlasDbConfig;
+import com.palantir.atlasdb.config.AtlasDbRuntimeConfig;
 
 public interface AtlasDbConfigurationProvider {
     AtlasDbConfig getAtlasDbConfig();
+
+    Optional<AtlasDbRuntimeConfig> getAtlasDbRuntimeConfig();
 }

atlasdb-dropwizard-bundle/src/test/java/com/palantir/atlasdb/dropwizard/commands/AtlasDbConsoleCommandTest.java

@@ -22,6 +22,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.Test;
@@ -30,6 +31,7 @@
 import com.google.common.base.Throwables;
 import com.palantir.atlasdb.config.AtlasDbConfig;
 import com.palantir.atlasdb.config.AtlasDbConfigs;
+import com.palantir.atlasdb.config.AtlasDbRuntimeConfig;
 import com.palantir.atlasdb.config.ImmutableAtlasDbConfig;
 import com.palantir.atlasdb.config.ImmutableLeaderConfig;
 import com.palantir.atlasdb.dropwizard.AtlasDbConfigurationProvider;
@@ -100,5 +102,10 @@ private class AtlasDbDropwizardConfig extends Configuration implements AtlasDbCo
         public AtlasDbConfig getAtlasDbConfig() {
             return atlasDbConfig;
         }
+
+        @Override
+        public Optional<AtlasDbRuntimeConfig> getAtlasDbRuntimeConfig() {
+            return Optional.empty();
+        }
     }
 }

atlasdb-ete-tests/src/main/java/com/palantir/atlasdb/AtlasDbEteConfiguration.java

@@ -15,21 +15,33 @@
  */
 package com.palantir.atlasdb;
 
+import java.util.Optional;
+
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.palantir.atlasdb.config.AtlasDbConfig;
+import com.palantir.atlasdb.config.AtlasDbRuntimeConfig;
 import com.palantir.atlasdb.dropwizard.AtlasDbConfigurationProvider;
 
 import io.dropwizard.Configuration;
 
 public class AtlasDbEteConfiguration extends Configuration implements AtlasDbConfigurationProvider {
     private final AtlasDbConfig atlasdb;
+    private final Optional<AtlasDbRuntimeConfig> atlasdbRuntime;
 
-    public AtlasDbEteConfiguration(@JsonProperty("atlasdb") AtlasDbConfig atlasdb) {
+    public AtlasDbEteConfiguration(@JsonProperty("atlasdb") AtlasDbConfig atlasdb,
+            @JsonProperty("atlasDbRuntime") Optional<AtlasDbRuntimeConfig> atlasDbRuntimeConfig) {
         this.atlasdb = atlasdb;
+        this.atlasdbRuntime = atlasDbRuntimeConfig;
     }
 
     @Override
     public AtlasDbConfig getAtlasDbConfig() {
         return atlasdb;
     }
+
+    @Override
+    public Optional<AtlasDbRuntimeConfig> getAtlasDbRuntimeConfig() {
+        return atlasdbRuntime;
+    }
+
 }

atlasdb-ete-tests/src/main/java/com/palantir/atlasdb/AtlasDbEteServer.java

@@ -15,6 +15,7 @@
  */
 package com.palantir.atlasdb;
 
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
@@ -28,6 +29,7 @@
 import com.palantir.atlasdb.cas.CheckAndSetSchema;
 import com.palantir.atlasdb.cas.SimpleCheckAndSetResource;
 import com.palantir.atlasdb.config.AtlasDbConfig;
+import com.palantir.atlasdb.config.AtlasDbRuntimeConfig;
 import com.palantir.atlasdb.dropwizard.AtlasDbBundle;
 import com.palantir.atlasdb.factory.TransactionManagers;
 import com.palantir.atlasdb.http.NotInitializedExceptionMapper;
@@ -77,18 +79,22 @@ public void run(AtlasDbEteConfiguration config, final Environment environment) t
     private TransactionManager tryToCreateTransactionManager(AtlasDbEteConfiguration config, Environment environment)
             throws InterruptedException {
         if (config.getAtlasDbConfig().initializeAsync()) {
-            return createTransactionManager(config.getAtlasDbConfig(), environment);
+            return createTransactionManager(config.getAtlasDbConfig(), config.getAtlasDbRuntimeConfig(), environment);
         } else {
-            return createTransactionManagerWithRetry(config.getAtlasDbConfig(), environment);
+            return createTransactionManagerWithRetry(config.getAtlasDbConfig(),
+                    config.getAtlasDbRuntimeConfig(),
+                    environment);
         }
     }
 
-    private TransactionManager createTransactionManagerWithRetry(AtlasDbConfig config, Environment environment)
+    private TransactionManager createTransactionManagerWithRetry(AtlasDbConfig config,
+            Optional<AtlasDbRuntimeConfig> atlasDbRuntimeConfig,
+            Environment environment)
             throws InterruptedException {
         Stopwatch sw = Stopwatch.createStarted();
         while (sw.elapsed(TimeUnit.SECONDS) < CREATE_TRANSACTION_MANAGER_MAX_WAIT_TIME_SECS) {
             try {
-                return createTransactionManager(config, environment);
+                return createTransactionManager(config, atlasDbRuntimeConfig, environment);
             } catch (RuntimeException e) {
                 log.warn("An error occurred while trying to create transaction manager. Retrying...", e);
                 Thread.sleep(CREATE_TRANSACTION_MANAGER_POLL_INTERVAL_SECS);
@@ -97,12 +103,14 @@ private TransactionManager createTransactionManagerWithRetry(AtlasDbConfig confi
         throw new IllegalStateException("Timed-out because we were unable to create transaction manager");
     }
 
-    private TransactionManager createTransactionManager(AtlasDbConfig config, Environment environment) {
+    private TransactionManager createTransactionManager(AtlasDbConfig config,
+            Optional<AtlasDbRuntimeConfig> atlasDbRuntimeConfigOptional, Environment environment) {
         return TransactionManagers.builder()
                 .config(config)
                 .schemas(ETE_SCHEMAS)
                 .registrar(environment.jersey()::register)
                 .userAgent("ete test")
+                .runtimeConfigSupplier(() -> atlasDbRuntimeConfigOptional)
                 .buildSerializable();
     }
 

atlasdb-ete-tests/src/test/java/com/palantir/atlasdb/ete/QosCassandraEteTestSetup.java

@@ -0,0 +1,154 @@
+/*
+ * Copyright 2017 Palantir Technologies, Inc. All rights reserved.
+ *
+ * Licensed under the BSD-3 License (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://opensource.org/licenses/BSD-3-Clause
+ *
+ * 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 com.palantir.atlasdb.ete;
+
+
+import java.net.InetSocketAddress;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Random;
+
+import org.awaitility.Awaitility;
+import org.awaitility.Duration;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.ClassRule;
+
+import com.codahale.metrics.MetricRegistry;
+import com.google.common.collect.ImmutableList;
+import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceConfig;
+import com.palantir.atlasdb.cassandra.ImmutableCassandraCredentialsConfig;
+import com.palantir.atlasdb.cassandra.ImmutableCassandraKeyValueServiceConfig;
+import com.palantir.atlasdb.config.AtlasDbConfig;
+import com.palantir.atlasdb.config.AtlasDbRuntimeConfig;
+import com.palantir.atlasdb.config.ImmutableAtlasDbConfig;
+import com.palantir.atlasdb.config.ImmutableAtlasDbRuntimeConfig;
+import com.palantir.atlasdb.config.ImmutableSweepConfig;
+import com.palantir.atlasdb.factory.TransactionManagers;
+import com.palantir.atlasdb.keyvalue.api.Cell;
+import com.palantir.atlasdb.qos.config.ImmutableQosClientConfig;
+import com.palantir.atlasdb.qos.config.ImmutableQosLimitsConfig;
+import com.palantir.atlasdb.table.description.ValueType;
+import com.palantir.atlasdb.todo.TodoSchema;
+import com.palantir.atlasdb.transaction.impl.SerializableTransactionManager;
+import com.palantir.atlasdb.util.AtlasDbMetrics;
+import com.palantir.docker.compose.DockerComposeRule;
+import com.palantir.docker.compose.configuration.ShutdownStrategy;
+import com.palantir.docker.compose.connection.Container;
+import com.palantir.docker.compose.connection.DockerPort;
+import com.palantir.docker.compose.logging.LogDirectory;
+import com.palantir.remoting.api.config.service.HumanReadableDuration;
+
+public class QosCassandraEteTestSetup {
+    private static final Random random = new Random();
+    protected static SerializableTransactionManager serializableTransactionManager;
+    protected static final int readBytesPerSecond = 10_000;
+    protected static final int writeBytesPerSecond = 10_000;
+    private static final int CASSANDRA_PORT_NUMBER = 9160;
+    protected static final int MAX_SOFT_LIMITING_SLEEP_MILLIS = 2000;
+
+    @ClassRule
+    public static DockerComposeRule docker = DockerComposeRule.builder()
+            .file("src/test/resources/cassandra-docker-compose.yml")
+            .waitingForService("cassandra", Container::areAllPortsOpen)
+            .saveLogsTo(LogDirectory.circleAwareLogDirectory(QosCassandraReadEteTest.class))
+            .shutdownStrategy(ShutdownStrategy.AGGRESSIVE_WITH_NETWORK_CLEANUP)
+            .build();
+
+    @Before
+    public void setup() {
+        AtlasDbMetrics.setMetricRegistry(new MetricRegistry());
+        ensureTransactionManagerIsCreated();
+    }
+
+    protected static void ensureTransactionManagerIsCreated() {
+        serializableTransactionManager = TransactionManagers.builder()
+                .config(getAtlasDbConfig())
+                .runtimeConfigSupplier(QosCassandraEteTestSetup::getAtlasDbRuntimeConfig)
+                .schemas(ImmutableList.of(TodoSchema.getSchema()))
+                .userAgent("qos-test")
+                .buildSerializable();
+
+        Awaitility.await()
+                .atMost(Duration.ONE_MINUTE)
+                .pollInterval(Duration.ONE_SECOND)
+                .until(serializableTransactionManager::isInitialized);
+    }
+
+    protected static void writeNTodosOfSize(int numTodos, int size) {
+        serializableTransactionManager.runTaskWithRetry((transaction) -> {
+            Map<Cell, byte[]> write = new HashMap<>();
+            for (int i = 0; i < numTodos; i++) {
+                Cell thisCell = Cell.create(ValueType.FIXED_LONG.convertFromJava(random.nextLong()),
+                        TodoSchema.todoTextColumn());
+                write.put(thisCell, ValueType.STRING.convertFromJava(getTodoOfSize(size)));
+            }
+            transaction.put(TodoSchema.todoTable(), write);
+            return null;
+        });
+    }
+
+    private static AtlasDbConfig getAtlasDbConfig() {
+        DockerPort cassandraPort = docker.containers()
+                .container("cassandra")
+                .port(CASSANDRA_PORT_NUMBER);
+
+        InetSocketAddress cassandraAddress = new InetSocketAddress(cassandraPort.getIp(),
+                cassandraPort.getExternalPort());
+
+        CassandraKeyValueServiceConfig kvsConfig = ImmutableCassandraKeyValueServiceConfig.builder()
+                .servers(ImmutableList.of(cassandraAddress))
+                .credentials(ImmutableCassandraCredentialsConfig.builder()
+                        .username("cassandra")
+                        .password("cassandra")
+                        .build())
+                .ssl(false)
+                .replicationFactor(1)
+                .autoRefreshNodes(false)
+                .build();
+
+        return ImmutableAtlasDbConfig.builder()
+                .namespace("qosete")
+                .keyValueService(kvsConfig)
+                .initializeAsync(true)
+                .build();
+    }
+
+    private static Optional<AtlasDbRuntimeConfig> getAtlasDbRuntimeConfig() {
+        return Optional.of(ImmutableAtlasDbRuntimeConfig.builder()
+                .sweep(ImmutableSweepConfig.builder().enabled(false).build())
+                .qos(ImmutableQosClientConfig.builder()
+                        .limits(ImmutableQosLimitsConfig.builder()
+                                .readBytesPerSecond(readBytesPerSecond)
+                                .writeBytesPerSecond(writeBytesPerSecond)
+                                .build())
+                        .maxBackoffSleepTime(HumanReadableDuration.milliseconds(MAX_SOFT_LIMITING_SLEEP_MILLIS))
+                        .build())
+                .build());
+    }
+
+    private static String getTodoOfSize(int size) {
+        // Note that the size of the cell for 1000 length text is actually 1050.
+        return String.join("", Collections.nCopies(size, "a"));
+    }
+
+    @After
+    public void after() {
+        serializableTransactionManager.close();
+    }
+}