WRStatsRateLimiter.php

Go to the documentation of this file.

<?php
 
namespace Wikimedia\WRStats;
 
class WRStatsRateLimiter {
    private $store;
    private $conditions;
    private $specs;
    private $prefix;
    private $now;
 
    public const BUCKET_COUNT = 30;
 
    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;
    }
 
    public function createBatch( $defaultAmount = 1 ) {
        return new LimitBatch( $this, $defaultAmount );
    }
 
    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];
    }
 
    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( "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 );
    }
 
    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];
    }
 
    public function tryIncrBatch( array $operations ) {
        $result = $this->peekBatch( $operations );
        if ( $result->isAllowed() ) {
            $this->incrBatch( $operations );
        }
        return $result;
    }
 
    public function incr(
        string $condName,
        EntityKey $entityKey = null,
        $amount = 1
    ) {
        $actions = [ new LimitOperation( $condName, $entityKey, $amount ) ];
        $this->incrBatch( $actions );
    }
 
    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();
    }
 
    public function setCurrentTime( $now ) {
        $this->now = $now;
    }
 
    public function resetCurrentTime() {
        $this->now = null;
    }
}