From 40da0ddecfe4917a25e45cdbbd5237283fbe8758 Mon Sep 17 00:00:00 2001
From: Gael Robin <robin.gael@gmail.com>
Date: Wed, 24 Jan 2024 16:06:18 +0100
Subject: [PATCH] Closes #6394: Reintroduce limit preload at job level

---
 inc/Engine/Preload/Database/Queries/Cache.php | 12 +++-
 .../Database/Queries/Cache/getPendingJobs.php | 25 ++++++++
 .../Database/Queries/Cache/getPendingJobs.php | 61 ++++++++++++++-----
 3 files changed, 82 insertions(+), 16 deletions(-)

diff --git a/inc/Engine/Preload/Database/Queries/Cache.php b/inc/Engine/Preload/Database/Queries/Cache.php
index 2feb6f6520..82d06fb935 100644
--- a/inc/Engine/Preload/Database/Queries/Cache.php
+++ b/inc/Engine/Preload/Database/Queries/Cache.php
@@ -338,8 +338,16 @@ public function remove_all_not_accessed_rows( float $delay = 1, string $unit = '
 	 * @return array
 	 */
 	public function get_pending_jobs( int $total = 45 ) {
+		$inprogress_count = $this->query(
+			[
+				'count'     => true,
+				'status'    => 'in-progress',
+				'is_locked' => false,
+			],
+			false
+		);
 
-		if ( $total <= 0 ) {
+		if ( $total <= 0 || $inprogress_count >= $total ) {
 			return [];
 		}
 
@@ -358,7 +366,7 @@ public function get_pending_jobs( int $total = 45 ) {
 
 		return $this->query(
 			[
-				'number'         => $total,
+				'number'         => ( $total - $inprogress_count ),
 				'status'         => 'pending',
 				'fields'         => [
 					'id',
diff --git a/tests/Fixtures/inc/Engine/Preload/Database/Queries/Cache/getPendingJobs.php b/tests/Fixtures/inc/Engine/Preload/Database/Queries/Cache/getPendingJobs.php
index 93529a3404..1c46cff0fe 100644
--- a/tests/Fixtures/inc/Engine/Preload/Database/Queries/Cache/getPendingJobs.php
+++ b/tests/Fixtures/inc/Engine/Preload/Database/Queries/Cache/getPendingJobs.php
@@ -8,6 +8,7 @@
 	'shouldReturnPendingRowsWhenInProgressLessThanTotal' => [
 		'config' => [
 			'total' => 45,
+			'in-progress' => 1,
 			'results' => [
 				$cache,
 			]
@@ -16,9 +17,33 @@
 			$cache
 		]
 	],
+	'shouldNotReturnEmptyPendingRowsWhenInProgressNotEqualsTotal' => [
+		'config' => [
+			'total' => 45,
+			'in-progress' => 44,
+			'results' => [
+				$cache,
+			],
+		],
+		'expected' => [
+			$cache
+		]
+	],
+	'shouldReturnEmptyPendingRowsWhenInProgressEqualsTotal' => [
+		'config' => [
+			'total' => 45,
+			'in-progress' => 45,
+			'results' => [
+				$cache,
+			],
+		],
+		'expected' => [
+		]
+	],
 	'negativeCountShouldReturnEmpty' => [
 		'config' => [
 			'total' => -10,
+			'in-progress' => 35,
 			'results' => [
 				$cache,
 			]
diff --git a/tests/Unit/inc/Engine/Preload/Database/Queries/Cache/getPendingJobs.php b/tests/Unit/inc/Engine/Preload/Database/Queries/Cache/getPendingJobs.php
index a7d8dee2ad..11ccf2c23f 100644
--- a/tests/Unit/inc/Engine/Preload/Database/Queries/Cache/getPendingJobs.php
+++ b/tests/Unit/inc/Engine/Preload/Database/Queries/Cache/getPendingJobs.php
@@ -24,22 +24,55 @@ protected function setUp(): void
 	 * @dataProvider configTestData
 	 */
 	public function testShouldReturnPending($config, $expected) {
+		$queue = new SplQueue();
 
-		if($config['total'] > 0) {
+		if($config['total'] > 0 && $config['in-progress'] < $config['total'] ) {
+			$queue->enqueue(
+				[
+					'params' =>
+						[
+							'count'     => true,
+							'status'    => 'in-progress',
+							'is_locked' => false,
+						],
+					'return' => $config['in-progress']
+				]
+			);
+
+			$queue->enqueue(
+				[
+					'params' =>
+						[
+							'number'         => $config['total']-$config['in-progress'],
+							'status'         => 'pending',
+							'fields'         => [
+								'id',
+								'url',
+							],
+							'job_id__not_in' => [
+								'not_in' => '',
+							],
+							'orderby'        => Filters\applied( 'rocket_preload_order' ) > 0 ? 'id' : 'modified',
+							'order'          => 'asc',
+							'is_locked' => false
+						],
+					'return' => $config['results']
+				]
+			);
+
+			$this->query->expects(self::exactly(2))->method('query')
+				->willReturnCallback(function($params) use ($queue) {
+					$dequeue = $queue->dequeue();
+					$this->assertEquals($dequeue['params'], $params);
+					return $dequeue['return'];
+				});
+		} else {
 			$this->query->expects(self::once())->method('query')->with([
-				'number'         => $config['total'],
-				'status'         => 'pending',
-				'fields'         => [
-					'id',
-					'url',
-				],
-				'job_id__not_in' => [
-					'not_in' => '',
-				],
-				'orderby'        => Filters\applied( 'rocket_preload_order' ) > 0 ? 'id' : 'modified',
-				'order'          => 'asc',
-				'is_locked' => false
-			])->willReturn($config['results']);
+				'count' => true,
+				'status' => 'in-progress',
+				'is_locked' => false,
+			])->willReturn($config['in-progress']);
+
 		}
 
 		$this->assertSame($expected, $this->query->get_pending_jobs($config['total']));