WRStats: Add rate limiter support

Bug: T261744
Change-Id: If3e66491306f22650b31f22ccd60f7d9468316c7

11 files changed, 786 insertions, 9 deletions

diff --git a/includes/libs/WRStats/LimitBatch.php b/includes/libs/WRStats/LimitBatch.php
new file mode 100644
index 000000000000..2ba555de42e5
--- /dev/null
+++ b/includes/libs/WRStats/LimitBatch.php
@@ -0,0 +1,114 @@
+<?php
+
+namespace Wikimedia\WRStats;
+
+/**
+ * A class representing a batch of increment/peek operations on a WRStatsRateLimiter
+ *
+ * @since 1.39
+ */
+class LimitBatch {
+	/** @var WRStatsRateLimiter */
+	private $limiter;
+
+	/** @var int */
+	private $defaultAmount;
+
+	/** @var LimitOperation[] */
+	private $operations = [];
+
+	/**
+	 * @internal
+	 *
+	 * @param WRStatsRateLimiter $limiter
+	 * @param int $defaultAmount
+	 */
+	public function __construct(
+		WRStatsRateLimiter $limiter,
+		$defaultAmount
+	) {
+		$this->limiter = $limiter;
+		$this->defaultAmount = $defaultAmount;
+	}
+
+	/**
+	 * Construct a local entity key and queue an operation for it.
+	 *
+	 * @param string $condName The condition name to be incremented/tested,
+	 *   which must match one of the ones passed to createRateLimiter()
+	 * @param mixed $components Entity key component or array of components
+	 * @param int|null $amount The amount to increment by, or null to use the default
+	 * @return $this
+	 */
+	public function localOp( $condName, $components = [], $amount = null ) {
+		if ( !is_array( $components ) ) {
+			$components = [ $components ];
+		}
+		$this->queueOp(
+			$condName,
+			new LocalEntityKey( array_merge( [ $condName ], $components ) ),
+			$amount
+		);
+		return $this;
+	}
+
+	/**
+	 * Construct a global entity key and queue an operation for it.
+	 *
+	 * @param string $condName The condition name to be incremented/tested,
+	 *   which must match one of the ones passed to createRateLimiter()
+	 * @param mixed $components Entity key components
+	 * @param int|null $amount The amount, or null to use the default
+	 * @return $this
+	 */
+	public function globalOp( $condName, $components = [], $amount = null ) {
+		if ( !is_array( $components ) ) {
+			$components = [ $components ];
+		}
+		$this->queueOp(
+			$condName,
+			new GlobalEntityKey( array_merge( [ $condName ], $components ) ),
+			$amount
+		);
+		return $this;
+	}
+
+	private function queueOp( $type, $entity, $amount ) {
+		$amount = $amount ?? $this->defaultAmount;
+		if ( isset( $this->operations[$type] ) ) {
+			throw new WRStatsError( __METHOD__ .
+				': cannot queue multiple actions of the same type, ' .
+				'since the result array is indexed by type' );
+		}
+		$this->operations[$type] = new LimitOperation( $type, $entity, $amount );
+	}
+
+	/**
+	 * Execute the batch, checking each operation against the defined limit,
+	 * but don't actually increment the metrics.
+	 *
+	 * @return LimitBatchResult
+	 */
+	public function peek() {
+		return $this->limiter->peekBatch( $this->operations );
+	}
+
+	/**
+	 * Execute the batch, unconditionally incrementing all the specified metrics.
+	 */
+	public function incr() {
+		$this->limiter->incrBatch( $this->operations );
+	}
+
+	/**
+	 * Execute the batch, checking each operation against the defined limit.
+	 * If all operations are allowed, all metrics will be incremented. If some
+	 * of the operations exceed the limit, none of the metrics will be
+	 * incremented.
+	 *
+	 * @return LimitBatchResult
+	 */
+	public function tryIncr() {
+		return $this->limiter->tryIncrBatch( $this->operations );
+	}
+}
diff --git a/includes/libs/WRStats/LimitBatchResult.php b/includes/libs/WRStats/LimitBatchResult.php
new file mode 100644
index 000000000000..796968489a54
--- /dev/null
+++ b/includes/libs/WRStats/LimitBatchResult.php
@@ -0,0 +1,91 @@
+<?php
+
+namespace Wikimedia\WRStats;
+
+/**
+ * A class representing the results from a batch operation.
+ *
+ * @since 1.39
+ */
+class LimitBatchResult {
+	/** @var LimitOperationResult[] */
+	private $results;
+
+	/** @var bool|null */
+	private $allowed;
+
+	/**
+	 * @internal
+	 *
+	 * @param LimitOperationResult[] $results
+	 */
+	public function __construct( $results ) {
+		$this->results = $results;
+	}
+
+	/**
+	 * Determine whether the batch as a whole is/was allowed
+	 *
+	 * @return bool
+	 */
+	public function isAllowed() {
+		if ( $this->allowed === null ) {
+			$this->allowed = true;
+			foreach ( $this->results as $result ) {
+				if ( !$result->isAllowed() ) {
+					$this->allowed = false;
+					break;
+				}
+			}
+		}
+		return $this->allowed;
+	}
+
+	/**
+	 * Get LimitOperationResult objects for operations exceeding the limit.
+	 *
+	 * The keys will match the input array. For input arrays constructed by
+	 * LimitBatch, the keys will be the condition names.
+	 *
+	 * @return LimitOperationResult[]
+	 */
+	public function getFailedResults() {
+		$failed = [];
+		foreach ( $this->results as $i => $result ) {
+			if ( !$result->isAllowed() ) {
+				$failed[$i] = $result;
+			}
+		}
+		return $failed;
+	}
+
+	/**
+	 * Get LimitOperationResult objects for operations not exceeding the limit.
+	 *
+	 * The keys will match the input array. For input arrays constructed by
+	 * LimitBatch, the keys will be the condition names.
+	 *
+	 * @return LimitOperationResult[]
+	 */
+	public function getPassedResults() {
+		$passed = [];
+		foreach ( $this->results as $i => $result ) {
+			if ( $result->isAllowed() ) {
+				$passed[$i] = $result;
+			}
+		}
+		return $passed;
+	}
+
+	/**
+	 * Get LimitOperationResult objects for all operations in the batch.
+	 *
+	 * The keys will match the input array. For input arrays constructed by
+	 * LimitBatch, the keys will be the condition names.
+	 *
+	 * @return LimitOperationResult[]
+	 */
+	public function getAllResults() {
+		return $this->results;
+	}
+}
diff --git a/includes/libs/WRStats/LimitCondition.php b/includes/libs/WRStats/LimitCondition.php
new file mode 100644
index 000000000000..4d660ee0cfb3
--- /dev/null
+++ b/includes/libs/WRStats/LimitCondition.php
@@ -0,0 +1,37 @@
+<?php
+
+namespace Wikimedia\WRStats;
+
+/**
+ * @since 1.39
+ * @newable
+ */
+class LimitCondition {
+	/** @var int The maximum number of events */
+	public $limit;
+	/** @var float|int The number of seconds over which the number of events may occur */
+	public $window;
+
+	/**
+	 * @param int|float|string $limit The maximum number of events
+	 * @param int|float|string $window The number of seconds over which the
+	 *   number of events may occur
+	 */
+	public function __construct( $limit, $window ) {
+		$this->limit = (int)$limit;
+		$this->window = +$window;
+		if ( $this->window <= 0 ) {
+			throw new WRStatsError( __METHOD__ .
+				': window must be positive' );
+		}
+	}
+
+	/**
+	 * Get the condition as a number of events per second
+	 *
+	 * @return float|int
+	 */
+	public function perSecond() {
+		return $this->limit / $this->window;
+	}
+}
diff --git a/includes/libs/WRStats/LimitOperation.php b/includes/libs/WRStats/LimitOperation.php
new file mode 100644
index 000000000000..c00470d3629b
--- /dev/null
+++ b/includes/libs/WRStats/LimitOperation.php
@@ -0,0 +1,33 @@
+<?php
+
+namespace Wikimedia\WRStats;
+
+/**
+ * Class representing one item in a limit batch
+ *
+ * @newable
+ * @since 1.39
+ */
+class LimitOperation {
+	/** @var string */
+	public $condName;
+	/** @var EntityKey */
+	public $entityKey;
+	/** @var int */
+	public $amount;
+
+	/**
+	 * @param string $condName
+	 * @param EntityKey|null $entityKey
+	 * @param int $amount
+	 */
+	public function __construct(
+		string $condName,
+		EntityKey $entityKey = null,
+		$amount = 1
+	) {
+		$this->condName = $condName;
+		$this->entityKey = $entityKey ?? new LocalEntityKey;
+		$this->amount = $amount;
+	}
+}
diff --git a/includes/libs/WRStats/LimitOperationResult.php b/includes/libs/WRStats/LimitOperationResult.php
new file mode 100644
index 000000000000..c9de55d3f45a
--- /dev/null
+++ b/includes/libs/WRStats/LimitOperationResult.php
@@ -0,0 +1,50 @@
+<?php
+
+namespace Wikimedia\WRStats;
+
+/**
+ * Information about the result of a single item in a limit batch
+ *
+ * @since 1.39
+ */
+class LimitOperationResult {
+	/** @var LimitCondition */
+	public $condition;
+
+	/** @var int The previous metric value before the current action was executed */
+	public $prevTotal;
+
+	/** @var int The value the metric would have if the increment operation were allowed */
+	public $newTotal;
+
+	/**
+	 * @internal
+	 *
+	 * @param LimitCondition $condition
+	 * @param int $prevTotal
+	 * @param int $newTotal
+	 */
+	public function __construct( LimitCondition $condition, $prevTotal, $newTotal ) {
+		$this->condition = $condition;
+		$this->prevTotal = $prevTotal;
+		$this->newTotal = $newTotal;
+	}
+
+	/**
+	 * Whether the operation was/is allowed.
+	 *
+	 * @return bool
+	 */
+	public function isAllowed() {
+		return $this->newTotal <= $this->condition->limit;
+	}
+
+	/**
+	 * Get a string representing the object, for testing or debugging
+	 *
+	 * @return string
+	 */
+	public function dump() {
+		return "LimitActionResult{{$this->newTotal}/{$this->condition->limit}}";
+	}
+}
diff --git a/includes/libs/WRStats/README.md b/includes/libs/WRStats/README.md
index 9e84bdaf35cf..589fb1b16a8a 100644
--- a/includes/libs/WRStats/README.md
+++ b/includes/libs/WRStats/README.md
@@ -1,5 +1,7 @@
 WRStats is a library for production-compatible metric storage and retrieval.
 
