Skip to content

Commit 33187e6

Browse files
fix: make direct workflow retry atomic (#3086)
Co-authored-by: homeboy-ci[bot] <266378653+homeboy-ci[bot]@users.noreply.github.com>
1 parent 625c12f commit 33187e6

5 files changed

Lines changed: 203 additions & 18 deletions

File tree

inc/Abilities/Job/RetryJobAbility.php

Lines changed: 64 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
namespace DataMachine\Abilities\Job;
1212

1313
use DataMachine\Core\DirectJobEnqueuer;
14+
use DataMachine\Core\DirectOperationRecoveryPolicy;
1415
use DataMachine\Core\JobRetryPolicy;
1516

1617
defined( 'ABSPATH' ) || exit;
@@ -137,9 +138,11 @@ public function execute( array $input ): array {
137138
}
138139

139140
if ( 'processing' === $previous_status ) {
140-
$generation = (int) ( $job['operation_generation'] ?? 0 );
141-
$token = (string) ( $job['operation_claim_token'] ?? '' );
142-
if ( $generation > 0 && '' !== $token && ( new DirectJobEnqueuer( $this->db_jobs ) )->hasLiveGenerationAction( $job_id, $generation, $token ) ) {
141+
$generation = (int) ( $job['operation_generation'] ?? 0 );
142+
$token = (string) ( $job['operation_claim_token'] ?? '' );
143+
$enqueuer = new DirectJobEnqueuer( $this->db_jobs );
144+
$live_execution = $generation > 0 && '' !== $token ? $enqueuer->liveGenerationExecution( $job_id, $generation, $token ) : 'none';
145+
if ( 'none' !== $live_execution ) {
143146
return array(
144147
'success' => false,
145148
'job_id' => $job_id,
@@ -149,6 +152,64 @@ public function execute( array $input ): array {
149152
'error' => sprintf( 'Job %d still has live execution for generation %d; retry after it exits or is recovered.', $job_id, $generation ),
150153
);
151154
}
155+
156+
$missing_action = DirectOperationRecoveryPolicy::diagnose(
157+
$job,
158+
$live_execution,
159+
DirectOperationRecoveryPolicy::recordedActionExists( (int) ( $job['operation_action_id'] ?? 0 ) )
160+
);
161+
if ( is_array( $missing_action ) ) {
162+
if ( ! empty( $job['operation_effects_begun_at'] ) ) {
163+
return array(
164+
'success' => false,
165+
'job_id' => $job_id,
166+
'previous_status' => $previous_status,
167+
'retryable' => false,
168+
'error_code' => 'job_effects_begun',
169+
'error' => sprintf( 'Job %d may have begun operation effects and cannot be safely retried.', $job_id ),
170+
);
171+
}
172+
173+
$requeue = $this->db_jobs->commit_missing_direct_operation_requeue(
174+
$job_id,
175+
(int) $missing_action['action_id'],
176+
(int) $missing_action['generation'],
177+
$token,
178+
'manual_retry',
179+
static fn( int $new_generation, string $new_token ): int => (int) as_schedule_single_action(
180+
time(),
181+
DirectJobEnqueuer::HOOK,
182+
array(
183+
'job_id' => $job_id,
184+
'flow_step_id' => $flow_step_id,
185+
'operation_generation' => $new_generation,
186+
'operation_claim_token' => $new_token,
187+
),
188+
DirectJobEnqueuer::GROUP,
189+
true
190+
)
191+
);
192+
if ( empty( $requeue['success'] ) ) {
193+
return array(
194+
'success' => false,
195+
'job_id' => $job_id,
196+
'previous_status' => $previous_status,
197+
'retryable' => true,
198+
'error_code' => (string) ( $requeue['reason'] ?? 'retry_enqueue_failed' ),
199+
'error' => sprintf( 'Job %d retry could not establish durable scheduler ownership.', $job_id ),
200+
);
201+
}
202+
203+
\DataMachine\Core\RunMetrics::increment( $job_id, 'retried' );
204+
return array(
205+
'success' => true,
206+
'job_id' => $job_id,
207+
'previous_status' => $previous_status,
208+
'prompt_requeued' => false,
209+
'direct_requeued' => true,
210+
'message' => sprintf( 'Job %d direct workflow retry enqueued.', $job_id ),
211+
);
212+
}
152213
$this->db_jobs->complete_job( $job_id, 'failed - manual_retry' );
153214
}
154215
if ( ! $this->db_jobs->reopen_failed_job( $job_id ) ) {

inc/Core/Database/Jobs/Jobs.php

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
use DataMachine\Core\Database\RunMetadata\RunMetadata;
2222
use DataMachine\Core\ExecutionQuery;
2323
use DataMachine\Core\ChildJobRecoveryPolicy;
24+
use DataMachine\Core\DirectOperationRecoveryPolicy;
2425
use DataMachine\Core\JobStatus;
2526
use DataMachine\Core\RunMetrics;
2627
use DataMachine\Core\RunLifecycleStore;
@@ -561,6 +562,11 @@ public function commit_missing_direct_operation_requeue( int $job_id, int $actio
561562
$result['reason'] = 'schedule_failed';
562563
return $result;
563564
}
565+
if ( ! DirectOperationRecoveryPolicy::recordedActionExists( $new_action_id ) ) {
566+
$this->wpdb->query( 'ROLLBACK' ); // phpcs:ignore WordPress.DB.DirectDatabaseQuery.DirectQuery,WordPress.DB.DirectDatabaseQuery.NoCaching
567+
$result['reason'] = 'action_receipt_missing';
568+
return $result;
569+
}
564570
$engine = is_array( $job['engine_data'] ?? null ) ? $job['engine_data'] : array();
565571
$engine['direct_operation_recovery'] = array(
566572
'schema' => 'datamachine.direct-operation-recovery.v1',

inc/Core/DirectJobEnqueuer.php

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,15 +19,19 @@ class DirectJobEnqueuer {
1919
private Jobs $jobs;
2020
private \Closure $scheduler;
2121
private \Closure $action_finder;
22+
private \Closure $action_receipt_exists;
2223

23-
public function __construct( ?Jobs $jobs = null, ?callable $scheduler = null, ?callable $action_finder = null ) {
24-
$this->jobs = $jobs ?? new Jobs();
25-
$this->scheduler = null !== $scheduler
24+
public function __construct( ?Jobs $jobs = null, ?callable $scheduler = null, ?callable $action_finder = null, ?callable $action_receipt_exists = null ) {
25+
$this->jobs = $jobs ?? new Jobs();
26+
$this->scheduler = null !== $scheduler
2627
? \Closure::fromCallable( $scheduler )
2728
: static fn( int $run_at, string $hook, array $args, string $group ) => as_schedule_single_action( $run_at, $hook, $args, $group );
28-
$this->action_finder = null !== $action_finder
29+
$this->action_finder = null !== $action_finder
2930
? \Closure::fromCallable( $action_finder )
3031
: fn( array $args ): int => $this->findScheduledActionId( $args );
32+
$this->action_receipt_exists = null !== $action_receipt_exists
33+
? \Closure::fromCallable( $action_receipt_exists )
34+
: static fn( int $action_id ): bool => DirectOperationRecoveryPolicy::recordedActionExists( $action_id );
3135
}
3236

3337
/**
@@ -93,6 +97,10 @@ public function enqueue( int $job_id, string $flow_step_id, ?int $timestamp = nu
9397
$this->jobs->finish_operation_enqueue( $job_id, 'enqueue_failed', 0, $token, $generation );
9498
return $this->failure( 'action_schedule_failed', $generation, true );
9599
}
100+
if ( ! ( $this->action_receipt_exists )( $action_id ) ) {
101+
$this->jobs->finish_operation_enqueue( $job_id, 'enqueue_failed', 0, $token, $generation );
102+
return $this->failure( 'action_receipt_missing', $generation, true );
103+
}
96104

97105
if ( ! $this->jobs->finish_operation_enqueue( $job_id, 'enqueued', $action_id, $token, $generation ) ) {
98106
return $this->failure( 'enqueue_claim_fenced', $generation, true, 'enqueuing' );

tests/Unit/Abilities/Job/DirectJobOwnershipTest.php

Lines changed: 110 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
use DataMachine\Core\Database\Agents\Agents;
2222
use DataMachine\Core\Database\Jobs\Jobs;
2323
use DataMachine\Core\DirectJobEnqueuer;
24+
use DataMachine\Core\DirectOperationRecoveryPolicy;
2425
use DataMachine\Core\JobRetryPolicy;
2526
use DataMachine\Core\JobStatus;
2627
use WP_UnitTestCase;
@@ -291,7 +292,7 @@ public function test_enqueue_failure_is_reclaimable_and_concurrent_calls_schedul
291292
)
292293
);
293294

294-
$failing = new DirectJobEnqueuer( $jobs, static fn() => false, static fn() => 0 );
295+
$failing = new DirectJobEnqueuer( $jobs, static fn() => false, static fn() => 0, static fn() => false );
295296
$this->assertFalse( $failing->enqueue( $job_id, 'ephemeral_step_0' )['success'] );
296297
$this->assertSame( 'enqueue_failed', $jobs->get_job( $job_id )['operation_state'] );
297298

@@ -306,7 +307,8 @@ static function () use ( &$schedule_count, &$scheduled_id ) {
306307
},
307308
static function () use ( &$scheduled_id ) {
308309
return $scheduled_id;
309-
}
310+
},
311+
static fn() => true
310312
);
311313
$first = $enqueuer->enqueue( $job_id, 'ephemeral_step_0' );
312314
$second = $enqueuer->enqueue( $job_id, 'ephemeral_step_0' );
@@ -329,11 +331,31 @@ static function () use ( &$scheduled_id ) {
329331
array( 'operation_claimed_at' => '2000-01-01 00:00:00' ),
330332
array( 'job_id' => $crashed_job_id )
331333
);
332-
$recovered = ( new DirectJobEnqueuer( $jobs, static fn() => 92, static fn() => 0 ) )->enqueue( $crashed_job_id, 'ephemeral_step_0' );
334+
$recovered = ( new DirectJobEnqueuer( $jobs, static fn() => 92, static fn() => 0, static fn() => true ) )->enqueue( $crashed_job_id, 'ephemeral_step_0' );
333335
$this->assertTrue( $recovered['success'] );
334336
$this->assertSame( 'enqueued', $jobs->get_job( $crashed_job_id )['operation_state'] );
335337
}
336338

339+
public function test_positive_scheduler_id_without_receipt_is_reclaimable(): void {
340+
$jobs = new Jobs();
341+
$job_id = $jobs->create_job(
342+
array(
343+
'pipeline_id' => 'direct',
344+
'flow_id' => 'direct',
345+
'operation_state' => 'preparing',
346+
'operation_step_id' => 'ephemeral_step_0',
347+
)
348+
);
349+
350+
$result = ( new DirectJobEnqueuer( $jobs, static fn() => 999999, static fn() => 0, static fn() => false ) )->enqueue( $job_id, 'ephemeral_step_0' );
351+
$job = $jobs->get_job( $job_id );
352+
353+
$this->assertFalse( $result['success'] );
354+
$this->assertSame( 'action_receipt_missing', $result['error'] );
355+
$this->assertSame( 'enqueue_failed', $job['operation_state'] );
356+
$this->assertSame( 0, (int) $job['operation_action_id'] );
357+
}
358+
337359
public function test_non_owner_gets_retryable_in_progress_until_action_is_durable(): void {
338360
$jobs = new Jobs();
339361
$job_id = $jobs->create_job(
@@ -346,7 +368,7 @@ public function test_non_owner_gets_retryable_in_progress_until_action_is_durabl
346368
);
347369
$this->assertIsArray( $jobs->claim_operation_enqueue( $job_id ) );
348370

349-
$result = ( new DirectJobEnqueuer( $jobs, static fn() => 99, static fn() => 0 ) )->enqueue( $job_id, 'ephemeral_step_0' );
371+
$result = ( new DirectJobEnqueuer( $jobs, static fn() => 99, static fn() => 0, static fn() => true ) )->enqueue( $job_id, 'ephemeral_step_0' );
350372

351373
$this->assertFalse( $result['success'] );
352374
$this->assertTrue( $result['retryable'] );
@@ -367,7 +389,7 @@ public function test_replay_commits_scheduled_action_after_submitter_crash(): vo
367389
$claim = $jobs->claim_operation_enqueue( $job_id );
368390
$this->assertIsArray( $claim );
369391

370-
$result = ( new DirectJobEnqueuer( $jobs, static fn() => 999, static fn() => 404 ) )->enqueue( $job_id, 'ephemeral_step_0' );
392+
$result = ( new DirectJobEnqueuer( $jobs, static fn() => 999, static fn() => 404, static fn() => true ) )->enqueue( $job_id, 'ephemeral_step_0' );
371393
$job = $jobs->get_job( $job_id );
372394

373395
$this->assertTrue( $result['success'] );
@@ -402,7 +424,8 @@ static function () use ( &$takeover_claim, $jobs, $wpdb, $job_id ) {
402424
$takeover_claim = $jobs->claim_operation_enqueue( $job_id );
403425
return 101;
404426
},
405-
static fn() => 0
427+
static fn() => 0,
428+
static fn() => true
406429
);
407430

408431
$slow_result = $slow->enqueue( $job_id, 'ephemeral_step_0' );
@@ -445,7 +468,8 @@ static function ( int $run_at, string $hook, array $args ) use ( &$worker_result
445468
$worker_result = ( new ExecuteStepAbility() )->execute( $args );
446469
return 303;
447470
},
448-
static fn() => 0
471+
static fn() => 0,
472+
static fn() => true
449473
) )->enqueue( $job_id, 'ephemeral_step_0' );
450474

451475
$this->assertTrue( $worker_result['deferred'] );
@@ -516,6 +540,85 @@ public function test_processing_direct_workflow_retry_waits_for_live_generation(
516540
$this->assertSame( $original_generation, (int) $job['operation_generation'] );
517541
}
518542

543+
public function test_processing_direct_retry_atomically_requeues_missing_action(): void {
544+
global $wpdb;
545+
546+
wp_set_current_user( $this->owner_id );
547+
$created = $this->execute( 'manual-missing-action-retry' );
548+
$jobs = new Jobs();
549+
$job_id = (int) $created['job_id'];
550+
$before = $jobs->get_job( $job_id );
551+
$this->assertTrue( $jobs->start_job( $job_id ) );
552+
$this->assertSame( 1, $wpdb->delete( $wpdb->prefix . 'actionscheduler_actions', array( 'action_id' => (int) $before['operation_action_id'] ), array( '%d' ) ) );
553+
$job_count = $jobs->get_jobs_count();
554+
555+
$retry = ( new RetryJobAbility() )->execute( array( 'job_id' => $job_id ) );
556+
$after = $jobs->get_job( $job_id );
557+
558+
$this->assertTrue( $retry['success'] );
559+
$this->assertTrue( $retry['direct_requeued'] );
560+
$this->assertSame( $job_id, (int) $retry['job_id'] );
561+
$this->assertSame( $job_count, $jobs->get_jobs_count() );
562+
$this->assertSame( JobStatus::PENDING, $after['status'] );
563+
$this->assertSame( 'enqueued', $after['operation_state'] );
564+
$this->assertGreaterThan( (int) $before['operation_generation'], (int) $after['operation_generation'] );
565+
$this->assertTrue( DirectOperationRecoveryPolicy::recordedActionExists( (int) $after['operation_action_id'] ) );
566+
$this->assertSame( 'manual_retry', $after['engine_data']['direct_operation_recovery']['trigger'] );
567+
$this->assertNull( $after['terminal_accounting_state'] );
568+
569+
$duplicate = ( new RetryJobAbility() )->execute( array( 'job_id' => $job_id, 'force' => true ) );
570+
$this->assertFalse( $duplicate['success'] );
571+
$this->assertSame( $job_count, $jobs->get_jobs_count() );
572+
}
573+
574+
public function test_missing_direct_action_requeue_rolls_back_unusable_receipt(): void {
575+
global $wpdb;
576+
577+
wp_set_current_user( $this->owner_id );
578+
$created = $this->execute( 'manual-missing-action-rollback' );
579+
$jobs = new Jobs();
580+
$job_id = (int) $created['job_id'];
581+
$before = $jobs->get_job( $job_id );
582+
$this->assertTrue( $jobs->start_job( $job_id ) );
583+
$this->assertSame( 1, $wpdb->delete( $wpdb->prefix . 'actionscheduler_actions', array( 'action_id' => (int) $before['operation_action_id'] ), array( '%d' ) ) );
584+
585+
$result = $jobs->commit_missing_direct_operation_requeue(
586+
$job_id,
587+
(int) $before['operation_action_id'],
588+
(int) $before['operation_generation'],
589+
(string) $before['operation_claim_token'],
590+
'manual_retry',
591+
static fn(): int => 999999
592+
);
593+
$after = $jobs->get_job( $job_id );
594+
595+
$this->assertFalse( $result['success'] );
596+
$this->assertSame( 'action_receipt_missing', $result['reason'] );
597+
$this->assertSame( JobStatus::PROCESSING, $after['status'] );
598+
$this->assertSame( (int) $before['operation_generation'], (int) $after['operation_generation'] );
599+
$this->assertSame( (int) $before['operation_action_id'], (int) $after['operation_action_id'] );
600+
$this->assertArrayNotHasKey( 'direct_operation_recovery', $after['engine_data'] );
601+
}
602+
603+
public function test_processing_direct_retry_rejects_missing_action_after_effects_begin(): void {
604+
global $wpdb;
605+
606+
wp_set_current_user( $this->owner_id );
607+
$created = $this->execute( 'manual-missing-action-effects' );
608+
$jobs = new Jobs();
609+
$job_id = (int) $created['job_id'];
610+
$before = $jobs->get_job( $job_id );
611+
$this->assertTrue( $jobs->start_job( $job_id ) );
612+
$this->assertTrue( $jobs->mark_operation_effects_begun( $job_id, (int) $before['operation_generation'], (string) $before['operation_claim_token'] ) );
613+
$this->assertSame( 1, $wpdb->delete( $wpdb->prefix . 'actionscheduler_actions', array( 'action_id' => (int) $before['operation_action_id'] ), array( '%d' ) ) );
614+
615+
$retry = ( new RetryJobAbility() )->execute( array( 'job_id' => $job_id ) );
616+
617+
$this->assertFalse( $retry['success'] );
618+
$this->assertSame( 'job_effects_begun', $retry['error_code'] );
619+
$this->assertSame( JobStatus::PROCESSING, $jobs->get_job( $job_id )['status'] );
620+
}
621+
519622
public function test_processing_multistep_retry_detects_live_action_for_different_step(): void {
520623
wp_set_current_user( $this->owner_id );
521624
$created = $this->execute( 'processing-multistep-retry' );

tests/direct-job-generation-smoke.php

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,15 +88,15 @@ function direct_generation_assert( bool $condition, string $message ): void {
8888
$blocked->job['operation_generation'] = 1;
8989
$blocked->job['operation_claim_token'] = 'owner-token';
9090
$blocked->claim_blocked = true;
91-
$blocked_result = ( new DirectJobEnqueuer( $blocked, static fn() => 100, static fn() => 0 ) )->enqueue( 42, 'ephemeral_step_0' );
91+
$blocked_result = ( new DirectJobEnqueuer( $blocked, static fn() => 100, static fn() => 0, static fn() => true ) )->enqueue( 42, 'ephemeral_step_0' );
9292
direct_generation_assert( false === $blocked_result['success'], 'non-owner does not acknowledge success without durable action' );
9393
direct_generation_assert( true === $blocked_result['retryable'] && 'enqueue_in_progress' === $blocked_result['error'], 'non-owner receives explicit retryable in-progress result' );
9494

9595
$crash_recovery = new DirectJobGenerationFakeJobs();
9696
$crash_recovery->job['operation_state'] = 'enqueuing';
9797
$crash_recovery->job['operation_generation'] = 1;
9898
$crash_recovery->job['operation_claim_token'] = 'token-1';
99-
$recovered = ( new DirectJobEnqueuer( $crash_recovery, static fn() => 999, static fn() => 404 ) )->enqueue( 42, 'ephemeral_step_0' );
99+
$recovered = ( new DirectJobEnqueuer( $crash_recovery, static fn() => 999, static fn() => 404, static fn() => true ) )->enqueue( 42, 'ephemeral_step_0' );
100100
direct_generation_assert( true === $recovered['success'], 'replay CAS-commits an action left by a crashed submitter' );
101101
direct_generation_assert( 'enqueued' === $crash_recovery->job['operation_state'] && 404 === $crash_recovery->job['operation_action_id'], 'crash reconciliation durably records the recovered action' );
102102

@@ -108,7 +108,8 @@ static function () use ( $interleaved, &$takeover ) {
108108
$takeover = $interleaved->forceTakeover();
109109
return 101;
110110
},
111-
static fn() => 0
111+
static fn() => 0,
112+
static fn() => true
112113
) )->enqueue( 42, 'ephemeral_step_0' );
113114
direct_generation_assert( false === $slow_result['success'] && 'enqueue_claim_fenced' === $slow_result['error'], 'expired slow generation cannot finish after takeover' );
114115
direct_generation_assert( 2 === $takeover['generation'], 'takeover advances enqueue generation' );
@@ -130,11 +131,17 @@ static function ( int $run_at, string $hook, array $args ) use ( &$seen_args ) {
130131
},
131132
static function ( array $args ): int {
132133
return 1 === (int) ( $args['operation_generation'] ?? 0 ) ? 111 : 0;
133-
}
134+
},
135+
static fn() => true
134136
) )->enqueue( 42, 'ephemeral_step_0' );
135137
direct_generation_assert( true === $retry_result['success'], 'retry generation schedules successfully while prior action exists' );
136138
direct_generation_assert( 2 === $seen_args['operation_generation'], 'retry action is keyed to the next generation' );
137139

140+
$missing_receipt = new DirectJobGenerationFakeJobs();
141+
$missing_result = ( new DirectJobEnqueuer( $missing_receipt, static fn() => 505, static fn() => 0, static fn() => false ) )->enqueue( 42, 'ephemeral_step_0' );
142+
direct_generation_assert( false === $missing_result['success'] && 'action_receipt_missing' === $missing_result['error'], 'positive scheduler ID without a durable receipt fails enqueue' );
143+
direct_generation_assert( 'enqueue_failed' === $missing_receipt->job['operation_state'], 'missing action receipt leaves the operation reclaimable' );
144+
138145
$jobs_source = file_get_contents( dirname( __DIR__ ) . '/inc/Core/Database/Jobs/Jobs.php' ) ?: '';
139146
$step_source = file_get_contents( dirname( __DIR__ ) . '/inc/Abilities/Engine/ExecuteStepAbility.php' ) ?: '';
140147
direct_generation_assert( str_contains( $jobs_source, 'operation_claim_token = %s' ) && str_contains( $jobs_source, 'operation_generation = %d' ), 'enqueue finish is fenced by token and generation' );

0 commit comments

Comments
 (0)