diff --git a/__tests__/unit-tests/test-lock.php b/__tests__/unit-tests/test-lock.php new file mode 100644 index 0000000..4c3247d --- /dev/null +++ b/__tests__/unit-tests/test-lock.php @@ -0,0 +1,179 @@ +locks as $lock ) { + Lock::reset_lock( $lock ); + } + } + + public function tearDown(): void { + foreach ( $this->locks as $lock ) { + Lock::reset_lock( $lock ); + } + + parent::tearDown(); + } + + public function test_concurrent_acquisition_with_limit_one_admits_one_worker() { + $lock = 'concurrent-acquisition'; + + $this->assertTrue( Lock::check_lock( $lock, 1 ) ); + $this->assertFalse( Lock::check_lock( $lock, 1 ) ); + $this->assertSame( 1, Lock::get_lock_value( $lock ) ); + } + + public function test_rejected_admission_does_not_change_lock_value() { + $lock = 'rejected-admission'; + + $this->assertTrue( Lock::check_lock( $lock, 1 ) ); + $this->assertFalse( Lock::check_lock( $lock, 1 ) ); + $this->assertSame( 1, Lock::get_lock_value( $lock ) ); + } + + public function test_failed_lock_release_can_be_retried() { + global $wpdb; + + $lock = 'failed-release'; + $reflection = new \ReflectionProperty( Lock::class, 'acquired_locks' ); + $reflection->setAccessible( true ); + $original_locks = $reflection->getValue(); + $original_wpdb = $wpdb; + $reflection->setValue( + array( + $lock => array( 1 ), + ) + ); + $wpdb = new class() { + public $options = 'wp_options'; + public $queries = 0; + + public function prepare( $query, ...$args ) { + return $query; + } + + public function query( $query ) { + ++$this->queries; + + return 1 === $this->queries ? false : 0; + } + }; + + try { + $this->assertFalse( Lock::free_lock( $lock ) ); + $this->assertSame( array( 1 ), $reflection->getValue()[ $lock ] ); + $this->assertTrue( Lock::free_lock( $lock ) ); + $this->assertSame( array(), $reflection->getValue()[ $lock ] ); + } finally { + $wpdb = $original_wpdb; + $reflection->setValue( $original_locks ); + } + } + + public function test_stale_lock_recovery_records_the_new_admission() { + global $wpdb; + + $lock = 'stale-lock'; + $wpdb->update( + $wpdb->options, + array( 'option_value' => '1:' . ( time() - MINUTE_IN_SECONDS ) . ':1' ), + array( 'option_name' => 'a8ccc_lock_stale-lock' ) + ); + + $this->assertTrue( Lock::check_lock( $lock, 1, 1 ) ); + $this->assertSame( 1, Lock::get_lock_value( $lock ) ); + $this->assertGreaterThan( time() - 2, Lock::get_lock_timestamp( $lock ) ); + } + + public function test_old_generation_cleanup_does_not_release_recovered_lock() { + global $wpdb; + + $lock = 'old-generation-cleanup'; + $this->assertTrue( Lock::check_lock( $lock, 1 ) ); + $wpdb->update( + $wpdb->options, + array( 'option_value' => '1:' . ( time() - MINUTE_IN_SECONDS ) . ':1' ), + array( 'option_name' => 'a8ccc_lock_old-generation-cleanup' ) + ); + $this->assertTrue( Lock::check_lock( $lock, 1, 1 ) ); + $state = explode( ':', $wpdb->get_var( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = 'a8ccc_lock_old-generation-cleanup'" ) ); + + $reflection = new \ReflectionProperty( Lock::class, 'acquired_locks' ); + $reflection->setAccessible( true ); + $reflection->setValue( + array( + $lock => array( (int) $state[2], 1 ), + ) + ); + + $this->assertTrue( Lock::free_lock( $lock ) ); + $this->assertSame( 1, Lock::get_lock_value( $lock ) ); + } + + public function test_old_generation_cleanup_does_not_release_a_recreated_lock() { + global $wpdb; + + $lock = 'recreated-lock'; + $this->assertTrue( Lock::check_lock( $lock, 1 ) ); + $state = explode( ':', $wpdb->get_var( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = 'a8ccc_lock_recreated-lock'" ) ); + $old_generation = (int) $state[2]; + $this->assertTrue( Lock::free_lock( $lock ) ); + $this->assertNull( $wpdb->get_var( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = 'a8ccc_lock_recreated-lock'" ) ); + + $this->assertTrue( Lock::check_lock( $lock, 1 ) ); + $state = explode( ':', $wpdb->get_var( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = 'a8ccc_lock_recreated-lock'" ) ); + $replacement_state = implode( ':', $state ); + $replacement_generation = (int) $state[2]; + + $reflection = new \ReflectionProperty( Lock::class, 'acquired_locks' ); + $reflection->setAccessible( true ); + $reflection->setValue( + array( + $lock => array( $replacement_generation, $old_generation ), + ) + ); + + $this->assertTrue( Lock::free_lock( $lock ) ); + $this->assertSame( $replacement_state, $wpdb->get_var( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = 'a8ccc_lock_recreated-lock'" ) ); + } + + public function test_old_generation_cleanup_does_not_release_a_reset_lock() { + global $wpdb; + + $lock = 'reset-lock'; + $this->assertTrue( Lock::check_lock( $lock, 1 ) ); + $state = explode( ':', $wpdb->get_var( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = 'a8ccc_lock_reset-lock'" ) ); + $old_generation = (int) $state[2]; + $this->assertTrue( Lock::reset_lock( $lock ) ); + $this->assertTrue( Lock::check_lock( $lock, 1 ) ); + $state = explode( ':', $wpdb->get_var( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = 'a8ccc_lock_reset-lock'" ) ); + $replacement_state = implode( ':', $state ); + $replacement_generation = (int) $state[2]; + + $reflection = new \ReflectionProperty( Lock::class, 'acquired_locks' ); + $reflection->setAccessible( true ); + $reflection->setValue( + array( + $lock => array( $replacement_generation, $old_generation ), + ) + ); + + $this->assertTrue( Lock::free_lock( $lock ) ); + $this->assertSame( $replacement_state, $wpdb->get_var( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = 'a8ccc_lock_reset-lock'" ) ); + } +} diff --git a/includes/class-events.php b/includes/class-events.php index 38508ba..d9e7d77 100644 --- a/includes/class-events.php +++ b/includes/class-events.php @@ -42,9 +42,6 @@ class Events extends Singleton { * Register hooks */ protected function class_init() { - // Prime lock cache if not present. - Lock::prime_lock( self::LOCK ); - // Prepare environment as early as possible. $earliest_action = did_action( 'muplugins_loaded' ) ? 'plugins_loaded' : 'muplugins_loaded'; add_action( $earliest_action, array( $this, 'prepare_environment' ) ); @@ -241,9 +238,6 @@ public function run_event( $timestamp, $action, $instance, $force = false ) { // Limit how many events are processed concurrently, unless explicitly bypassed. if ( ! $force ) { - // Prepare event-level lock. - $this->prime_event_action_lock( $event ); - if ( ! $this->can_run_event( $event ) ) { /* translators: 1: Event action, 2: Event arguments */ $error_message = sprintf( __( 'No resources available to run the job with action `%1$s` and arguments `%2$s`.', 'automattic-cron-control' ), $event->get_action(), maybe_serialize( $event->get_args() ) ); @@ -300,10 +294,6 @@ public function run_event( $timestamp, $action, $instance, $force = false ) { return $return; } - private function prime_event_action_lock( Event $event ): void { - Lock::prime_lock( $this->get_lock_key_for_event_action( $event ), JOB_LOCK_EXPIRY_IN_MINUTES * \MINUTE_IN_SECONDS ); - } - // Checks concurrency locks, deciding if the event can be run at this moment. private function can_run_event( Event $event ): bool { // Limit to one concurrent execution of a specific action by default. @@ -348,11 +338,7 @@ private function reset_event_lock( Event $event ): bool { $lock_key = $this->get_lock_key_for_event_action( $event ); $expires = JOB_LOCK_EXPIRY_IN_MINUTES * \MINUTE_IN_SECONDS; - if ( isset( $this->concurrent_action_whitelist[ $event->get_action() ] ) ) { - return Lock::free_lock( $lock_key, $expires ); - } else { - return Lock::reset_lock( $lock_key, $expires ); - } + return Lock::free_lock( $lock_key, $expires ); } /** diff --git a/includes/class-lock.php b/includes/class-lock.php index b2da5e2..5d33f7b 100644 --- a/includes/class-lock.php +++ b/includes/class-lock.php @@ -11,6 +11,8 @@ * Lock class */ class Lock { + private static $acquired_locks = array(); + /** * Set a lock and limit how many concurrent jobs are permitted * @@ -20,29 +22,60 @@ class Lock { * @return bool */ public static function check_lock( $lock, $limit = null, $timeout = null ) { - // Timeout, should a process die before its lock is freed. + global $wpdb; + if ( ! is_numeric( $timeout ) ) { $timeout = LOCK_DEFAULT_TIMEOUT_IN_MINUTES * \MINUTE_IN_SECONDS; } - // Check for, and recover from, deadlock. - if ( self::get_lock_timestamp( $lock ) < time() - $timeout ) { - self::reset_lock( $lock ); - return true; - } - - // Default limit for concurrent events. if ( ! is_numeric( $limit ) ) { $limit = LOCK_DEFAULT_LIMIT; } - // Check if process can run. - if ( self::get_lock_value( $lock ) >= $limit ) { + $now = time(); + $stale_before = $now - $timeout; + $generation = self::get_next_generation(); + + // LAST_INSERT_ID() retains the generation from this atomic upsert on this database connection. + $result = $wpdb->query( + $wpdb->prepare( + "INSERT INTO `$wpdb->options` (`option_name`, `option_value`, `autoload`) + VALUES (%s, CONCAT(1, ':', %d, ':', LAST_INSERT_ID(%d)), 'no') + ON DUPLICATE KEY UPDATE `option_value` = IF( + CAST(SUBSTRING_INDEX(SUBSTRING_INDEX(`option_value`, ':', 2), ':', -1) AS UNSIGNED) < %d, + CONCAT(1, ':', %d, ':', LAST_INSERT_ID(%d)), + IF( + CAST(SUBSTRING_INDEX(`option_value`, ':', 1) AS UNSIGNED) < %d, + CONCAT(CAST(SUBSTRING_INDEX(`option_value`, ':', 1) AS UNSIGNED) + 1, ':', %d, ':', LAST_INSERT_ID(CAST(SUBSTRING_INDEX(`option_value`, ':', -1) AS UNSIGNED))), + `option_value` + ) + )", + self::get_key( $lock ), + $now, + $generation, + $stale_before, + $now, + $generation, + $limit, + $now + ) + ); + + if ( $result <= 0 ) { + return false; + } + + $generation = (int) $wpdb->insert_id; + if ( $generation < 1 ) { return false; - } else { - wp_cache_incr( self::get_key( $lock ) ); - return true; } + + if ( ! isset( self::$acquired_locks[ $lock ] ) ) { + self::$acquired_locks[ $lock ] = array(); + } + self::$acquired_locks[ $lock ][] = $generation; + + return true; } /** @@ -53,49 +86,69 @@ public static function check_lock( $lock, $limit = null, $timeout = null ) { * @return bool */ public static function free_lock( $lock, $expires = 0 ) { - if ( self::get_lock_value( $lock ) > 1 ) { - wp_cache_decr( self::get_key( $lock ) ); - } else { - wp_cache_set( self::get_key( $lock ), 0, null, $expires ); + global $wpdb; + + if ( empty( self::$acquired_locks[ $lock ] ) ) { + return false; } - wp_cache_set( self::get_key( $lock, 'timestamp' ), time(), null, $expires ); + $generation = end( self::$acquired_locks[ $lock ] ); + $now = time(); - return true; + $result = $wpdb->query( + $wpdb->prepare( + "UPDATE `$wpdb->options` + SET `option_value` = CONCAT(GREATEST(CAST(SUBSTRING_INDEX(`option_value`, ':', 1) AS UNSIGNED) - 1, 0), ':', %d, ':', %d) + WHERE `option_name` = %s + AND CAST(SUBSTRING_INDEX(`option_value`, ':', 1) AS UNSIGNED) > 0 + AND CAST(SUBSTRING_INDEX(`option_value`, ':', -1) AS UNSIGNED) = %d", + $now, + $generation, + self::get_key( $lock ), + $generation + ) + ); + + if ( $result > 0 ) { + $wpdb->query( + $wpdb->prepare( + "DELETE FROM `$wpdb->options` + WHERE `option_name` = %s + AND `option_value` = %s", + self::get_key( $lock ), + self::build_lock_state( 0, $now, $generation ) + ) + ); + } + + if ( false !== $result ) { + array_pop( self::$acquired_locks[ $lock ] ); + } + + return false !== $result; } /** * Build cache key * * @param string $lock Lock name. - * @param string $type Key type, either lock or timestamp. - * @return string|bool + * @return string */ - private static function get_key( $lock, $type = 'lock' ) { - switch ( $type ) { - case 'lock': - return "a8ccc_lock_{$lock}"; - break; - - case 'timestamp': - return "a8ccc_lock_ts_{$lock}"; - break; - } - - return false; + private static function get_key( $lock ) { + return "a8ccc_lock_{$lock}"; } /** - * Ensure lock entries are initially set + * Ensure lock entries are initially set. + * + * Lock acquisition creates a missing row atomically, so priming no longer + * needs to perform a database write. * * @param string $lock Lock name. * @param int $expires Lock expiration timestamp. * @return null */ public static function prime_lock( $lock, $expires = 0 ) { - wp_cache_add( self::get_key( $lock ), 0, null, $expires ); - wp_cache_add( self::get_key( $lock, 'timestamp' ), time(), null, $expires ); - return null; } @@ -106,7 +159,8 @@ public static function prime_lock( $lock, $expires = 0 ) { * @return int */ public static function get_lock_value( $lock ) { - return (int) wp_cache_get( self::get_key( $lock ), null, true ); + $state = self::get_lock_state( $lock ); + return null === $state ? 0 : $state['count']; } /** @@ -116,7 +170,8 @@ public static function get_lock_value( $lock ) { * @return int */ public static function get_lock_timestamp( $lock ) { - return (int) wp_cache_get( self::get_key( $lock, 'timestamp' ), null, true ); + $state = self::get_lock_state( $lock ); + return null === $state ? 0 : $state['timestamp']; } /** @@ -127,9 +182,87 @@ public static function get_lock_timestamp( $lock ) { * @return bool */ public static function reset_lock( $lock, $expires = 0 ) { - wp_cache_set( self::get_key( $lock ), 0, null, $expires ); - wp_cache_set( self::get_key( $lock, 'timestamp' ), time(), null, $expires ); + global $wpdb; - return true; + $now = time(); + $generation = self::get_next_generation(); + $result = $wpdb->query( + $wpdb->prepare( + "INSERT INTO `$wpdb->options` (`option_name`, `option_value`, `autoload`) + VALUES (%s, %s, 'no') + ON DUPLICATE KEY UPDATE `option_value` = CONCAT(0, ':', %d, ':', %d)", + self::get_key( $lock ), + self::build_lock_state( 0, $now, $generation ), + $now, + $generation + ) + ); + + unset( self::$acquired_locks[ $lock ] ); + return false !== $result; + } + + /** + * Get the lock state from its option row. + * + * @param string $lock Lock name. + * @return array|null + */ + private static function get_lock_state( $lock ) { + global $wpdb; + + $value = $wpdb->get_var( + $wpdb->prepare( + "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = %s", + self::get_key( $lock ) + ) + ); + + return self::parse_lock_state( $value ); + } + + /** + * Parse a lock state from its stored value. + * + * @param string|null $value Stored lock value. + * @return array|null + */ + private static function parse_lock_state( $value ) { + $state = explode( ':', (string) $value ); + if ( 3 !== count( $state ) || ! ctype_digit( $state[0] ) || ! ctype_digit( $state[1] ) || ! ctype_digit( $state[2] ) ) { + return null; + } + + return array( + 'count' => (int) $state[0], + 'timestamp' => (int) $state[1], + 'generation' => (int) $state[2], + ); + } + + /** + * Build a value suitable for storage in the lock option. + * + * @param int $count Lock count. + * @param int $timestamp Lock timestamp. + * @param int $generation Lock generation. + * @return string + */ + private static function build_lock_state( $count, $timestamp, $generation ) { + return "{$count}:{$timestamp}:{$generation}"; + } + + /** + * Get a non-zero lock-row generation. + * + * @return int + */ + private static function get_next_generation() { + if ( \function_exists( 'wp_rand' ) ) { + return \wp_rand( 1, \PHP_INT_MAX ); + } + + // Lock may initialize before WordPress procedural functions are loaded. + return \mt_rand( 1, \mt_getrandmax() ); } }