Qos Service config

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraClientPoolImpl.java

@@ -75,6 +75,7 @@
 import com.palantir.atlasdb.keyvalue.api.TableReference;
 import com.palantir.atlasdb.keyvalue.cassandra.CassandraClientFactory.ClientCreationFailedException;
 import com.palantir.atlasdb.qos.AtlasDbQosClient;
+import com.palantir.atlasdb.qos.ImmutableQosServiceRuntimeConfig;
 import com.palantir.atlasdb.qos.QosService;
 import com.palantir.atlasdb.qos.QosServiceResource;
 import com.palantir.atlasdb.util.AtlasDbMetrics;
@@ -185,7 +186,7 @@ public static CassandraClientPool create(CassandraKeyValueServiceConfig config,
     private static CassandraClientPoolImpl create(CassandraKeyValueServiceConfig config,
             StartupChecks startupChecks, boolean initializeAsync) {
         // TODO eventually we'll want to pass this in from somewhere
-        QosService qosResource = new QosServiceResource();
+        QosService qosResource = new QosServiceResource(ImmutableQosServiceRuntimeConfig.builder().build());
 
         ScheduledExecutorService scheduler = new InstrumentedScheduledExecutorService(
                 Executors.newSingleThreadScheduledExecutor(),

atlasdb-client/src/main/java/com/palantir/atlasdb/qos/AtlasDbQosClient.java

@@ -26,7 +26,7 @@
 public class AtlasDbQosClient {
     QosService qosService;
 
-    volatile int credits;
+    volatile long credits;
 
     private AtlasDbQosClient(QosService qosService) {
         this.qosService = qosService;

atlasdb-client/src/test/java/com/palantir/atlasdb/qos/AtlasDbQosClientTest.java

@@ -34,7 +34,7 @@ public class AtlasDbQosClientTest {
 
     @Before
     public void setUp() {
-        when(qosService.getLimit("test-client")).thenReturn(1);
+        when(qosService.getLimit("test-client")).thenReturn(1L);
     }
 
     @Test

qos-service-api/build.gradle

@@ -12,5 +12,12 @@ libsDirName = file('build/artifacts')
 dependencies {
     compile group: 'javax.ws.rs', name: 'javax.ws.rs-api'
     compile group: 'com.palantir.safe-logging', name: 'safe-logging'
+    compile group: 'com.fasterxml.jackson.core', name: 'jackson-databind'
+
+    processor group: 'org.immutables', name: 'value'
+
+    testCompile group: 'org.assertj', name: 'assertj-core'
+    testCompile group: 'com.fasterxml.jackson.dataformat', name: 'jackson-dataformat-yaml'
+    testCompile group: 'com.google.guava', name: 'guava'
 }
 

qos-service-api/src/main/java/com/palantir/atlasdb/qos/QosService.java

@@ -31,5 +31,5 @@ public interface QosService {
     @POST
     @Consumes(MediaType.APPLICATION_JSON)
     @Produces(MediaType.APPLICATION_JSON)
-    int getLimit(@Safe @PathParam("client") String client);
+    long getLimit(@Safe @PathParam("client") String client);
 }

qos-service-api/src/main/java/com/palantir/atlasdb/qos/QosServiceResource.java

@@ -17,8 +17,19 @@
 package com.palantir.atlasdb.qos;
 
 public class QosServiceResource implements QosService {
+
+    private QosServiceRuntimeConfig config;
+
+    public QosServiceResource(QosServiceRuntimeConfig config) {
+        this.config = config;
+    }
+
     @Override
-    public int getLimit(String client) {
-        return Integer.MAX_VALUE;
+    public long getLimit(String client) {
+        Long currentClientLimit = config.clientLimits().get(client);
+        if (currentClientLimit == null) {
+            currentClientLimit = Long.MAX_VALUE;
+        }
+        return currentClientLimit;
     }
 }

qos-service-api/src/main/java/com/palantir/atlasdb/qos/QosServiceRuntimeConfig.java

@@ -0,0 +1,31 @@
+/*
+ * 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 java.util.Map;
+
+import org.immutables.value.Value;
+
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+
+@JsonDeserialize(as = ImmutableQosServiceRuntimeConfig.class)
+@JsonSerialize(as = ImmutableQosServiceRuntimeConfig.class)
+@Value.Immutable
+public abstract class QosServiceRuntimeConfig {
+    abstract Map<String, Long> clientLimits();
+}

qos-service-api/src/test/java/com/palantir/atlasdb/qos/QosRuntimeConfigDeserializationTest.java

@@ -0,0 +1,46 @@
+/*
+ * 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 org.junit.Test;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
+import com.fasterxml.jackson.dataformat.yaml.YAMLGenerator;
+import com.google.common.collect.ImmutableMap;
+
+public class QosRuntimeConfigDeserializationTest {
+
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(new YAMLFactory()
+            .disable(YAMLGenerator.Feature.USE_NATIVE_TYPE_ID)
+            .disable(YAMLGenerator.Feature.WRITE_DOC_START_MARKER));
+
+    @Test
+    public void canDeserializeQosServerConfiguration() throws IOException {
+        File testConfigFile = new File(QosServiceRuntimeConfig.class.getResource("/qos.yml").getPath());
+        QosServiceRuntimeConfig configuration = OBJECT_MAPPER.readValue(testConfigFile, QosServiceRuntimeConfig.class);
+
+        assertThat(configuration).isEqualTo(ImmutableQosServiceRuntimeConfig.builder()
+                .clientLimits(ImmutableMap.of("test", 10L, "test2", 20L))
+                .build());
+    }
+}

qos-service-api/src/test/java/com/palantir/atlasdb/qos/QosServiceResourceTest.java

@@ -23,7 +23,7 @@
 public class QosServiceResourceTest {
     @Test
     public void canGetLimit() {
-        QosService resource = new QosServiceResource();
+        QosService resource = new QosServiceResource(ImmutableQosServiceRuntimeConfig.builder().build());
         assertEquals(Integer.MAX_VALUE, resource.getLimit("test-client"));
     }
 }

qos-service-api/src/test/java/com/palantir/atlasdb/qos/QosServiceRuntimeConfigTest.java

@@ -0,0 +1,42 @@
+/*
+ * 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 org.junit.Test;
+
+import com.google.common.collect.ImmutableMap;
+
+public class QosServiceRuntimeConfigTest {
+    @Test
+    public void canBuildFromEmptyClientLimits() {
+        ImmutableQosServiceRuntimeConfig.builder().clientLimits(ImmutableMap.of()).build();
+    }
+
+    @Test
+    public void canBuildFromSingleClientLimit() {
+        ImmutableQosServiceRuntimeConfig.builder()
+                .clientLimits(ImmutableMap.of("test_client", 10L))
+                .build();
+    }
+
+    @Test
+    public void canBuildFromMultipleClientLimits() {
+        ImmutableQosServiceRuntimeConfig.builder()
+                .clientLimits(ImmutableMap.of("test_client", 10L, "test_client2", 100L))
+                .build();
+    }
+}

qos-service-api/src/test/resources/qos.yml

@@ -0,0 +1,3 @@
+clientLimits:
+  test: 10
+  test2: 20

timelock-server/src/main/java/com/palantir/atlasdb/timelock/TimeLockServerLauncher.java

@@ -52,7 +52,6 @@ public void initialize(Bootstrap<TimeLockServerConfiguration> bootstrap) {
     public void run(TimeLockServerConfiguration configuration, Environment environment) {
         environment.getObjectMapper().registerModule(new Jdk8Module());
         environment.jersey().register(HttpRemotingJerseyFeature.INSTANCE);
-        environment.jersey().register(new QosServiceResource());
 
         CombinedTimeLockServerConfiguration combined = TimeLockConfigMigrator.convert(configuration, environment);
         Consumer<Object> registrar = component -> environment.jersey().register(component);