Feature/qos service api

atlasdb-api/build.gradle

@@ -6,7 +6,7 @@ apply from: "../gradle/shared.gradle"
 dependencies {
   compile project(":atlasdb-commons")
   compile project(":timestamp-api")
-  compile 'javax.ws.rs:javax.ws.rs-api:2.0.1'
+  compile group: 'javax.ws.rs', name: 'javax.ws.rs-api'
   compile group: 'org.apache.commons', name: 'commons-lang3'
   compile group: 'com.fasterxml.jackson.core', name: 'jackson-annotations'
   compile group: 'com.fasterxml.jackson.core', name: 'jackson-databind'

atlasdb-api/versions.lock

@@ -73,7 +73,6 @@
         },
         "javax.ws.rs:javax.ws.rs-api": {
             "locked": "2.0.1",
-            "requested": "2.0.1",
             "transitive": [
                 "com.palantir.atlasdb:atlasdb-commons",
                 "com.palantir.atlasdb:timestamp-api"
@@ -170,7 +169,6 @@
         },
         "javax.ws.rs:javax.ws.rs-api": {
             "locked": "2.0.1",
-            "requested": "2.0.1",
             "transitive": [
                 "com.palantir.atlasdb:atlasdb-commons",
                 "com.palantir.atlasdb:timestamp-api"

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

@@ -0,0 +1,104 @@
+/*
+ * 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.keyvalue.cassandra;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Callable;
+
+import org.apache.cassandra.thrift.AutoDelegate_Client;
+import org.apache.cassandra.thrift.Cassandra;
+import org.apache.cassandra.thrift.ColumnOrSuperColumn;
+import org.apache.cassandra.thrift.ColumnParent;
+import org.apache.cassandra.thrift.Compression;
+import org.apache.cassandra.thrift.ConsistencyLevel;
+import org.apache.cassandra.thrift.CqlResult;
+import org.apache.cassandra.thrift.InvalidRequestException;
+import org.apache.cassandra.thrift.KeyRange;
+import org.apache.cassandra.thrift.KeySlice;
+import org.apache.cassandra.thrift.Mutation;
+import org.apache.cassandra.thrift.SchemaDisagreementException;
+import org.apache.cassandra.thrift.SlicePredicate;
+import org.apache.cassandra.thrift.TimedOutException;
+import org.apache.cassandra.thrift.UnavailableException;
+import org.apache.thrift.TException;
+
+import com.palantir.atlasdb.qos.AtlasDbQosClient;
+import com.palantir.common.base.Throwables;
+import com.palantir.processors.AutoDelegate;
+
+/**
+ * Wrapper for Cassandra.Client.
+ */
+@AutoDelegate(typeToExtend = Cassandra.Client.class)
+@SuppressWarnings({"checkstyle:all", "DuplicateThrows"}) // :'(
+public class CassandraClient extends AutoDelegate_Client {
+    private final Cassandra.Client delegate;
+    private final AtlasDbQosClient qosClient;
+
+    public CassandraClient(Cassandra.Client delegate, AtlasDbQosClient qosClient) {
+        super(delegate.getInputProtocol());
+        this.delegate = delegate;
+        this.qosClient = qosClient;
+    }
+
+    @Override
+    public Cassandra.Client delegate() {
+        return delegate;
+    }
+
+    @Override
+    public Map<ByteBuffer, List<ColumnOrSuperColumn>> multiget_slice(List<ByteBuffer> keys, ColumnParent column_parent,
+            SlicePredicate predicate, ConsistencyLevel consistency_level)
+            throws InvalidRequestException, UnavailableException, TimedOutException, TException {
+        return checkLimitAndCall(() -> super.multiget_slice(keys, column_parent, predicate, consistency_level));
+    }
+
+    @Override
+    public void batch_mutate(Map<ByteBuffer, Map<String, List<Mutation>>> mutation_map,
+            ConsistencyLevel consistency_level)
+            throws InvalidRequestException, UnavailableException, TimedOutException, TException {
+        checkLimitAndCall(() -> {
+            super.batch_mutate(mutation_map, consistency_level);
+            return null;
+        });
+    }
+
+    @Override
+    public CqlResult execute_cql3_query(ByteBuffer query, Compression compression, ConsistencyLevel consistency)
+            throws InvalidRequestException, UnavailableException, TimedOutException, SchemaDisagreementException,
+            TException {
+        return checkLimitAndCall(() -> super.execute_cql3_query(query, compression, consistency));
+    }
+
+    @Override
+    public List<KeySlice> get_range_slices(ColumnParent column_parent, SlicePredicate predicate, KeyRange range,
+            ConsistencyLevel consistency_level)
+            throws InvalidRequestException, UnavailableException, TimedOutException, TException {
+        return checkLimitAndCall(() -> super.get_range_slices(column_parent, predicate, range, consistency_level));
+    }
+
+    private <T> T checkLimitAndCall(Callable<T> callable) {
+        try {
+            qosClient.checkLimit();
+            return callable.call();
+        } catch (Exception ex) {
+            throw Throwables.throwUncheckedException(ex);
+        }
+    }
+}

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

@@ -46,6 +46,7 @@
 import com.google.common.collect.Maps;
 import com.palantir.atlasdb.cassandra.CassandraCredentialsConfig;
 import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceConfig;
+import com.palantir.atlasdb.qos.AtlasDbQosClient;
 import com.palantir.common.exception.AtlasDbDependencyException;
 import com.palantir.logsafe.SafeArg;
 import com.palantir.logsafe.UnsafeArg;
@@ -66,18 +67,22 @@ public SSLSocketFactory load(InetSocketAddress host) throws Exception {
                 }
             });
 
+    private final AtlasDbQosClient qosClient;
     private final InetSocketAddress addr;
     private final CassandraKeyValueServiceConfig config;
 
-    public CassandraClientFactory(InetSocketAddress addr, CassandraKeyValueServiceConfig config) {
+    public CassandraClientFactory(AtlasDbQosClient qosClient,
+            InetSocketAddress addr,
+            CassandraKeyValueServiceConfig config) {
+        this.qosClient = qosClient;
         this.addr = addr;
         this.config = config;
     }
 
     @Override
     public Client create() throws Exception {
         try {
-            return getClient(addr, config);
+            return getClient(qosClient, addr, config);
         } catch (Exception e) {
             String message = String.format("Failed to construct client for %s/%s", addr, config.getKeyspaceOrThrow());
             if (config.usingSsl()) {
@@ -87,9 +92,10 @@ public Client create() throws Exception {
         }
     }
 
-    private static Cassandra.Client getClient(InetSocketAddress addr,
-                                              CassandraKeyValueServiceConfig config) throws Exception {
-        Client ret = getClientInternal(addr, config);
+    private static Cassandra.Client getClient(AtlasDbQosClient qosClient,
+            InetSocketAddress addr,
+            CassandraKeyValueServiceConfig config) throws Exception {
+        Client ret = getWrappedClient(qosClient, addr, config);
         try {
             ret.set_keyspace(config.getKeyspaceOrThrow());
             log.debug("Created new client for {}/{}{}{}",
@@ -105,6 +111,13 @@ private static Cassandra.Client getClient(InetSocketAddress addr,
         }
     }
 
+    private static Client getWrappedClient(AtlasDbQosClient qosClient,
+            InetSocketAddress addr,
+            CassandraKeyValueServiceConfig config)
+            throws TException {
+        return new CassandraClient(getClientInternal(addr, config), qosClient);
+    }
+
     static Cassandra.Client getClientInternal(InetSocketAddress addr, CassandraKeyValueServiceConfig config)
             throws TException {
         TSocket thriftSocket = new TSocket(addr.getHostString(), addr.getPort(), config.socketTimeoutMillis());

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

@@ -31,6 +31,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.TreeMap;
+import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ThreadLocalRandom;
@@ -50,6 +51,7 @@
 import org.slf4j.LoggerFactory;
 import org.slf4j.helpers.MessageFormatter;
 
+import com.codahale.metrics.InstrumentedScheduledExecutorService;
 import com.codahale.metrics.Meter;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -72,6 +74,10 @@
 import com.palantir.atlasdb.keyvalue.api.InsufficientConsistencyException;
 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.QosService;
+import com.palantir.atlasdb.qos.QosServiceResource;
+import com.palantir.atlasdb.util.AtlasDbMetrics;
 import com.palantir.atlasdb.util.MetricsManager;
 import com.palantir.common.base.FunctionCheckedException;
 import com.palantir.common.base.Throwables;
@@ -150,6 +156,7 @@ public void shutdown() {
     private final CassandraKeyValueServiceConfig config;
     private final Map<InetSocketAddress, CassandraClientPoolingContainer> currentPools = Maps.newConcurrentMap();
     private final StartupChecks startupChecks;
+    private final AtlasDbQosClient qosClient;
     private final ScheduledExecutorService refreshDaemon;
     private final MetricsManager metricsManager = new MetricsManager();
     private final RequestMetrics aggregateMetrics = new RequestMetrics(null);
@@ -177,14 +184,24 @@ public static CassandraClientPool create(CassandraKeyValueServiceConfig config,
 
     private static CassandraClientPoolImpl create(CassandraKeyValueServiceConfig config,
             StartupChecks startupChecks, boolean initializeAsync) {
-        CassandraClientPoolImpl cassandraClientPool = new CassandraClientPoolImpl(config, startupChecks);
+        // TODO eventually we'll want to pass this in from somewhere
+        QosService qosResource = new QosServiceResource();
+
+        ScheduledExecutorService scheduler = new InstrumentedScheduledExecutorService(
+                Executors.newSingleThreadScheduledExecutor(),
+                AtlasDbMetrics.getMetricRegistry(),
+                "qos-client-executor");
+        AtlasDbQosClient qosClient = AtlasDbQosClient.create(scheduler, qosResource, config.getKeyspaceOrThrow());
+        CassandraClientPoolImpl cassandraClientPool = new CassandraClientPoolImpl(config, startupChecks, qosClient);
         cassandraClientPool.wrapper.initialize(initializeAsync);
         return cassandraClientPool;
     }
 
-    private CassandraClientPoolImpl(CassandraKeyValueServiceConfig config, StartupChecks startupChecks) {
+    private CassandraClientPoolImpl(CassandraKeyValueServiceConfig config, StartupChecks startupChecks,
+            AtlasDbQosClient qosClient) {
         this.config = config;
         this.startupChecks = startupChecks;
+        this.qosClient = qosClient;
         this.refreshDaemon = Tracers.wrap(PTExecutors.newScheduledThreadPool(1, new ThreadFactoryBuilder()
                 .setDaemon(true)
                 .setNameFormat("CassandraClientPoolRefresh-%d")
@@ -287,7 +304,7 @@ private synchronized void refreshPool() {
     @VisibleForTesting
     void addPool(InetSocketAddress server) {
         int currentPoolNumber = cassandraHosts.indexOf(server) + 1;
-        addPool(server, new CassandraClientPoolingContainer(server, config, currentPoolNumber));
+        addPool(server, new CassandraClientPoolingContainer(qosClient, server, config, currentPoolNumber));
     }
 
     @VisibleForTesting

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

@@ -40,6 +40,7 @@
 import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
 import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceConfig;
+import com.palantir.atlasdb.qos.AtlasDbQosClient;
 import com.palantir.atlasdb.util.MetricsManager;
 import com.palantir.common.base.FunctionCheckedException;
 import com.palantir.common.pooling.PoolingContainer;
@@ -49,17 +50,19 @@
 public class CassandraClientPoolingContainer implements PoolingContainer<Client> {
     private static final Logger log = LoggerFactory.getLogger(CassandraClientPoolingContainer.class);
 
+    private final AtlasDbQosClient qosClient;
     private final InetSocketAddress host;
-    private CassandraKeyValueServiceConfig config;
+    private final CassandraKeyValueServiceConfig config;
     private final MetricsManager metricsManager = new MetricsManager();
     private final AtomicLong count = new AtomicLong();
     private final AtomicInteger openRequests = new AtomicInteger();
     private final GenericObjectPool<Client> clientPool;
 
-    public CassandraClientPoolingContainer(
+    public CassandraClientPoolingContainer(AtlasDbQosClient qosClient,
             InetSocketAddress host,
             CassandraKeyValueServiceConfig config,
             int poolNumber) {
+        this.qosClient = qosClient;
         this.host = host;
         this.config = config;
         this.clientPool = createClientPool(poolNumber);
@@ -235,7 +238,7 @@ public String toString() {
      * @param poolNumber number of the pool for metric registration.
      */
     private GenericObjectPool<Client> createClientPool(int poolNumber) {
-        CassandraClientFactory cassandraClientFactory = new CassandraClientFactory(host, config);
+        CassandraClientFactory cassandraClientFactory = new CassandraClientFactory(qosClient, host, config);
         GenericObjectPoolConfig poolConfig = new GenericObjectPoolConfig();
 
         poolConfig.setMinIdle(config.poolSize());

atlasdb-cassandra/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraClientTest.java

@@ -0,0 +1,94 @@
+/*
+ * 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.keyvalue.cassandra;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+
+import javax.naming.LimitExceededException;
+
+import org.apache.cassandra.thrift.Cassandra;
+import org.apache.cassandra.thrift.ColumnParent;
+import org.apache.cassandra.thrift.Compression;
+import org.apache.cassandra.thrift.ConsistencyLevel;
+import org.apache.cassandra.thrift.KeyRange;
+import org.apache.cassandra.thrift.SlicePredicate;
+import org.apache.thrift.TException;
+import org.junit.Before;
+import org.junit.Test;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.palantir.atlasdb.encoding.PtBytes;
+import com.palantir.atlasdb.keyvalue.cassandra.thrift.SlicePredicates;
+import com.palantir.atlasdb.qos.AtlasDbQosClient;
+
+public class CassandraClientTest {
+    private static final ByteBuffer ROW_KEY = ByteBuffer.wrap(PtBytes.toBytes("key"));
+
+    private final Cassandra.Client mockClient = mock(Cassandra.Client.class);
+    private final AtlasDbQosClient qosClient = mock(AtlasDbQosClient.class);
+
+    private static final ColumnParent TEST_TABLE = new ColumnParent("table");
+    private static final SlicePredicate SLICE_PREDICATE = SlicePredicates.create(SlicePredicates.Range.ALL,
+            SlicePredicates.Limit.ONE);
+
+    private CassandraClient client;
+
+    @Before
+    public void setUp() {
+        client = new CassandraClient(mockClient, qosClient);
+    }
+
+    @Test
+    public void multigetSliceChecksLimit() throws TException, LimitExceededException {
+        client.multiget_slice(ImmutableList.of(ROW_KEY), TEST_TABLE, SLICE_PREDICATE, ConsistencyLevel.ANY);
+
+        verify(qosClient, times(1)).checkLimit();
+        verifyNoMoreInteractions(qosClient);
+    }
+
+    @Test
+    public void batchMutateChecksLimit() throws TException, LimitExceededException {
+        client.batch_mutate(ImmutableMap.of(), ConsistencyLevel.ANY);
+
+        verify(qosClient, times(1)).checkLimit();
+        verifyNoMoreInteractions(qosClient);
+    }
+
+    @Test
+    public void executeCqlQueryChecksLimit() throws TException, LimitExceededException {
+        ByteBuffer query = ByteBuffer.wrap("SELECT * FROM test_table LIMIT 1".getBytes(StandardCharsets.UTF_8));
+        client.execute_cql3_query(query, Compression.NONE, ConsistencyLevel.ANY);
+
+        verify(qosClient, times(1)).checkLimit();
+        verifyNoMoreInteractions(qosClient);
+    }
+
+    @Test
+    public void getRangeSlicesChecksLimit() throws TException, LimitExceededException {
+        client.get_range_slices(TEST_TABLE, SLICE_PREDICATE, new KeyRange(), ConsistencyLevel.ANY);
+
+        verify(qosClient, times(1)).checkLimit();
+        verifyNoMoreInteractions(qosClient);
+    }
+}