fix: let WP-CLI continue a rollback started from the dashboard

inc/cli/cli_media.php

@@ -43,10 +43,16 @@ private function update_images_template( $action ) {
 			'limit_exceeded' => __( 'You have reached the maximum offloading limit of images. To increase the offload limit and for more information, contact our support.', 'optimole-wp' ),
 		];
 		$settings = new Optml_Settings();
-		if ( $settings->get( 'offload_media' ) === 'disabled' ) {
+		if ( ! $settings->can_transfer_images( $action ) ) {
 			return \WP_CLI::error( __( 'You need to have the offload_media option enabled in order to use this command', 'optimole-wp' ) );
 		}
 
+		// One run moves images at a time. A background step that is moving images right now holds this
+		// lock; while this command holds it, the background steps do nothing.
+		if ( ! Optml_Media_Offload::acquire_moving_images_lock() ) {
+			return \WP_CLI::error( 'Another image transfer is running. Wait for it to finish, then run this command again.' );
+		}
+
 		Optml_Media_Offload::clear_offload_errors_meta();
 		if ( $action === 'rollback' ) {
 			Optml_Media_Offload::clear_rollback_errors_meta();
@@ -60,6 +66,7 @@ private function update_images_template( $action ) {
 		Optml_Media_Offload::record_process_meta( $number_of_images );
 		$batch = 5;
 		if ( $number_of_images === 0 ) {
+			Optml_Media_Offload::release_moving_images_lock();
 			WP_CLI::success( $strings[ $action ]['success'] );
 
 			return;
@@ -88,14 +95,17 @@ private function update_images_template( $action ) {
 				$action === 'rollback' ? Optml_Media_Offload::instance()->rollback_images( $batch ) : Optml_Media_Offload::instance()->upload_images( $batch );
 			}
 
-			if ( Optml_Media_Offload::instance()->settings->is_offload_limit_reached() ) {
+			// The offload limit stops uploads only: moving images back lowers the usage.
+			if ( $action === 'offload' && Optml_Media_Offload::instance()->settings->is_offload_limit_reached() ) {
+				Optml_Media_Offload::release_moving_images_lock();
 				WP_CLI::error( $strings['limit_exceeded'], true );
 			}
 
 			$progress->tick();
 			++$tick;
 		}
 		$progress->finish();
+		Optml_Media_Offload::release_moving_images_lock();
 		WP_CLI::line( $strings[ $action ]['success'] );
 	}
 }

inc/media_offload.php

@@ -2043,6 +2043,65 @@ public static function renew_transfer_lock( $token, $action ) {
 		return true;
 	}
 
+	/**
+	 * Take the lock of the run that is moving images right now, without waiting.
+	 *
+	 * The transfer lock above says which chain owns a transfer. It is taken when a chain is scheduled,
+	 * and again on every request while a transfer is in progress, so it cannot tell whether images are
+	 * being moved at this moment. This one is a database advisory lock: taking it is atomic, and the
+	 * database frees it when the request ends, also after a crash. A background step holds it for its
+	 * batch, and WP-CLI for its whole run, so the two never move the same images at the same time.
+	 *
+	 * @return bool False while another run holds it. True when it was taken, and also where the
+	 *              database has no advisory locks, so transfers work there as before.
+	 */
+	public static function acquire_moving_images_lock(): bool {
+		global $wpdb;
+
+		// The lock lives on one database server. Where reads are spread over several (HyperDB,
+		// LudicrousDB), every run has to ask the same one: the primary.
+		if ( method_exists( $wpdb, 'send_reads_to_masters' ) ) {
+			$wpdb->send_reads_to_masters();
+		}
+
+		$suppress = $wpdb->suppress_errors();
+		$result   = $wpdb->get_var( $wpdb->prepare( 'SELECT GET_LOCK( %s, 0 )', self::get_moving_images_lock_name() ) );
+		$wpdb->suppress_errors( $suppress );
+
+		// A fatal error skips the cleanup of the caller. On a persistent connection the lock would then
+		// outlive the request, so it is also released when the request shuts down.
+		add_action( 'shutdown', [ __CLASS__, 'release_moving_images_lock' ] );
+
+		return null === $result || 1 === (int) $result;
+	}
+
+	/**
+	 * Release the lock taken with {@see self::acquire_moving_images_lock()}.
+	 *
+	 * @return void
+	 */
+	public static function release_moving_images_lock(): void {
+		global $wpdb;
+
+		$suppress = $wpdb->suppress_errors();
+		$wpdb->get_var( $wpdb->prepare( 'SELECT RELEASE_LOCK( %s )', self::get_moving_images_lock_name() ) );
+		$wpdb->suppress_errors( $suppress );
+	}
+
+	/**
+	 * Name of the advisory lock. Such locks are shared by every database on the server, so the name
+	 * holds the database and the table prefix (which is per site on a network).
+	 *
+	 * @return string
+	 */
+	public static function get_moving_images_lock_name(): string {
+		global $wpdb;
+
+		$database = defined( 'DB_NAME' ) ? DB_NAME : '';
+
+		return 'optml_moving_images_' . md5( $database . $wpdb->prefix );
+	}
+
 	/**
 	 * Release the transfer lock if we still own it, allowing another worker to acquire it.
 	 *
@@ -2151,6 +2210,14 @@ public function start_processing_images( $action, $batch, $page, $total, $step,
 			return;
 		}
 
+		// Only one run moves images at a time. WP-CLI holds this lock while its transfer runs: it does the
+		// work of this step, and the chain is scheduled again once the transfer lock expires.
+		if ( ! self::acquire_moving_images_lock() ) {
+			self::$instance->logger->add_log( $type, 'Another run is moving images, this step is skipped.' );
+
+			return;
+		}
+
 		set_time_limit( 0 );
 
 		try {
@@ -2201,6 +2268,8 @@ public function start_processing_images( $action, $batch, $page, $total, $step,
 					$lock_token,
 				]
 			);
+		} finally {
+			self::release_moving_images_lock();
 		}
 	}
 

inc/settings.php

@@ -881,6 +881,23 @@ public function is_offload_enabled() {
 		return $this->get( 'offload_media' ) === 'enabled' || $this->get( 'rollback_status' ) !== 'disabled';
 	}
 
+	/**
+	 * Whether a bulk image transfer can run. Moving images to Optimole needs offload_media on.
+	 * Moving them back is also allowed while a rollback is in progress: "Transfer back" in the
+	 * dashboard turns offload_media off when it starts the rollback.
+	 *
+	 * @param string $action The transfer, 'offload' or 'rollback'.
+	 *
+	 * @return bool
+	 */
+	public function can_transfer_images( $action ) {
+		if ( 'rollback' === $action ) {
+			return $this->is_offload_enabled();
+		}
+
+		return $this->get( 'offload_media' ) !== 'disabled';
+	}
+
 	/**
 	 * Get cloud sites whitelist for current domain only.
 	 *