+## Data model
+
 Currently, only counters are supported, in the RRDTool sense of an
 increment-only value with some support for calculating derivatives (rates).
 
@@ -7,10 +9,16 @@ Memcached is the intended data store. Since memcached does not support
 floating-point increment, metrics can be transparently scaled by a resolution
 factor.
 
-The counter is split into time window "buckets". To read a rate, all buckets in
-the requested time range are fetched. When the boundary of the time range does
-not exactly coincide with the boundary of a bucket, interpolation is applied,
-to scale down the counter value from the bucket.
+Each metric can contain multiple time-series sequences. A sequence is a set of
+counter values, equally spaced in time, with a fixed retention period. So for
+example, a metric might contain a sequence with a value for each second,
+expiring after one minute, and a sequence with a value for each minute,
+expiring after one hour.
+
+To read a rate, a sequence is selected, then all buckets in the requested time
+range are fetched. When the boundary of the time range does not exactly
+coincide with the boundary of a bucket, interpolation is applied, to scale down
+the counter value from the bucket.
 
 Interpolation assumes that the rate of events in the time window is constant.
 If the end of the bucket would be after the current time, the current time is
@@ -18,7 +26,31 @@ used as the bucket end time for interpolation purposes. In other words, the
 number of events in the future is assumed to be zero; all events are assumed to
 be in the past.
 
