From 2318f6930e4e52dab96e9991012011ada45de0e9 Mon Sep 17 00:00:00 2001 From: Nathan Ziebart Date: Wed, 15 Nov 2017 21:30:07 +0000 Subject: [PATCH] qos config --- .../atlasdb/config/AtlasDbRuntimeConfig.java | 7 +- .../atlasdb/factory/TransactionManagers.java | 30 ++------ .../factory/TransactionManagersTest.java | 3 +- .../qos/QosServiceIntegrationTest.java | 4 +- .../atlasdb/qos/client/AtlasDbQosClient.java | 18 ++--- .../atlasdb/qos/config/QosClientConfig.java | 47 +++++++++++++ .../atlasdb/qos/config/QosLimitsConfig.java | 41 +++++++++++ .../atlasdb/qos/ratelimit/QosRateLimiter.java | 2 +- .../qos/ratelimit/QosRateLimiters.java | 43 ++++++++++++ .../QosClientConfigDeserializationTest.java | 68 +++++++++++++++++++ .../QosRuntimeConfigDeserializationTest.java | 2 +- .../{ => client}/AtlasDbQosClientTest.java | 24 ++++--- .../src/test/resources/qos-client.yml | 10 +++ .../resources/{qos.yml => qos-server.yml} | 0 14 files changed, 250 insertions(+), 49 deletions(-) create mode 100644 qos-service-impl/src/main/java/com/palantir/atlasdb/qos/config/QosClientConfig.java create mode 100644 qos-service-impl/src/main/java/com/palantir/atlasdb/qos/config/QosLimitsConfig.java create mode 100644 qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiters.java create mode 100644 qos-service-impl/src/test/java/com/palantir/atlasdb/qos/QosClientConfigDeserializationTest.java rename qos-service-impl/src/test/java/com/palantir/atlasdb/qos/{ => client}/AtlasDbQosClientTest.java (79%) create mode 100644 qos-service-impl/src/test/resources/qos-client.yml rename qos-service-impl/src/test/resources/{qos.yml => qos-server.yml} (100%) diff --git a/atlasdb-config/src/main/java/com/palantir/atlasdb/config/AtlasDbRuntimeConfig.java b/atlasdb-config/src/main/java/com/palantir/atlasdb/config/AtlasDbRuntimeConfig.java index 30ca96ce1cc..df5fbf01319 100644 --- a/atlasdb-config/src/main/java/com/palantir/atlasdb/config/AtlasDbRuntimeConfig.java +++ b/atlasdb-config/src/main/java/com/palantir/atlasdb/config/AtlasDbRuntimeConfig.java @@ -23,7 +23,7 @@ import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.fasterxml.jackson.databind.annotation.JsonSerialize; import com.palantir.atlasdb.AtlasDbConstants; -import com.palantir.remoting.api.config.service.ServiceConfiguration; +import com.palantir.atlasdb.qos.config.QosClientConfig; @JsonDeserialize(as = ImmutableAtlasDbRuntimeConfig.class) @JsonSerialize(as = ImmutableAtlasDbRuntimeConfig.class) @@ -61,7 +61,10 @@ public long getTimestampCacheSize() { return AtlasDbConstants.DEFAULT_TIMESTAMP_CACHE_SIZE; } - public abstract Optional getQosServiceConfiguration(); + @Value.Default + public QosClientConfig qos() { + return QosClientConfig.DEFAULT; + } /** * Runtime live-reloadable parameters for communicating with TimeLock. diff --git a/atlasdb-config/src/main/java/com/palantir/atlasdb/factory/TransactionManagers.java b/atlasdb-config/src/main/java/com/palantir/atlasdb/factory/TransactionManagers.java index e57ee08c3d9..fd662a8fae3 100644 --- a/atlasdb-config/src/main/java/com/palantir/atlasdb/factory/TransactionManagers.java +++ b/atlasdb-config/src/main/java/com/palantir/atlasdb/factory/TransactionManagers.java @@ -18,8 +18,6 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import java.util.function.Supplier; @@ -30,7 +28,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import com.codahale.metrics.InstrumentedScheduledExecutorService; import com.codahale.metrics.MetricRegistry; import com.fasterxml.jackson.annotation.JsonIgnore; import com.google.common.annotations.VisibleForTesting; @@ -72,11 +69,10 @@ import com.palantir.atlasdb.persistentlock.KvsBackedPersistentLockService; import com.palantir.atlasdb.persistentlock.NoOpPersistentLockService; import com.palantir.atlasdb.persistentlock.PersistentLockService; -import com.palantir.atlasdb.qos.FakeQosClient; import com.palantir.atlasdb.qos.QosClient; -import com.palantir.atlasdb.qos.QosService; import com.palantir.atlasdb.qos.client.AtlasDbQosClient; -import com.palantir.atlasdb.qos.ratelimit.QosRateLimiter; +import com.palantir.atlasdb.qos.config.QosClientConfig; +import com.palantir.atlasdb.qos.ratelimit.QosRateLimiters; import com.palantir.atlasdb.schema.generated.SweepTableFactory; import com.palantir.atlasdb.spi.AtlasDbFactory; import com.palantir.atlasdb.spi.KeyValueServiceConfig; @@ -117,9 +113,6 @@ import com.palantir.lock.impl.LockServiceImpl; import com.palantir.lock.v2.TimelockService; import com.palantir.logsafe.UnsafeArg; -import com.palantir.remoting.api.config.service.ServiceConfiguration; -import com.palantir.remoting3.clients.ClientConfigurations; -import com.palantir.remoting3.jaxrs.JaxRsClient; import com.palantir.timestamp.TimestampService; import com.palantir.timestamp.TimestampStoreInvalidator; import com.palantir.util.OptionalResolver; @@ -316,8 +309,7 @@ SerializableTransactionManager serializable() { java.util.function.Supplier runtimeConfigSupplier = () -> runtimeConfigSupplier().get().orElse(defaultRuntime); - - QosClient qosClient = getQosClient(runtimeConfigSupplier.get().getQosServiceConfiguration()); + QosClient qosClient = getQosClient(runtimeConfigSupplier.get().qos()); ServiceDiscoveringAtlasSupplier atlasFactory = new ServiceDiscoveringAtlasSupplier( @@ -411,20 +403,10 @@ SerializableTransactionManager serializable() { return transactionManager; } - private QosClient getQosClient(Optional serviceConfiguration) { - return serviceConfiguration.map(this::createAtlasDbQosClient).orElse(FakeQosClient.INSTANCE); - } - - private QosClient createAtlasDbQosClient(ServiceConfiguration serviceConfiguration) { - QosService qosService = JaxRsClient.create(QosService.class, - userAgent(), - ClientConfigurations.of(serviceConfiguration)); + private QosClient getQosClient(QosClientConfig config) { // TODO(nziebart): create a RefreshingRateLimiter - ScheduledExecutorService scheduler = new InstrumentedScheduledExecutorService( - Executors.newSingleThreadScheduledExecutor(), - AtlasDbMetrics.getMetricRegistry(), - "qos-client-executor"); - return AtlasDbQosClient.create(QosRateLimiter.create()); + QosRateLimiters rateLimiters = QosRateLimiters.create(config.limits()); + return AtlasDbQosClient.create(rateLimiters); } private static boolean areTransactionManagerInitializationPrerequisitesSatisfied( diff --git a/atlasdb-config/src/test/java/com/palantir/atlasdb/factory/TransactionManagersTest.java b/atlasdb-config/src/test/java/com/palantir/atlasdb/factory/TransactionManagersTest.java index 729b4f0e69f..5a6258b43bd 100644 --- a/atlasdb-config/src/test/java/com/palantir/atlasdb/factory/TransactionManagersTest.java +++ b/atlasdb-config/src/test/java/com/palantir/atlasdb/factory/TransactionManagersTest.java @@ -73,6 +73,7 @@ import com.palantir.atlasdb.config.TimeLockClientConfig; import com.palantir.atlasdb.factory.startup.TimeLockMigrator; import com.palantir.atlasdb.memory.InMemoryAtlasDbConfig; +import com.palantir.atlasdb.qos.config.QosClientConfig; import com.palantir.atlasdb.table.description.GenericTestSchema; import com.palantir.atlasdb.transaction.impl.SerializableTransactionManager; import com.palantir.atlasdb.util.MetricsRule; @@ -193,7 +194,7 @@ public void setup() throws JsonProcessingException { runtimeConfig = mock(AtlasDbRuntimeConfig.class); when(runtimeConfig.timestampClient()).thenReturn(ImmutableTimestampClientConfig.of(false)); - when(runtimeConfig.getQosServiceConfiguration()).thenReturn(Optional.empty()); + when(runtimeConfig.qos()).thenReturn(QosClientConfig.DEFAULT); when(runtimeConfig.timelockRuntime()).thenReturn(Optional.empty()); environment = mock(Consumer.class); diff --git a/qos-service-impl/src/integTest/java/com/palantir/atlasdb/qos/QosServiceIntegrationTest.java b/qos-service-impl/src/integTest/java/com/palantir/atlasdb/qos/QosServiceIntegrationTest.java index 51642b0c450..4d355fa2978 100644 --- a/qos-service-impl/src/integTest/java/com/palantir/atlasdb/qos/QosServiceIntegrationTest.java +++ b/qos-service-impl/src/integTest/java/com/palantir/atlasdb/qos/QosServiceIntegrationTest.java @@ -50,8 +50,8 @@ public class QosServiceIntegrationTest { @Test public void returnsConfiguredLimits() { - assertThat(service.getLimit("test")).isEqualTo(10L); - assertThat(service.getLimit("test2")).isEqualTo(20L); + assertThat(service.getLimit("test")).isEqualTo(10); + assertThat(service.getLimit("test2")).isEqualTo(20); } } diff --git a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/client/AtlasDbQosClient.java b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/client/AtlasDbQosClient.java index 997c75cce36..09dc7ca64fc 100644 --- a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/client/AtlasDbQosClient.java +++ b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/client/AtlasDbQosClient.java @@ -26,24 +26,24 @@ import com.google.common.base.Ticker; import com.palantir.atlasdb.qos.QosClient; import com.palantir.atlasdb.qos.QosMetrics; -import com.palantir.atlasdb.qos.ratelimit.QosRateLimiter; +import com.palantir.atlasdb.qos.ratelimit.QosRateLimiters; public class AtlasDbQosClient implements QosClient { private static final Logger log = LoggerFactory.getLogger(AtlasDbQosClient.class); - private final QosRateLimiter rateLimiter; + private final QosRateLimiters rateLimiters; private final QosMetrics metrics; private final Ticker ticker; - public static AtlasDbQosClient create(QosRateLimiter rateLimiter) { - return new AtlasDbQosClient(rateLimiter, new QosMetrics(), Ticker.systemTicker()); + public static AtlasDbQosClient create(QosRateLimiters rateLimiters) { + return new AtlasDbQosClient(rateLimiters, new QosMetrics(), Ticker.systemTicker()); } @VisibleForTesting - AtlasDbQosClient(QosRateLimiter rateLimiter, QosMetrics metrics, Ticker ticker) { + AtlasDbQosClient(QosRateLimiters rateLimiters, QosMetrics metrics, Ticker ticker) { this.metrics = metrics; - this.rateLimiter = rateLimiter; + this.rateLimiters = rateLimiters; this.ticker = ticker; } @@ -53,7 +53,7 @@ public T executeRead( ReadQuery query, Function weigher) throws E { int estimatedWeight = getWeight(estimatedWeigher, 1); - rateLimiter.consumeWithBackoff(estimatedWeight); + rateLimiters.read().consumeWithBackoff(estimatedWeight); // TODO(nziebart): decide what to do if we encounter a timeout exception long startTimeNanos = ticker.read(); @@ -64,7 +64,7 @@ public T executeRead( metrics.updateReadCount(); metrics.updateBytesRead(actualWeight); metrics.updateReadTimeMicros(TimeUnit.NANOSECONDS.toMicros(totalTimeNanos)); - rateLimiter.recordAdjustment(actualWeight - estimatedWeight); + rateLimiters.read().recordAdjustment(actualWeight - estimatedWeight); return result; } @@ -72,7 +72,7 @@ public T executeRead( @Override public void executeWrite(Supplier weigher, WriteQuery query) throws E { int weight = getWeight(weigher, 1); - rateLimiter.consumeWithBackoff(weight); + rateLimiters.write().consumeWithBackoff(weight); // TODO(nziebart): decide what to do if we encounter a timeout exception long startTimeNanos = ticker.read(); diff --git a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/config/QosClientConfig.java b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/config/QosClientConfig.java new file mode 100644 index 00000000000..313a02f3d7d --- /dev/null +++ b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/config/QosClientConfig.java @@ -0,0 +1,47 @@ +/* + * 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.qos.config; + +import java.util.Optional; + +import org.immutables.value.Value; + +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.annotation.JsonSerialize; +import com.palantir.remoting.api.config.service.HumanReadableDuration; +import com.palantir.remoting.api.config.service.ServiceConfiguration; + +@Value.Immutable +@JsonDeserialize(as = ImmutableQosClientConfig.class) +@JsonSerialize(as = ImmutableQosClientConfig.class) +public abstract class QosClientConfig { + + public static final QosClientConfig DEFAULT = ImmutableQosClientConfig.builder().build(); + + public abstract Optional qosService(); + + @Value.Default + public HumanReadableDuration maxBackoffSleepTime() { + return HumanReadableDuration.seconds(10); + } + + @Value.Default + public QosLimitsConfig limits() { + return QosLimitsConfig.DEFAULT_NO_LIMITS; + } + +} diff --git a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/config/QosLimitsConfig.java b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/config/QosLimitsConfig.java new file mode 100644 index 00000000000..e932632ddaa --- /dev/null +++ b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/config/QosLimitsConfig.java @@ -0,0 +1,41 @@ +/* + * 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.qos.config; + +import org.immutables.value.Value; + +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.annotation.JsonSerialize; + +@Value.Immutable +@JsonDeserialize(as = ImmutableQosLimitsConfig.class) +@JsonSerialize(as = ImmutableQosLimitsConfig.class) +public abstract class QosLimitsConfig { + + public static final QosLimitsConfig DEFAULT_NO_LIMITS = ImmutableQosLimitsConfig.builder().build(); + + @Value.Default + public int readBytesPerSecond() { + return Integer.MAX_VALUE; + } + + @Value.Default + public int writeBytesPerSecond() { + return Integer.MAX_VALUE; + } + +} diff --git a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java index 4112e03d9e1..4ea8fc5e227 100644 --- a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java +++ b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java @@ -80,7 +80,7 @@ public Duration consumeWithBackoff(int estimatedNumUnits) { /** * Records an adjustment to the original estimate of units consumed passed to {@link #consumeWithBackoff(int)}. This * should be called after a query returns, when the exact number of units consumed is known. This value may be - * positive (if the original estimate was too small) or negative (if the original estimate was too large. + * positive (if the original estimate was too small) or negative (if the original estimate was too large). */ public void recordAdjustment(int adjustmentUnits) { if (adjustmentUnits > 0) { diff --git a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiters.java b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiters.java new file mode 100644 index 00000000000..503b5b79d39 --- /dev/null +++ b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiters.java @@ -0,0 +1,43 @@ +/* + * 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.qos.ratelimit; + +import org.immutables.value.Value; + +import com.palantir.atlasdb.qos.config.QosLimitsConfig; + +@Value.Immutable +public interface QosRateLimiters { + + static QosRateLimiters create(QosLimitsConfig config) { + QosRateLimiter readLimiter = QosRateLimiter.create(); + readLimiter.updateRate(config.readBytesPerSecond()); + + QosRateLimiter writeLimiter = QosRateLimiter.create(); + writeLimiter.updateRate(config.writeBytesPerSecond()); + + return ImmutableQosRateLimiters.builder() + .read(readLimiter) + .write(writeLimiter) + .build(); + } + + QosRateLimiter read(); + + QosRateLimiter write(); + +} diff --git a/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/QosClientConfigDeserializationTest.java b/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/QosClientConfigDeserializationTest.java new file mode 100644 index 00000000000..1edfe1e9b16 --- /dev/null +++ b/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/QosClientConfigDeserializationTest.java @@ -0,0 +1,68 @@ +/* + * 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.qos; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Paths; + +import org.junit.Test; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import com.fasterxml.jackson.datatype.guava.GuavaModule; +import com.fasterxml.jackson.datatype.jdk8.Jdk8Module; +import com.palantir.atlasdb.qos.config.ImmutableQosClientConfig; +import com.palantir.atlasdb.qos.config.ImmutableQosLimitsConfig; +import com.palantir.atlasdb.qos.config.QosClientConfig; +import com.palantir.atlasdb.qos.config.QosServiceRuntimeConfig; +import com.palantir.remoting.api.config.service.HumanReadableDuration; +import com.palantir.remoting.api.config.service.ServiceConfiguration; +import com.palantir.remoting.api.config.ssl.SslConfiguration; +import com.palantir.remoting3.ext.jackson.ShimJdk7Module; + +public class QosClientConfigDeserializationTest { + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(new YAMLFactory()) + .registerModule(new GuavaModule()) + .registerModule(new ShimJdk7Module()) + .registerModule(new Jdk8Module()); + + @Test + public void canDeserializeFromYaml() throws IOException { + QosClientConfig expected = ImmutableQosClientConfig.builder() + .qosService( + ServiceConfiguration.builder() + .addUris("http://localhost:8080") + .security(SslConfiguration.of(Paths.get("trustStore.jks"))) + .build()) + .maxBackoffSleepTime(HumanReadableDuration.seconds(20)) + .limits(ImmutableQosLimitsConfig.builder() + .readBytesPerSecond(123) + .writeBytesPerSecond(456) + .build()) + .build(); + + File configFile = new File(QosServiceRuntimeConfig.class.getResource("/qos-client.yml").getPath()); + QosClientConfig config = OBJECT_MAPPER.readValue(configFile, QosClientConfig.class); + + assertThat(config).isEqualTo(expected); + } + +} diff --git a/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/QosRuntimeConfigDeserializationTest.java b/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/QosRuntimeConfigDeserializationTest.java index 2b7bb0449ae..55a662bb1cb 100644 --- a/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/QosRuntimeConfigDeserializationTest.java +++ b/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/QosRuntimeConfigDeserializationTest.java @@ -37,7 +37,7 @@ public class QosRuntimeConfigDeserializationTest { @Test public void canDeserializeQosServerConfiguration() throws IOException { - File testConfigFile = new File(QosServiceRuntimeConfig.class.getResource("/qos.yml").getPath()); + File testConfigFile = new File(QosServiceRuntimeConfig.class.getResource("/qos-server.yml").getPath()); QosServiceRuntimeConfig configuration = OBJECT_MAPPER.readValue(testConfigFile, QosServiceRuntimeConfig.class); assertThat(configuration).isEqualTo(ImmutableQosServiceRuntimeConfig.builder() diff --git a/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/AtlasDbQosClientTest.java b/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/client/AtlasDbQosClientTest.java similarity index 79% rename from qos-service-impl/src/test/java/com/palantir/atlasdb/qos/AtlasDbQosClientTest.java rename to qos-service-impl/src/test/java/com/palantir/atlasdb/qos/client/AtlasDbQosClientTest.java index f6f1ff61c74..27a8ddbda4a 100644 --- a/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/AtlasDbQosClientTest.java +++ b/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/client/AtlasDbQosClientTest.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package com.palantir.atlasdb.qos; +package com.palantir.atlasdb.qos.client; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Mockito.mock; @@ -26,8 +26,11 @@ import org.junit.Test; import com.google.common.base.Ticker; -import com.palantir.atlasdb.qos.client.AtlasDbQosClient; +import com.palantir.atlasdb.qos.QosMetrics; +import com.palantir.atlasdb.qos.QosService; +import com.palantir.atlasdb.qos.ratelimit.ImmutableQosRateLimiters; import com.palantir.atlasdb.qos.ratelimit.QosRateLimiter; +import com.palantir.atlasdb.qos.ratelimit.QosRateLimiters; public class AtlasDbQosClientTest { @@ -38,11 +41,14 @@ public class AtlasDbQosClientTest { private static final long TOTAL_TIME_MICROS = 4; private QosService qosService = mock(QosService.class); - private QosRateLimiter rateLimiter = mock(QosRateLimiter.class); + private QosRateLimiter readLimiter = mock(QosRateLimiter.class); + private QosRateLimiter writeLimiter = mock(QosRateLimiter.class); + private QosRateLimiters rateLimiters = ImmutableQosRateLimiters.builder() + .read(readLimiter).write(writeLimiter).build(); private QosMetrics metrics = mock(QosMetrics.class); private Ticker ticker = mock(Ticker.class); - private AtlasDbQosClient qosClient = new AtlasDbQosClient(rateLimiter, metrics, ticker); + private AtlasDbQosClient qosClient = new AtlasDbQosClient(rateLimiters, metrics, ticker); @Before public void setUp() { @@ -55,9 +61,9 @@ public void setUp() { public void consumesSpecifiedNumUnitsForReads() { qosClient.executeRead(() -> ESTIMATED_BYTES, () -> "foo", ignored -> ACTUAL_BYTES); - verify(rateLimiter).consumeWithBackoff(ESTIMATED_BYTES); - verify(rateLimiter).recordAdjustment(ACTUAL_BYTES - ESTIMATED_BYTES); - verifyNoMoreInteractions(rateLimiter); + verify(readLimiter).consumeWithBackoff(ESTIMATED_BYTES); + verify(readLimiter).recordAdjustment(ACTUAL_BYTES - ESTIMATED_BYTES); + verifyNoMoreInteractions(readLimiter, writeLimiter); } @Test @@ -74,8 +80,8 @@ public void recordsReadMetrics() throws TestCheckedException { public void consumesSpecifiedNumUnitsForWrites() { qosClient.executeWrite(() -> ACTUAL_BYTES, () -> { }); - verify(rateLimiter).consumeWithBackoff(ACTUAL_BYTES); - verifyNoMoreInteractions(rateLimiter); + verify(writeLimiter).consumeWithBackoff(ACTUAL_BYTES); + verifyNoMoreInteractions(readLimiter, writeLimiter); } @Test diff --git a/qos-service-impl/src/test/resources/qos-client.yml b/qos-service-impl/src/test/resources/qos-client.yml new file mode 100644 index 00000000000..163da250c04 --- /dev/null +++ b/qos-service-impl/src/test/resources/qos-client.yml @@ -0,0 +1,10 @@ +qosService: + uris: + - http://localhost:8080 + security: + trustStorePath: trustStore.jks +maxBackoffSleepTime: 20 seconds +limits: + readBytesPerSecond: 123 + writeBytesPerSecond: 456 + diff --git a/qos-service-impl/src/test/resources/qos.yml b/qos-service-impl/src/test/resources/qos-server.yml similarity index 100% rename from qos-service-impl/src/test/resources/qos.yml rename to qos-service-impl/src/test/resources/qos-server.yml