Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
96.88% covered (success)
96.88%
62 / 64
80.00% covered (warning)
80.00%
8 / 10
 CRAP
0.00% covered (danger)
0.00%
0 / 1
WRStatsRateLimiter
96.88% covered (success)
96.88%
62 / 64
80.00% covered (warning)
80.00%
8 / 10
 19
0.00% covered (danger)
0.00%
0 / 1
 __construct
100.00% covered (success)
100.00%
13 / 13
100.00% covered (success)
100.00%
1 / 1
 2
 createBatch
100.00% covered (success)
100.00%
1 / 1
100.00% covered (success)
100.00%
1 / 1
 1
 peek
100.00% covered (success)
100.00%
3 / 3
100.00% covered (success)
100.00%
1 / 1
 1
 peekBatch
96.15% covered (success)
96.15%
25 / 26
0.00% covered (danger)
0.00%
0 / 1
 6
 tryIncr
100.00% covered (success)
100.00%
3 / 3
100.00% covered (success)
100.00%
1 / 1
 1
 tryIncrBatch
100.00% covered (success)
100.00%
4 / 4
100.00% covered (success)
100.00%
1 / 1
 2
 incr
100.00% covered (success)
100.00%
2 / 2
100.00% covered (success)
100.00%
1 / 1
 1
 incrBatch
100.00% covered (success)
100.00%
10 / 10
100.00% covered (success)
100.00%
1 / 1
 3
 setCurrentTime
100.00% covered (success)
100.00%
1 / 1
100.00% covered (success)
100.00%
1 / 1
 1
 resetCurrentTime
0.00% covered (danger)
0.00%
0 / 1
0.00% covered (danger)
0.00%
0 / 1
 2