-Usage:
+## Storage keys
+
+Storage keys are composed of an array of components, conventionally joined by
+the storage class into a string key separated by colons:
+
+    <prefix> : <metric name> : <sequence name> : <time> [ : <entity key> ]
+
+* The prefix is a string or array of strings which are fixed for a given
+  factory. The prefix is supposed to identify the caller.
+* The metric name identifies the metric in the configuration.
+* The sequence name identifies the sequence. If there is only one unnamed
+  sequence, its name is the empty string. Subsequent sequences will be named
+  after their key in the configuration array. An explicit sequence name can be
+  given if desired for stability.
+* The time bucket is identified by an integer.
+* Additional key components can be supplied. This "entity key" is used to
+  distinguish different instances of a metric which share the same
+  configuration. For example, if you want to count the number of edits by
+  each user, the entity key could include the user ID.
+
+Entity keys may either be local (LocalEntityKey) or global (GlobalEntityKey).
+This flag is passed down to the storage class. If a local entity key is
+supplied, MediaWiki will prefix the key with the name of the wiki.
+
+## Usage
 
 ```php
 $specs = [
@@ -39,15 +71,43 @@ $specs = [
       ] ]
    ],
 ];
+$entity = new LocalEntityKey( [ 'user_id' => 1 ] );
+
 $wrstatsFactory = MediaWikiServices::getInstance()->getWRStatsFactory();
 $writer = $wrstatsFactory->createWriter( $specs );
-$writer->incr( 'a', null, 1 );
-$writer->incr( 'b', null, 25 );
+$writer->incr( 'a', $entity, 1 );
+$writer->incr( 'b', $entity, 25 );
 $writer->flush();
 
 $reader = $wrstatsFactory->createReader( $specs );
-$rateA = $reader->getRate( 'a', null, $reader->latest( 60 ) );
-$rateB = $reader->getRate( 'b', null, $reader->latest( 120 ) );
+$rateA = $reader->getRate( 'a', $entity, $reader->latest( 60 ) );
+$rateB = $reader->getRate( 'b', $entity, $reader->latest( 120 ) );
 print $rateA->perSecond();
 print $rateB->perMinute();
 ```
