Add per partition rate limit table options

Added `maxReadsPerSecond` and `maxWritesPerSecond` options to `TableOptions`
which will allow convenient creation of CREATE and ALTER statements with
rate limit options for read and write operations.
Both maxReadsPerSecond and maxWritesPerSecond are optional - omitting one of
them means "no limit" for that type of operation.

Fixes scylladb#166

driver-core/src/main/java/com/datastax/driver/core/schemabuilder/TableOptions.java

@@ -66,6 +66,10 @@ public abstract class TableOptions<T extends TableOptions> extends SchemaStateme
 
  private Optional<Boolean> replicateOnWrite = Optional.absent();
 
+ private Optional<Integer> maxReadsPerSecond = Optional.absent();
+
+ private Optional<Integer> maxWritesPerSecond = Optional.absent();
+
  private Optional<SpeculativeRetryValue> speculativeRetry = Optional.absent();
 
  private Optional<Boolean> cdc = Optional.absent();
@@ -348,6 +352,36 @@ public T replicateOnWrite(Boolean replicateOnWrite) {
  return self;
  }
 
+ /**
+ * Sets rate limit for read operations in table option "per_partition_rate_limit". NOTE: Due to
+ * ScyllaDB’s distributed nature, tracking per-partition request rates is not perfect and the
+ * actual rate of accepted requests may be higher up to a factor of keyspace’s RF. This feature
+ * should not be used to enforce precise limits but rather serve as an overload protection
+ * feature.
+ *
+ * @param maxReadsPerSecond rate limit for read operations
+ * @return this {@code TableOptions} object.
+ */
+ public T maxReadsPerSecond(int maxReadsPerSecond) {
+ this.maxReadsPerSecond = Optional.of(maxReadsPerSecond);
+ return self;
+ }
+
+ /**
+ * Sets rate limit for write operations in table option "per_partition_rate_limit". NOTE: Due to
+ * ScyllaDB’s distributed nature, tracking per-partition request rates is not perfect and the
+ * actual rate of accepted requests may be higher up to a factor of keyspace’s RF. This feature
+ * should not be used to enforce precise limits but rather serve as an overload protection
+ * feature.
+ *
+ * @param maxWritesPerSecond rate limit for write operations
+ * @return this {@code TableOptions} object.
+ */
+ public T maxWritesPerSecond(int maxWritesPerSecond) {
+ this.maxWritesPerSecond = Optional.of(maxWritesPerSecond);
+ return self;
+ }
+
  /**
  * To override normal read timeout when read_repair_chance is not 1.0, sending another request to
  * read, choose one of these values and use the property to create or alter the table:
@@ -518,6 +552,25 @@ private List<String> buildCommonOptions() {
  options.add("replicate_on_write = " + replicateOnWrite.get());
  }
 
+ if (maxReadsPerSecond.isPresent() || maxWritesPerSecond.isPresent()) {
+ StringBuilder sBuilder = new StringBuilder("per_partition_rate_limit = {");
+
+ if (maxReadsPerSecond.isPresent()) {
+ sBuilder.append("'max_reads_per_second': ").append(maxReadsPerSecond.get());
+
+ if (maxWritesPerSecond.isPresent()) {
+ sBuilder.append(", ");
+ }
+ }
+
+ if (maxWritesPerSecond.isPresent()) {
+ sBuilder.append("'max_writes_per_second': ").append(maxWritesPerSecond.get());
+ }
+
+ sBuilder.append("}");
+ options.add(sBuilder.toString());
+ }
+
  if (speculativeRetry.isPresent()) {
  options.add("speculative_retry = " + speculativeRetry.get().value());
  }

driver-core/src/test/java/com/datastax/driver/core/schemabuilder/AlterTest.java

@@ -121,6 +121,8 @@ public void should_alter_table_options() throws Exception {
  .memtableFlushPeriodInMillis(12)
  .populateIOCacheOnFlush(true)
  .replicateOnWrite(true)
+ .maxReadsPerSecond(123)
+ .maxWritesPerSecond(456)
  .readRepairChance(0.42)
  .speculativeRetry(always())
  .cdc(true);
@@ -147,6 +149,7 @@ public void should_alter_table_options() throws Exception {
  + "AND populate_io_cache_on_flush = true "
  + "AND read_repair_chance = 0.42 "
  + "AND replicate_on_write = true "
+ + "AND per_partition_rate_limit = {'max_reads_per_second': 123, 'max_writes_per_second': 456} "
  + "AND speculative_retry = 'ALWAYS' "
  + "AND cdc = true");
 

driver-core/src/test/java/com/datastax/driver/core/schemabuilder/CreateTest.java

@@ -351,6 +351,30 @@ public void should_create_table_with_compact_storage() throws Exception {
  + "WITH COMPACT STORAGE");
  }
 
+ @Test(groups = "unit")
+ public void should_omit_write_rate_limit_option() throws Exception {
+ // When
+ SchemaStatement statement =
+ createTable("test")
+ .addPartitionKey("id", DataType.bigint())
+ .addClusteringColumn("col1", DataType.uuid())
+ .addClusteringColumn("col2", DataType.uuid())
+ .addColumn("name", DataType.text())
+ .withOptions()
+ .maxReadsPerSecond(123);
+
+ // Then
+ assertThat(statement.getQueryString())
+ .isEqualTo(
+ "\n\tCREATE TABLE test(\n\t\t"
+ + "id bigint,\n\t\t"
+ + "col1 uuid,\n\t\t"
+ + "col2 uuid,\n\t\t"
+ + "name text,\n\t\t"
+ + "PRIMARY KEY(id, col1, col2))\n\t"
+ + "WITH per_partition_rate_limit = {'max_reads_per_second': 123}");
+ }
+
  @Test(groups = "unit")
  public void should_create_table_with_all_options() throws Exception {
  // When
@@ -378,6 +402,8 @@ public void should_create_table_with_all_options() throws Exception {
  .populateIOCacheOnFlush(true)
  .readRepairChance(0.05)
  .replicateOnWrite(true)
+ .maxReadsPerSecond(123)
+ .maxWritesPerSecond(456)
  .speculativeRetry(always())
  .cdc(true);
 
@@ -403,6 +429,7 @@ public void should_create_table_with_all_options() throws Exception {
  + "AND populate_io_cache_on_flush = true "
  + "AND read_repair_chance = 0.05 "
  + "AND replicate_on_write = true "
+ + "AND per_partition_rate_limit = {'max_reads_per_second': 123, 'max_writes_per_second': 456} "
  + "AND speculative_retry = 'ALWAYS' "
  + "AND cdc = true AND CLUSTERING ORDER BY(col1 ASC, col2 DESC) AND COMPACT STORAGE");
  }