Prevent duplicate scheduling of image processing actions

inc/media_offload.php

@@ -56,6 +56,17 @@ class Optml_Media_Offload extends Optml_App_Replacer {
 	const POST_OFFLOADED_FLAG = 'optimole_offload_post';
 	const POST_ROLLBACK_FLAG = 'optimole_rollback_post';
 	const RETRYABLE_META_COUNTER = '_optimole_retryable_errors';
+
+	/**
+	 * Transient name for the transfer lock.
+	 */
+	const TRANSFER_LOCK_TRANSIENT = 'optml_transfer_lock';
+
+	/**
+	 * Time to live for the transfer lock, in seconds.
+	 */
+	const TRANSFER_LOCK_TTL = 600;
+
 	/**
 	 * Flag used inside wp_get_attachment url filter.
 	 *
@@ -173,7 +184,7 @@ public static function instance() {
 					add_filter( 'wp_insert_attachment_data', [ self::$instance, 'insert' ], 10, 4 );
 				}
 
-				add_action( 'optml_start_processing_images', [ self::$instance, 'start_processing_images' ], 10, 5 );
+				add_action( 'optml_start_processing_images', [ self::$instance, 'start_processing_images' ], 10, 6 );
 				add_action(
 					'optml_move_images_by_id',
 					[
@@ -200,8 +211,9 @@ public static function instance() {
 	 * @return void
 	 */
 	public function maybe_reschedule() {
+		$lock = get_transient( self::TRANSFER_LOCK_TRANSIENT );
 		// If this is in pending, we do nothing.
-		if ( self::is_scheduled( 'optml_start_processing_images' ) ) {
+		if ( false !== $lock ) {
 			return;
 		}
 		// If there is no transfer in progress, we do nothing.
@@ -1899,18 +1911,25 @@ public static function move_images( $action, $refresh ) {
 					'action' => $type,
 				];
 			}
-			$total = ceil( $count / $batch );
-			self::schedule_action(
-				time(),
-				'optml_start_processing_images',
-				[
-					$action,
-					$batch,
-					1,
-					$total,
-					$step,
-				]
-			);
+
+			// We acquire a lock to prevent multiple workers from running the same action concurrently.
+			$lock_token = self::acquire_transfer_lock( $action );
+
+			if ( false !== $lock_token ) {
+				$total = ceil( $count / $batch );
+				self::schedule_action(
+					time(),
+					'optml_start_processing_images',
+					[
+						$action,
+						$batch,
+						1,
+						$total,
+						$step,
+						$lock_token,
+					]
+				);
+			}
 		}
 
 		$response = [
@@ -1970,6 +1989,77 @@ public static function is_scheduled( $hook ) {
 		}
 	}
 
+	/**
+	 * Attempt to acquire the transfer lock for a given action.
+	 *
+	 * @param string $action The transfer action ('offload_images'|'rollback_images').
+	 *
+	 * @return string|false The lock token on success, false if another worker already holds the lock.
+	 */
+	public static function acquire_transfer_lock( $action ) {
+		$lock = get_transient( self::TRANSFER_LOCK_TRANSIENT );
+		if ( false !== $lock ) {
+			return false;
+		}
+
+		$token = wp_generate_uuid4();
+
+		set_transient(
+			self::TRANSFER_LOCK_TRANSIENT,
+			[
+				'token'  => $token,
+				'action' => $action,
+			],
+			self::TRANSFER_LOCK_TTL
+		);
+
+		return $token;
+	}
+
+	/**
+	 * Renew the transfer lock if we still own it, extending its expiration.
+	 *
+	 * @param string $token The lock token this worker was given when it started the chain.
+	 * @param string $action The transfer action currently being processed.
+	 *
+	 * @return bool True if we still own the lock and renewed it, false if ownership was lost.
+	 */
+	public static function renew_transfer_lock( $token, $action ) {
+		$lock = get_transient( self::TRANSFER_LOCK_TRANSIENT );
+
+		if ( ! is_array( $lock ) || ! isset( $lock['token'] ) || $lock['token'] !== $token ) {
+			return false;
+		}
+
+		set_transient(
+			self::TRANSFER_LOCK_TRANSIENT,
+			[
+				'token'  => $token,
+				'action' => $action,
+			],
+			self::TRANSFER_LOCK_TTL
+		);
+
+		return true;
+	}
+
+	/**
+	 * Release the transfer lock if we still own it, allowing another worker to acquire it.
+	 *
+	 * @param string $token The lock token to release.
+	 *
+	 * @return void
+	 */
+	public static function release_transfer_lock( $token ) {
+		$lock = get_transient( self::TRANSFER_LOCK_TRANSIENT );
+
+		if ( ! is_array( $lock ) || ! isset( $lock['token'] ) || $lock['token'] !== $token ) {
+			return;
+		}
+
+		delete_transient( self::TRANSFER_LOCK_TRANSIENT );
+	}
+
 	/**
 	 * Start Processing Images by IDs
 	 *
@@ -2029,14 +2119,21 @@ public function move_single_image( $action, $id ) {
 	 * @param int    $page The page of images to process.
 	 * @param int    $total The total number of pages.
 	 * @param int    $step The current step.
+	 * @param string $lock_token The transfer lock token owned by this processing chain.
 	 *
 	 * @return void
 	 */
-	public function start_processing_images( $action, $batch, $page, $total, $step ) {
+	public function start_processing_images( $action, $batch, $page, $total, $step, $lock_token = '' ) {
 		$option = 'offload_images' === $action ? 'offloading_status' : 'rollback_status';
 		$type   = 'offload_images' === $action ? 'offload' : 'rollback';
 
 		if ( self::$instance->settings->get( $option ) === 'disabled' ) {
+			self::release_transfer_lock( $lock_token );
+			return;
+		}
+
+		// If we don't own the lock anymore, stop processing.
+		if ( ! self::renew_transfer_lock( $lock_token, $action ) ) {
 			return;
 		}
 
@@ -2048,6 +2145,9 @@ public function start_processing_images( $action, $batch, $page, $total, $step )
 
 			self::$instance->settings->update( 'show_offload_finish_notice', $type );
 
+			// Transfer completed successfully: release the lock.
+			self::release_transfer_lock( $lock_token );
+
 			return;
 		}
 
@@ -2080,10 +2180,12 @@ public function start_processing_images( $action, $batch, $page, $total, $step )
 					$page,
 					$total,
 					$step,
+					$lock_token,
 				]
 			);
 		} catch ( Exception $e ) {
 			// Reschedule the cron to run again after a delay. Sometimes memory limit is exausted.
+			// This is a retryable error, so the lock is kept rather than released.
 			$delay_in_seconds = 10;
 			self::$instance->logger->add_log( $type, $e->getMessage() );
 
@@ -2096,6 +2198,7 @@ public function start_processing_images( $action, $batch, $page, $total, $step )
 					$page,
 					$total,
 					$step,
+					$lock_token,
 				]
 			);
 		}