+
+## Rate limiter
+
+WRStatsRateLimiter uses underlying WRStats metrics to implement rate limiting
+functionality. The rate limiter allows you to limit the number of actions in a
+rolling time window.
+
+WRStats does not require a synchronous data store, so the counter may overshoot
+its limit slightly before actions are prevented.
+
+Typical usage:
+
+```php
+$wrstatsFactory = MediaWikiServices::getInstance()->getWRStatsFactory();
+$rateLimiter = $wrstatsFactory->createRateLimiter(
+    'by_user' => [
+        new LimitCondition( 10, 60 ) // 10 edits per minute
+    ],
+    'MyLimitCaller'
+);
+$entity = new LocalEntityKey( [ 'user_id', 1 ] );
+
+if ( $rateLimiter->tryIncr( 'by_user', $entity )->isAllowed() ) {
+    // perform the action
+}
+```
diff --git a/includes/libs/WRStats/WRStatsFactory.php b/includes/libs/WRStats/WRStatsFactory.php
index a7d38dfe278e..1410f7cf490a 100644
--- a/includes/libs/WRStats/WRStatsFactory.php
+++ b/includes/libs/WRStats/WRStatsFactory.php
@@ -64,4 +64,24 @@ class WRStatsFactory {
 	public function createReader( $specs, $prefix = 'WRStats' ) {
 		return new WRStatsReader( $this->store, $specs, $prefix );
 	}
+
+	/**
+	 * Create a rate limiter.
+	 *
+	 * @param LimitCondition[] $conditions An array in which the key is the
+	 *   condition name, and the value is a LimitCondition describing the limit.
+	 * @param string|string[] $prefix A string or array of strings to prefix
+	 *   before storage keys.
+	 * @param array $options An associative array of options:
+	 *   - bucketCount: Each window is divided into this many time buckets.
+	 *     Fetching the current count will typically result in a request for
+	 *     this many keys.
+	 * @return WRStatsRateLimiter
+	 */
+	public function createRateLimiter(
+		$conditions, $prefix = 'WRStats', $options = []
+	) {
+		return new WRStatsRateLimiter(
+			$this->store, $conditions, $prefix, $options );
+	}
 }
