Homec4science
Diffusion Phabricator (feature/better-testing) Browse

phab/support/startup/PhabricatorClientLimit.php
8567e3cacaa3feature/better-testing

Last Change
Raw File

PhabricatorClientLimit.php
View Options

<?php
abstract class PhabricatorClientLimit {
private $limitKey;
private $clientKey;
private $limit;
final public function setLimitKey($limit_key) {
$this->limitKey = $limit_key;
return $this;
}
final public function getLimitKey() {
return $this->limitKey;
}
final public function setClientKey($client_key) {
$this->clientKey = $client_key;
return $this;
}
final public function getClientKey() {
return $this->clientKey;
}
final public function setLimit($limit) {
$this->limit = $limit;
return $this;
}
final public function getLimit() {
return $this->limit;
}
final public function didConnect() {
// NOTE: We can not use pht() here because this runs before libraries
// load.
if (!function_exists('apc_fetch') && !function_exists('apcu_fetch')) {
throw new Exception(
'You can not configure connection rate limits unless APC/APCu are '.
'available. Rate limits rely on APC/APCu to track clients and '.
'connections.');
}
if ($this->getClientKey() === null) {
throw new Exception(
'You must configure a client key when defining a rate limit.');
}
if ($this->getLimitKey() === null) {
throw new Exception(
'You must configure a limit key when defining a rate limit.');
}
if ($this->getLimit() === null) {
throw new Exception(
'You must configure a limit when defining a rate limit.');
}
$points = $this->getConnectScore();
if ($points) {
$this->addScore($points);
}
$score = $this->getScore();
if (!$this->shouldRejectConnection($score)) {
// Client has not hit the limit, so continue processing the request.
return null;
}
$penalty = $this->getPenaltyScore();
if ($penalty) {
$this->addScore($penalty);
$score += $penalty;
}
return $this->getRateLimitReason($score);
}
final public function didDisconnect(array $request_state) {
$score = $this->getDisconnectScore($request_state);
if ($score) {
$this->addScore($score);
}
}
/**
* Get the number of seconds for each rate bucket.
*
* For example, a value of 60 will create one-minute buckets.
*
* @return int Number of seconds per bucket.
*/
abstract protected function getBucketDuration();
/**
* Get the total number of rate limit buckets to retain.
*
* @return int Total number of rate limit buckets to retain.
*/
abstract protected function getBucketCount();
/**
* Get the score to add when a client connects.
*
* @return double Connection score.
*/
abstract protected function getConnectScore();
/**
* Get the number of penalty points to add when a client hits a rate limit.
*
* @return double Penalty score.
*/
abstract protected function getPenaltyScore();
/**
* Get the score to add when a client disconnects.
*
* @return double Connection score.
*/
abstract protected function getDisconnectScore(array $request_state);
/**
* Get a human-readable explanation of why the client is being rejected.
*
* @return string Brief rejection message.
*/
abstract protected function getRateLimitReason($score);
/**
* Determine whether to reject a connection.
*
* @return bool True to reject the connection.
*/
abstract protected function shouldRejectConnection($score);
/**
* Get the APC key for the smallest stored bucket.
*
* @return string APC key for the smallest stored bucket.
* @task ratelimit
*/
private function getMinimumBucketCacheKey() {
$limit_key = $this->getLimitKey();
return "limit:min:{$limit_key}";
}
/**
* Get the current bucket ID for storing rate limit scores.
*
* @return int The current bucket ID.
*/
private function getCurrentBucketID() {
return (int)(time() / $this->getBucketDuration());
}
/**
* Get the APC key for a given bucket.
*
* @param int Bucket to get the key for.
* @return string APC key for the bucket.
*/
private function getBucketCacheKey($bucket_id) {
$limit_key = $this->getLimitKey();
return "limit:bucket:{$limit_key}:{$bucket_id}";
}
/**
* Add points to the rate limit score for some client.
*
* @param string Some key which identifies the client making the request.
* @param float The cost for this request; more points pushes them toward
* the limit faster.
* @return this
*/
private function addScore($score) {
$is_apcu = (bool)function_exists('apcu_fetch');
$current = $this->getCurrentBucketID();
$bucket_key = $this->getBucketCacheKey($current);
// There's a bit of a race here, if a second process reads the bucket
// before this one writes it, but it's fine if we occasionally fail to
// record a client's score. If they're making requests fast enough to hit
// rate limiting, we'll get them soon enough.
if ($is_apcu) {
$bucket = apcu_fetch($bucket_key);
} else {
$bucket = apc_fetch($bucket_key);
}
if (!is_array($bucket)) {
$bucket = array();
}
$client_key = $this->getClientKey();
if (empty($bucket[$client_key])) {
$bucket[$client_key] = 0;
}
$bucket[$client_key] += $score;
if ($is_apcu) {
apcu_store($bucket_key, $bucket);
} else {
apc_store($bucket_key, $bucket);
}
return $this;
}
/**
* Get the current rate limit score for a given client.
*
* @return float The client's current score.
* @task ratelimit
*/
private function getScore() {
$is_apcu = (bool)function_exists('apcu_fetch');
// Identify the oldest bucket stored in APC.
$min_key = $this->getMinimumBucketCacheKey();
if ($is_apcu) {
$min = apcu_fetch($min_key);
} else {
$min = apc_fetch($min_key);
}
// If we don't have any buckets stored yet, store the current bucket as
// the oldest bucket.
$cur = $this->getCurrentBucketID();
if (!$min) {
if ($is_apcu) {
apcu_store($min_key, $cur);
} else {
apc_store($min_key, $cur);
}
$min = $cur;
}
// Destroy any buckets that are older than the minimum bucket we're keeping
// track of. Under load this normally shouldn't do anything, but will clean
// up an old bucket once per minute.
$count = $this->getBucketCount();
for ($cursor = $min; $cursor < ($cur - $count); $cursor++) {
$bucket_key = $this->getBucketCacheKey($cursor);
if ($is_apcu) {
apcu_delete($bucket_key);
apcu_store($min_key, $cursor + 1);
} else {
apc_delete($bucket_key);
apc_store($min_key, $cursor + 1);
}
}
$client_key = $this->getClientKey();
// Now, sum up the client's scores in all of the active buckets.
$score = 0;
for (; $cursor <= $cur; $cursor++) {
$bucket_key = $this->getBucketCacheKey($cursor);
if ($is_apcu) {
$bucket = apcu_fetch($bucket_key);
} else {
$bucket = apc_fetch($bucket_key);
}
if (isset($bucket[$client_key])) {
$score += $bucket[$client_key];
}
}
return $score;
}
}
c4science · Help