Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
90.91% covered (success)
90.91%
50 / 55
70.00% covered (warning)
70.00%
7 / 10
 CRAP
0.00% covered (danger)
0.00%
0 / 1
SearchAfter
90.91% covered (success)
90.91%
50 / 55
70.00% covered (warning)
70.00%
7 / 10
 23.40
0.00% covered (danger)
0.00%
0 / 1
 __construct
85.71% covered (warning)
85.71%
6 / 7
0.00% covered (danger)
0.00%
0 / 1
 3.03
 current
100.00% covered (success)
100.00%
3 / 3
100.00% covered (success)
100.00%
1 / 1
 2
 key
100.00% covered (success)
100.00%
1 / 1
100.00% covered (success)
100.00%
1 / 1
 1
 next
80.00% covered (warning)
80.00%
8 / 10
0.00% covered (danger)
0.00%
0 / 1
 4.13
 initializeSearchAfter
100.00% covered (success)
100.00%
1 / 1
100.00% covered (success)
100.00%
1 / 1
 1
 runSearch
84.62% covered (warning)
84.62%
11 / 13
0.00% covered (danger)
0.00%
0 / 1
 4.06
 doSearch
100.00% covered (success)
100.00%
8 / 8
100.00% covered (success)
100.00%
1 / 1
 3
 rewind
100.00% covered (success)
100.00%
7 / 7
100.00% covered (success)
100.00%
1 / 1
 2
 valid
100.00% covered (success)
100.00%
1 / 1
100.00% covered (success)
100.00%
1 / 1
 1
 calcBackoff
100.00% covered (success)
100.00%
4 / 4
100.00% covered (success)
100.00%
1 / 1
 2
1<?php
2
3declare( strict_types = 1 );
4namespace CirrusSearch\Elastica;
5
6use CirrusSearch\LogChannel;
7use Elastica\Exception\ExceptionInterface as ElasticaExceptionInterface;
8use Elastica\Exception\RuntimeException;
9use Elastica\Query;
10use Elastica\ResultSet;
11use Elastica\Search;
12use InvalidArgumentException;
13use MediaWiki\Logger\LoggerFactory;
14
15class SearchAfter implements \Iterator {
16    private const MAX_BACKOFF_SEC = 120;
17    private const MICROSEC_PER_SEC = 1_000_000;
18    /** @var Search */
19    private $search;
20    /** @var Query */
21    private $baseQuery;
22    /** @var ?ResultSet */
23    private $currentResultSet;
24    /** @var ?int */
25    private $currentPage;
26    /** @var float[] Sequence of second length backoffs to use for retries */
27    private $backoff;
28    /** @var array Initial value for search_after */
29    private array $initialSearchAfter = [];
30
31    /**
32     * @param Search $search
33     * @param int $numRetries The number of retries to perform on each iteration
34     * @param float $backoffFactor Scales the backoff duration, backoff calculated as
35     *   {backoffFactor} * 2^({retry} - 1) which gives, with no scaling, [0.5, 1, 2, 4, 8, ...]
36     */
37    public function __construct( Search $search, int $numRetries = 12, float $backoffFactor = 1. ) {
38        $this->search = $search;
39        $this->baseQuery = clone $search->getQuery();
40        if ( !$this->baseQuery->hasParam( 'sort' ) ) {
41            throw new InvalidArgumentException( 'ScrollAfter query must have a sort' );
42        }
43        if ( $numRetries < 0 ) {
44            throw new InvalidArgumentException( '$numRetries must be >= 0' );
45        }
46        $this->backoff = $this->calcBackoff( $numRetries, $backoffFactor );
47    }
48
49    public function current(): ResultSet {
50        if ( $this->currentResultSet === null ) {
51            throw new RuntimeException( 'Iterator is in an invalid state and must be rewound' );
52        }
53        return $this->currentResultSet;
54    }
55
56    public function key(): int {
57        return $this->currentPage ?? 0;
58    }
59
60    public function next(): void {
61        if ( $this->currentResultSet !== null ) {
62            if ( count( $this->currentResultSet ) === 0 ) {
63                return;
64            }
65            $lastHit = $this->currentResultSet[count( $this->currentResultSet ) - 1];
66            $this->search->getQuery()->setParam( 'search_after', $lastHit->getSort() );
67        } elseif ( $this->currentPage !== -1 ) {
68            // iterator is in failed state
69            return;
70        }
71        // ensure if runSearch throws the iterator becomes invalid
72        $this->currentResultSet = null;
73        $this->currentResultSet = $this->runSearch();
74        $this->currentPage++;
75    }
76
77    public function initializeSearchAfter( array $searchAfter ): void {
78        $this->initialSearchAfter = $searchAfter;
79    }
80
81    private function runSearch(): ResultSet {
82        foreach ( $this->backoff as $backoffSec ) {
83            try {
84                return $this->doSearch();
85            } catch ( ElasticaExceptionInterface $e ) {
86                LoggerFactory::getInstance( LogChannel::DEFAULT )->warning(
87                    "Exception thrown during SearchAfter iteration. Retrying in {backoffSec}s.",
88                    [
89                        'exception' => $e,
90                        'backoffSec' => $backoffSec,
91                    ]
92                );
93                if ( !defined( 'MW_PHPUNIT_TEST' ) ) {
94                    usleep( (int)( $backoffSec * self::MICROSEC_PER_SEC ) );
95                }
96            }
97        }
98        // Final attempt after exhausting retries.
99        return $this->doSearch();
100    }
101
102    private function doSearch(): ResultSet {
103        $rs = $this->search->search();
104        if ( $rs->getResponse()->getStatus() >= 400 ) {
105            throw new RuntimeException( "Search request returned HTTP {$rs->getResponse()->getStatus()}: " .
106                                        print_r( $rs->getResponse()->getData(), true ) );
107        }
108        if ( !isset( $rs->getResponse()->getData()['_shards'] ) ) {
109            throw new RuntimeException( 'Incoherent search response: ' .
110                                        print_r( $rs->getResponse()->getData(), true ) );
111        }
112        return $rs;
113    }
114
115    public function rewind(): void {
116        // Use -1 so that on increment the first page is 0
117        $this->currentPage = -1;
118        $this->currentResultSet = null;
119        $query = clone $this->baseQuery;
120        if ( $this->initialSearchAfter ) {
121            $query->setParam( 'search_after', $this->initialSearchAfter );
122        }
123        $this->search->setQuery( $query );
124        // rewind performs the first query
125        $this->next();
126    }
127
128    public function valid(): bool {
129        return count( $this->currentResultSet ?? [] ) > 0;
130    }
131
132    private function calcBackoff( int $maxRetries, float $backoffFactor ): array {
133        $backoff = [];
134        for ( $retry = 0; $retry < $maxRetries; $retry++ ) {
135            $backoff[$retry] = min( $backoffFactor * pow( 2, $retry - 1 ), self::MAX_BACKOFF_SEC );
136        }
137        return $backoff;
138    }
139}