diff --git a/includes/libs/WRStats/WRStatsRateLimiter.php b/includes/libs/WRStats/WRStatsRateLimiter.php
new file mode 100644
index 000000000000..1b915e3ce868
--- /dev/null
+++ b/includes/libs/WRStats/WRStatsRateLimiter.php
@@ -0,0 +1,216 @@
+<?php
+
+namespace Wikimedia\WRStats;
+
+/**
+ * A rate limiter with a WRStats backend
+ */
+class WRStatsRateLimiter {
+	/** @var StatsStore */
+	private $store;
+	/** @var LimitCondition[] */
+	private $conditions;
+	/** @var array */
+	private $specs;
+	/** @var string|string[] */
+	private $prefix;
+	/** @var float|int|null */
+	private $now;
+
+	/** Default number of time buckets per action */
+	public const BUCKET_COUNT = 30;
+
+	/**
+	 * @internal Use WRStatsFactory
+	 *
+	 * @see WRStatsFactory::createRateLimiter for full parameter documentation.
+	 *
+	 * @param StatsStore $store
+	 * @param LimitCondition[] $conditions
+	 * @param string|string[] $prefix
+	 * @param array $options
+	 */
+	public function __construct(
+		StatsStore $store,
+		$conditions,
+		$prefix = 'WRLimit',
+		$options = []
+	) {
+		$this->store = $store;
+		$this->conditions = $conditions;
+		$this->prefix = $prefix;
+		$bucketCount = $options['bucketCount'] ?? self::BUCKET_COUNT;
+
+		$specs = [];
+		foreach ( $conditions as $name => $condition ) {
+			$specs[$name] = [
+				'sequences' => [ [
+					'timeStep' => $condition->window / $bucketCount,
+					'expiry' => $condition->window
+				] ]
+			];
+		}
+		$this->specs = $specs;
+	}
+
+	/**
+	 * Create a batch object for rate limiting of multiple metrics.
+	 *
+	 * @param int $defaultAmount The amount to increment each metric by, if no
+	 *   amount is passed to localOp/globalOp
+	 * @return LimitBatch
+	 */
+	public function createBatch( $defaultAmount = 1 ) {
+		return new LimitBatch( $this, $defaultAmount );
+	}
+
+	/**
+	 * Check whether executing a single operation would exceed the defined limit,
+	 * without incrementing the count.
+	 *
+	 * @param string $condName
+	 * @param EntityKey|null $entityKey
+	 * @param int $amount
+	 * @return LimitOperationResult
+	 */
+	public function peek(
+		string $condName,
+		EntityKey $entityKey = null,
+		$amount = 1
+	): LimitOperationResult {
+		$actions = [ new LimitOperation( $condName, $entityKey, $amount ) ];
+		$result = $this->peekBatch( $actions );
+		return $result->getAllResults()[0];
+	}
+
+	/**
+	 * Check whether executing a given set of increment operations would exceed
+	 * any defined limit, without actually performing the increment.
+	 *
+	 * @param LimitOperation[] $operations
+	 * @return LimitBatchResult
+	 */
+	public function peekBatch( array $operations ) {
+		$reader = new WRStatsReader( $this->store, $this->specs, $this->prefix );
+		if ( $this->now !== null ) {
+			$reader->setCurrentTime( $this->now );
+		}
+
+		$rates = [];
+		$amounts = [];
+		foreach ( $operations as $operation ) {
+			$name = $operation->condName;
+			$cond = $this->conditions[$name] ?? null;
+			if ( $cond === null ) {
+				throw new WRStatsError( __METHOD__ .
+					": unrecognized metric \"$name\"" );
+			}
+			if ( !isset( $rates[$name] ) ) {
+				$range = $reader->latest( $cond->window );
+				$rates[$name] = $reader->getRate( $name, $operation->entityKey, $range );
+				$amounts[$name] = 0;
+			}
+			$amounts[$name] += $operation->amount;
+		}
+
+		$results = [];
+		foreach ( $operations as $i => $operation ) {
+			$name = $operation->condName;
+			$total = $rates[$name]->total();
+			$cond = $this->conditions[$name];
+			$results[$i] = new LimitOperationResult(
+				$cond,
+				$total,
+				$total + $amounts[$name]
+			);
+		}
+		return new LimitBatchResult( $results );
+	}
+
+	/**
+	 * Check if the limit would be exceeded by incrementing the specified
+	 * metric. If not, increment it.
+	 *
+	 * @param string $condName
+	 * @param EntityKey|null $entityKey
+	 * @param int $amount
+	 * @return LimitOperationResult
+	 */
+	public function tryIncr(
+		string $condName,
+		EntityKey $entityKey = null,
+		$amount = 1
+	): LimitOperationResult {
+		$actions = [ new LimitOperation( $condName, $entityKey, $amount ) ];
+		$result = $this->tryIncrBatch( $actions );
+		return $result->getAllResults()[0];
+	}
+
+	/**
+	 * Check if the limit would be exceeded by execution of the given set of
+	 * increment operations. If not, perform the increments.
+	 *
+	 * @param LimitOperation[] $operations
+	 * @return LimitBatchResult
+	 */
+	public function tryIncrBatch( array $operations ) {
+		$result = $this->peekBatch( $operations );
+		if ( $result->isAllowed() ) {
+			$this->incrBatch( $operations );
+		}
+		return $result;
+	}
+
+	/**
+	 * Unconditionally increment a metric.
+	 *
+	 * @param string $condName
+	 * @param EntityKey|null $entityKey
+	 * @param int $amount
+	 * @return void
+	 */
+	public function incr(
+		string $condName,
+		EntityKey $entityKey = null,
+		$amount = 1
+	) {
+		$actions = [ new LimitOperation( $condName, $entityKey, $amount ) ];
+		$this->incrBatch( $actions );
+	}
+
+	/**
+	 * Unconditionally increment a set of metrics.
+	 *
+	 * @param LimitOperation[] $operations
+	 */
+	public function incrBatch( array $operations ) {
+		$writer = new WRStatsWriter( $this->store, $this->specs, $this->prefix );
+		if ( $this->now !== null ) {
+			$writer->setCurrentTime( $this->now );
+		}
+		foreach ( $operations as $operation ) {
+			$writer->incr(
+				$operation->condName,
+				$operation->entityKey,
+				$operation->amount
+			);
+		}
+		$writer->flush();
+	}
+
+	/**
+	 * Set the current time.
+	 *
+	 * @param float|int $now
+	 */
+	public function setCurrentTime( $now ) {
+		$this->now = $now;
+	}
+
+	/**
+	 * Forget a time set with setCurrentTime(). Use the actual current time.
+	 */
+	public function resetCurrentTime() {
+		$this->now = null;
+	}
+}
diff --git a/includes/libs/WRStats/WRStatsReader.php b/includes/libs/WRStats/WRStatsReader.php
index 3af1ebbb0c84..4a5a19a7af25 100644
--- a/includes/libs/WRStats/WRStatsReader.php
+++ b/includes/libs/WRStats/WRStatsReader.php
@@ -36,6 +36,10 @@ class WRStatsReader {
 			$this->metricSpecs[$name] = new MetricSpec( $spec );
 		}
 		$this->prefixComponents = is_array( $prefix ) ? $prefix : [ $prefix ];