1<?php
2
3namespace Wikimedia\WRStats;
4
5/**
6 * A rate limiter with a WRStats backend
7 *
8 * @since 1.39
9 */
10class WRStatsRateLimiter {
11    /** @var StatsStore */
12    private $store;
13    /** @var array<string,LimitCondition> */
14    private $conditions;
15    /** @var array<string,array> */
16    private $specs;
17    /** @var string|string[] */
18    private $prefix;
19    /** @var float|int|null */
20    private $now;
21
22    /** Default number of time buckets per action */
23    public const BUCKET_COUNT = 30;
24
25    /**
26     * @internal Use WRStatsFactory::createRateLimiter instead
27     * @param StatsStore $store
28     * @param array<string,LimitCondition> $conditions
29     * @param string|string[] $prefix
30     * @param array $options
31     */
32    public function __construct(
33        StatsStore $store,
34        $conditions,
35        $prefix = 'WRLimit',
36        $options = []
37    ) {
38        $this->store = $store;
39        $this->conditions = $conditions;
40        $this->prefix = $prefix;
41        $bucketCount = $options['bucketCount'] ?? self::BUCKET_COUNT;
42
43        $specs = [];
44        foreach ( $conditions as $name => $condition ) {
45            $specs[$name] = [
46                'sequences' => [ [
47                    'timeStep' => $condition->window / $bucketCount,
48                    'expiry' => $condition->window
49                ] ]
50            ];
51        }
52        $this->specs = $specs;
53    }
54
55    /**
56     * Create a batch object for rate limiting of multiple metrics.
57     *
58     * @param int $defaultAmount The amount to increment each metric by, if no
59     *   amount is passed to localOp/globalOp
60     * @return LimitBatch
61     */
62    public function createBatch( $defaultAmount = 1 ) {
63        return new LimitBatch( $this, $defaultAmount );
64    }
65
66    /**
67     * Check whether executing a single operation would exceed the defined limit,
68     * without incrementing the count.
69     *
70     * @param string $condName
71     * @param EntityKey|null $entityKey
72     * @param int $amount
73     * @return LimitOperationResult
74     */
75    public function peek(
76        string $condName,
77        ?EntityKey $entityKey = null,
78        $amount = 1
79    ): LimitOperationResult {
80        $actions = [ new LimitOperation( $condName, $entityKey, $amount ) ];
81        $result = $this->peekBatch( $actions );
82        return $result->getAllResults()[0];
83    }
84
85    /**
86     * Check whether executing a given set of increment operations would exceed
87     * any defined limit, without actually performing the increment.
88     *
89     * @param LimitOperation[] $operations
90     * @return LimitBatchResult
91     */
92    public function peekBatch( array $operations ) {
93        $reader = new WRStatsReader( $this->store, $this->specs, $this->prefix );
94        if ( $this->now !== null ) {
95            $reader->setCurrentTime( $this->now );
96        }
97
98        $rates = [];
99        $amounts = [];
100        foreach ( $operations as $operation ) {
101            $name = $operation->condName;
102            $cond = $this->conditions[$name] ?? null;
103            if ( $cond === null ) {
104                throw new WRStatsError( "Unrecognized metric \"$name\"" );
105            }
106            if ( !isset( $rates[$name] ) ) {
107                $range = $reader->latest( $cond->window );
108                $rates[$name] = $reader->getRate( $name, $operation->entityKey, $range );
109                $amounts[$name] = 0;
110            }
111            $amounts[$name] += $operation->amount;
112        }
113
114        $results = [];
115        foreach ( $operations as $i => $operation ) {
116            $name = $operation->condName;
117            $total = $rates[$name]->total();
118            $cond = $this->conditions[$name];
119            $results[$i] = new LimitOperationResult(
120                $cond,
121                $total,
122                $total + $amounts[$name]
123            );
124        }
125        return new LimitBatchResult( $results );
126    }
127
128    /**
129     * Check if the limit would be exceeded by incrementing the specified
130     * metric. If not, increment it.
131     *
132     * @param string $condName
133     * @param EntityKey|null $entityKey
134     * @param int $amount
135     * @return LimitOperationResult
136     */
137    public function tryIncr(
138        string $condName,
139        ?EntityKey $entityKey = null,
140        $amount = 1
141    ): LimitOperationResult {
142        $actions = [ new LimitOperation( $condName, $entityKey, $amount ) ];
143        $result = $this->tryIncrBatch( $actions );
144        return $result->getAllResults()[0];
145    }
146
147    /**
148     * Check if the limit would be exceeded by execution of the given set of
149     * increment operations. If not, perform the increments.
150     *
151     * @param LimitOperation[] $operations
152     * @return LimitBatchResult
153     */
154    public function tryIncrBatch( array $operations ) {
155        $result = $this->peekBatch( $operations );
156        if ( $result->isAllowed() ) {
157            $this->incrBatch( $operations );
158        }
159        return $result;
160    }
161
162    /**
163     * Unconditionally increment a metric.
164     *
165     * @param string $condName
166     * @param EntityKey|null $entityKey
167     * @param int $amount
168     * @return void
169     */
170    public function incr(
171        string $condName,
172        ?EntityKey $entityKey = null,
173        $amount = 1
174    ) {
175        $actions = [ new LimitOperation( $condName, $entityKey, $amount ) ];
176        $this->incrBatch( $actions );
177    }
178
179    /**
180     * Unconditionally increment a set of metrics.
181     *
182     * @param LimitOperation[] $operations
183     */
184    public function incrBatch( array $operations ) {
185        $writer = new WRStatsWriter( $this->store, $this->specs, $this->prefix );
186        if ( $this->now !== null ) {
187            $writer->setCurrentTime( $this->now );
188        }
189        foreach ( $operations as $operation ) {
190            $writer->incr(
191                $operation->condName,
192                $operation->entityKey,
193                $operation->amount
194            );
195        }
196        $writer->flush();
197    }
198
199    /**
200     * Set the current time.
201     *
202     * @param float|int $now
203     */
204    public function setCurrentTime( $now ) {
205        $this->now = $now;
206    }
207
208    /**
209     * Forget a time set with setCurrentTime(). Use the actual current time.
210     */
211    public function resetCurrentTime() {
212        $this->now = null;
213    }
214}