tests/test-media.php

@@ -639,4 +639,104 @@ public function test_alter_attachment_image_src() {
 			}
 		}
 	}
+
+	/**
+	 * Two overlapping "start transfer" requests must not be able to claim the transfer lock.
+	 */
+	public function test_acquire_transfer_lock_blocks_concurrent_start() {
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+
+		$first  = Optml_Media_Offload::acquire_transfer_lock( 'offload_images' );
+		$second = Optml_Media_Offload::acquire_transfer_lock( 'offload_images' );
+
+		$this->assertNotFalse( $first );
+		$this->assertFalse( $second );
+
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+	}
+
+	/**
+	 * A lock left behind by a worker that died mid-batch (fatal error, killed process) must
+	 * eventually become reclaimable, otherwise a crashed transfer would be stuck forever.
+	 */
+	public function test_acquire_transfer_lock_reclaims_expired_lock() {
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+
+		$first = Optml_Media_Offload::acquire_transfer_lock( 'offload_images' );
+		$this->assertNotFalse( $first );
+
+		// Simulate a worker that died mid-batch: force the transient's timeout into the past.
+		update_option( '_transient_timeout_' . Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT, time() - 10 );
+
+		$this->assertFalse( false !== get_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT ) );
+
+		$second = Optml_Media_Offload::acquire_transfer_lock( 'offload_images' );
+
+		$this->assertNotFalse( $second );
+		$this->assertNotEquals( $first, $second );
+
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+	}
+
+	/**
+	 * Only the current lock owner (matching token) may renew it - otherwise a worker that
+	 * already lost the lock to a reclaiming worker could revive its stale ownership.
+	 */
+	public function test_renew_transfer_lock_requires_owned_token() {
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+
+		$token = Optml_Media_Offload::acquire_transfer_lock( 'offload_images' );
+		$this->assertNotFalse( $token );
+
+		$this->assertFalse( Optml_Media_Offload::renew_transfer_lock( 'not-the-owner', 'offload_images' ) );
+
+		$this->assertTrue( Optml_Media_Offload::renew_transfer_lock( $token, 'offload_images' ) );
+
+		$lock_after = get_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+		$this->assertSame( $token, $lock_after['token'] );
+
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+	}
+
+	/**
+	 * Only the current lock owner (matching token) may release it - a stray/late release call
+	 * from a worker that already lost ownership must not tear down another worker's lock.
+	 */
+	public function test_release_transfer_lock_only_releases_owned_lock() {
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+
+		$token = Optml_Media_Offload::acquire_transfer_lock( 'offload_images' );
+		$this->assertNotFalse( $token );
+
+		Optml_Media_Offload::release_transfer_lock( 'not-the-owner' );
+		$this->assertTrue( false !== get_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT ) );
+
+		Optml_Media_Offload::release_transfer_lock( $token );
+		$this->assertFalse( false !== get_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT ) );
+	}
+
+	/**
+	 * Test that two overlapping move_images() calls do not schedule two processing chains for the same transfer.
+	 */
+	public function test_move_images_does_not_schedule_duplicate_processing_chain() {
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+
+		Optml_Media_Offload::instance()->settings->update( 'rollback_status', 'enabled' );
+		Optml_Media_Offload::instance()->settings->update( 'offloading_status', 'disabled' );
+
+		Optml_Media_Offload::move_images( 'rollback_images', false );
+		Optml_Media_Offload::move_images( 'rollback_images', false );
+
+		// Count the number of scheduled processing chains for the transfer. There should be exactly one.
+		$scheduled_count = 0;
+		foreach ( _get_cron_array() as $timestamp => $hooks ) {
+			if ( isset( $hooks['optml_start_processing_images'] ) ) {
+				$scheduled_count += count( $hooks['optml_start_processing_images'] );
+			}
+		}
+
+		$this->assertEquals( 1, $scheduled_count );
+
+		delete_transient( Optml_Media_Offload::TRANSFER_LOCK_TRANSIENT );
+	}
 }