+		if ( !count( $this->prefixComponents ) ) {
+			throw new WRStatsError( __METHOD__ .
+				': there must be at least one prefix component' );
+		}
 	}
 
 	/**
diff --git a/includes/libs/WRStats/WRStatsWriter.php b/includes/libs/WRStats/WRStatsWriter.php
index c083142bf714..4431d3508111 100644
--- a/includes/libs/WRStats/WRStatsWriter.php
+++ b/includes/libs/WRStats/WRStatsWriter.php
@@ -34,6 +34,10 @@ class WRStatsWriter {
 			$this->metricSpecs[$name] = new MetricSpec( $spec );
 		}
 		$this->prefixComponents = is_array( $prefix ) ? $prefix : [ $prefix ];
+		if ( !count( $this->prefixComponents ) ) {
+			throw new WRStatsError( __METHOD__ .
+				': there must be at least one prefix component' );
+		}
 	}
 
 	/**
diff --git a/tests/phpunit/unit/includes/libs/WRStats/WRStatsRateLimiterTest.php b/tests/phpunit/unit/includes/libs/WRStats/WRStatsRateLimiterTest.php
new file mode 100644
index 000000000000..9f3bba23bb51
--- /dev/null
+++ b/tests/phpunit/unit/includes/libs/WRStats/WRStatsRateLimiterTest.php
@@ -0,0 +1,148 @@
+<?php
+
+namespace Wikimedia\WRStats;
+
+use PHPUnit\Framework\TestCase;
+
+/**
+ * @covers \Wikimedia\WRStats\WRStatsRateLimiter
+ * @covers \Wikimedia\WRStats\LimitOperationResult
+ * @covers \Wikimedia\WRStats\LimitBatch
+ * @covers \Wikimedia\WRStats\LimitBatchResult
+ * @covers \Wikimedia\WRStats\LimitCondition
+ * @covers \Wikimedia\WRStats\LimitOperation
+ */
+class WRStatsRateLimiterTest extends TestCase {
+	public function testTryIncrBatch() {
+		$store = new ArrayStatsStore;
+		$conds = [
+			'cond1' => new LimitCondition( 3, 60 ),
+		];
+		$rateLimiter = new WRStatsRateLimiter( $store, $conds );
+		$rateLimiter->setCurrentTime( 1000 );
+		$actions = [ new LimitOperation( 'cond1' ) ];
+
+		$batchResult = $rateLimiter->tryIncrBatch( $actions );
+		$this->assertTrue( $batchResult->isAllowed() );
+		$this->assertSame( "LimitActionResult{1/3}",
+			$batchResult->getAllResults()[0]->dump() );
+
+		$this->assertTrue( $rateLimiter->tryIncrBatch( $actions )->isAllowed() );
+		$this->assertTrue( $rateLimiter->tryIncrBatch( $actions )->isAllowed() );
+
+		$batchResult = $rateLimiter->tryIncrBatch( $actions );
+		$this->assertFalse( $batchResult->isAllowed() );
+		$this->assertCount( 0, $batchResult->getPassedResults() );
+		$failed = $batchResult->getFailedResults();
+		$this->assertCount( 1, $failed );
+		$this->assertSame( "LimitActionResult{4/3}", $failed[0]->dump() );
+		$this->assertSame( 3, $failed[0]->prevTotal );
+		$this->assertSame( 4, $failed[0]->newTotal );
+	}
+
+	public function testMultiEntity() {
+		$store = new ArrayStatsStore;
+		$conds = [
+			'cond1' => new LimitCondition( 1, 60 ),
+		];
+		$entity1 = new LocalEntityKey( [ 'user_id', 1 ] );
+		$entity2 = new LocalEntityKey( [ 'user_id', 2 ] );
+		$rateLimiter = new WRStatsRateLimiter( $store, $conds );
+		$rateLimiter->setCurrentTime( 1000 );
+		$this->assertTrue(
+			$rateLimiter->tryIncrBatch( [
+				new LimitOperation( 'cond1', $entity1 )
+			] )->isAllowed()
+		);
+		$this->assertTrue(
+			$rateLimiter->tryIncrBatch( [
+				new LimitOperation( 'cond1', $entity2 )
+			] )->isAllowed()
+		);
+		$rateLimiter->incr( 'cond1', $entity1, 10 );
+		$rateLimiter->incr( 'cond1', $entity2, 20 );
+		$res = $rateLimiter->tryIncr( 'cond1', $entity1 );
+		$this->assertSame( 'LimitActionResult{12/1}', $res->dump() );
+		$res = $rateLimiter->tryIncr( 'cond1', $entity2 );
+		$this->assertSame( 'LimitActionResult{22/1}', $res->dump() );
+	}
+
+	public function testMultiEntityBatch() {
+		$store = new ArrayStatsStore;
+		$conds = [
+			'id' => new LimitCondition( 2, 60 ),
+			'ip' => new LimitCondition( 3, 60 )
+		];
+		$rateLimiter = new WRStatsRateLimiter( $store, $conds );
+		$rateLimiter->setCurrentTime( 1000 );
+
+		// Increment both metrics to 1
+		$res = $rateLimiter->createBatch()
+			->localOp( 'id', 1 )
+			->globalOp( 'ip', '127.0.0.1' )
+			->tryIncr();
+		$this->assertTrue( $res->isAllowed() );
+		$this->assertSame( 'LimitActionResult{1/2}',
+			$res->getAllResults()['id']->dump() );
+		$this->assertSame( 'LimitActionResult{1/3}',
+			$res->getAllResults()['ip']->dump() );
+
+		// Increment to 2, testing peek/incr
+		$batch = $rateLimiter->createBatch()
+			->localOp( 'id', 1 )
+			->globalOp( 'ip', '127.0.0.1' );
+		$this->assertTrue( $batch->peek()->isAllowed() );
+		$batch->incr();
+
+		// Increment to 3, in which the id metric is expected to fail
+		$res = $rateLimiter->createBatch()
+			->localOp( 'id', 1 )
+			->globalOp( 'ip', '127.0.0.1' )
+			->tryIncr();
+		$this->assertFalse( $res->isAllowed() );
+		$this->assertCount( 1, $res->getPassedResults() );
+		$failed = $res->getFailedResults();
+		$this->assertCount( 1, $failed );
+		$all = $res->getAllResults();
+		$this->assertSame( 'LimitActionResult{3/2}', $all['id']->dump() );
+		$this->assertSame( 'LimitActionResult{3/3}', $all['ip']->dump() );
+	}
+
+	public function testTimeDependent() {
+		$store = new ArrayStatsStore;
+		$conds = [
+			'cond1' => new LimitCondition( 3, 100 ),
+		];
+		$rateLimiter = new WRStatsRateLimiter( $store, $conds, 'WRLimit',
+			[ 'bucketCount' => 10 ] );
+
+		$rateLimiter->setCurrentTime( 1000 );
+		$this->assertTrue( $rateLimiter->tryIncr( 'cond1' )->isAllowed() );
+		$rateLimiter->setCurrentTime( 1010 );
+		$this->assertTrue( $rateLimiter->tryIncr( 'cond1' )->isAllowed() );
+		$rateLimiter->setCurrentTime( 1020 );
+		$this->assertTrue( $rateLimiter->tryIncr( 'cond1' )->isAllowed() );
+		$rateLimiter->setCurrentTime( 1090 );
+		$this->assertSame( 'LimitActionResult{4/3}',
+			$rateLimiter->peek( 'cond1' )->dump() );
+
+		$rateLimiter->setCurrentTime( 1100 );
+		// The range goes back to 1000 so the count is still 3
+		$this->assertSame( 3,
+			$rateLimiter->peek( 'cond1' )->prevTotal );
+
+		$rateLimiter->setCurrentTime( 1103 );
+		// The range goes back to 1000 so the count should be 2.7, rounded up to 3
+		$this->assertSame( 3,
+			$rateLimiter->peek( 'cond1' )->prevTotal );
+
+		$rateLimiter->setCurrentTime( 1109 );
+		// The range goes back to 1009 so the first bucket is interpolated and
+		// rounded down to zero, so the count is 2 and there is room for
+		// another action
+		$this->assertTrue( $rateLimiter->peek( 'cond1' )->isAllowed() );
+		// But not room for 2 actions
+		$this->assertSame( 'LimitActionResult{4/3}',
+			$rateLimiter->tryIncr( 'cond1', null, 2 )->dump() );
+	}
+}