Move job queue cleanup to separate function (#12)

* Move job queue cleanup to separate function

* Add comments, more test cleanup

* Re-instate onInit so migrations are completed before processing the job queue

Author: adam-fowler

Sources/JobsPostgres/Migrations/UpdateJobDelay.swift

@@ -1,10 +1,3 @@
-//
-//  UpdateJobDelay.swift
-//  swift-jobs-postgres
-//
-//  Created by Stevenson Michel on 2/17/25.
-//
-
 //===----------------------------------------------------------------------===//
 //
 // This source file is part of the Hummingbird server framework project

Sources/JobsPostgres/PostgresJobsQueue.swift

@@ -46,7 +46,7 @@ public final class PostgresJobQueue: JobQueueDriver {
     public typealias JobID = UUID
 
     /// what to do with failed/processing jobs from last time queue was handled
-    public enum JobInitialization: Sendable {
+    public enum JobCleanup: Sendable {
         case doNothing
         case rerun
         case remove
@@ -73,20 +73,14 @@ public final class PostgresJobQueue: JobQueueDriver {
 
     /// Queue configuration
     public struct Configuration: Sendable {
-        let pendingJobsInitialization: JobInitialization
-        let failedJobsInitialization: JobInitialization
-        let processingJobsInitialization: JobInitialization
+        /// Queue poll time to wait if queue empties
         let pollTime: Duration
 
+        ///  Initialize configuration
+        /// - Parameter pollTime: Queue poll time to wait if queue empties
         public init(
-            pendingJobsInitialization: JobInitialization = .doNothing,
-            failedJobsInitialization: JobInitialization = .rerun,
-            processingJobsInitialization: JobInitialization = .rerun,
             pollTime: Duration = .milliseconds(100)
         ) {
-            self.pendingJobsInitialization = pendingJobsInitialization
-            self.failedJobsInitialization = failedJobsInitialization
-            self.processingJobsInitialization = processingJobsInitialization
             self.pollTime = pollTime
         }
     }
@@ -115,20 +109,47 @@ public final class PostgresJobQueue: JobQueueDriver {
         await migrations.add(UpdateJobDelay())
     }
 
-    /// Run on initialization of the job queue
     public func onInit() async throws {
+        self.logger.info("Waiting for JobQueue migrations")
+        /// Need migrations to have completed before job queue processing can start
+        try await self.migrations.waitUntilCompleted()
+    }
+
+    ///  Cleanup job queues
+    ///
+    /// This function is used to re-run or delete jobs in a certain state. Failed jobs can be
+    /// pushed back into the pending queue to be re-run or removed. When called at startup in
+    /// theory no job should be set to processing, or set to pending but not in the queue. but if
+    /// your job server crashes these states are possible, so we also provide options to re-queue
+    /// these jobs so they are run again.
+    ///
+    /// The job queue needs to be running when you call cleanup. You can call `cleanup` with
+    /// `failedJobs`` set to whatever you like at any point to re-queue failed jobs. Moving processing
+    /// or pending jobs should only be done if you are certain there is nothing else processing
+    /// the job queue.
+    ///
+    /// - Parameters:
+    ///   - failedJobs: What to do with jobs tagged as failed
+    ///   - processingJobs: What to do with jobs tagged as processing
+    ///   - pendingJobs: What to do with jobs tagged as pending
+    /// - Throws:
+    public func cleanup(
+        failedJobs: JobCleanup = .doNothing,
+        processingJobs: JobCleanup = .doNothing,
+        pendingJobs: JobCleanup = .doNothing
+    ) async throws {
         do {
-            self.logger.info("Waiting for JobQueue migrations")
+            /// wait for migrations to complete before running job queue cleanup
             try await self.migrations.waitUntilCompleted()
             _ = try await self.client.withConnection { connection in
-                self.logger.info("Update Jobs at initialization")
-                try await self.updateJobsOnInit(withStatus: .pending, onInit: self.configuration.pendingJobsInitialization, connection: connection)
+                self.logger.info("Update Jobs")
+                try await self.updateJobsOnInit(withStatus: .pending, onInit: pendingJobs, connection: connection)
                 try await self.updateJobsOnInit(
                     withStatus: .processing,
-                    onInit: self.configuration.processingJobsInitialization,
+                    onInit: processingJobs,
                     connection: connection
                 )
-                try await self.updateJobsOnInit(withStatus: .failed, onInit: self.configuration.failedJobsInitialization, connection: connection)
+                try await self.updateJobsOnInit(withStatus: .failed, onInit: failedJobs, connection: connection)
             }
         } catch let error as PSQLError {
             logger.error(
@@ -321,7 +342,7 @@ public final class PostgresJobQueue: JobQueueDriver {
         return jobs
     }
 
-    func updateJobsOnInit(withStatus status: Status, onInit: JobInitialization, connection: PostgresConnection) async throws {
+    func updateJobsOnInit(withStatus status: Status, onInit: JobCleanup, connection: PostgresConnection) async throws {
         switch onInit {
         case .remove:
             try await connection.query(

Tests/JobsPostgresTests/JobsTests.swift

@@ -51,7 +51,7 @@ final class JobsTests: XCTestCase {
 
     func createJobQueue(
         numWorkers: Int,
-        configuration: PostgresJobQueue.Configuration,
+        configuration: PostgresJobQueue.Configuration = .init(),
         function: String = #function
     ) async throws -> JobQueue<PostgresJobQueue> {
         let logger = {
@@ -87,6 +87,9 @@ final class JobsTests: XCTestCase {
     /// shutdown correctly
     @discardableResult public func testJobQueue<T>(
         jobQueue: JobQueue<PostgresJobQueue>,
+        failedJobsInitialization: PostgresJobQueue.JobCleanup = .remove,
+        processingJobsInitialization: PostgresJobQueue.JobCleanup = .remove,
+        pendingJobsInitialization: PostgresJobQueue.JobCleanup = .remove,
         revertMigrations: Bool = false,
         test: (JobQueue<PostgresJobQueue>) async throws -> T
     ) async throws -> T {
@@ -110,6 +113,7 @@ final class JobsTests: XCTestCase {
                         try await migrations.revert(client: client, groups: [.jobQueue], logger: logger, dryRun: false)
                     }
                     try await migrations.apply(client: client, groups: [.jobQueue], logger: logger, dryRun: false)
+                    try await jobQueue.queue.cleanup(failedJobs: failedJobsInitialization, processingJobs: processingJobsInitialization)
                     let value = try await test(jobQueue)
                     await serviceGroup.triggerGracefulShutdown()
                     return value
@@ -134,13 +138,22 @@ final class JobsTests: XCTestCase {
     /// shutdown correctly
     @discardableResult public func testJobQueue<T>(
         numWorkers: Int,
-        configuration: PostgresJobQueue.Configuration = .init(failedJobsInitialization: .remove, processingJobsInitialization: .remove),
+        failedJobsInitialization: PostgresJobQueue.JobCleanup = .remove,
+        processingJobsInitialization: PostgresJobQueue.JobCleanup = .remove,
+        pendingJobsInitialization: PostgresJobQueue.JobCleanup = .remove,
         revertMigrations: Bool = true,
         function: String = #function,
         test: (JobQueue<PostgresJobQueue>) async throws -> T
     ) async throws -> T {
-        let jobQueue = try await self.createJobQueue(numWorkers: numWorkers, configuration: configuration, function: function)
-        return try await self.testJobQueue(jobQueue: jobQueue, revertMigrations: revertMigrations, test: test)
+        let jobQueue = try await self.createJobQueue(numWorkers: numWorkers, configuration: .init(), function: function)
+        return try await self.testJobQueue(
+            jobQueue: jobQueue,
+            failedJobsInitialization: failedJobsInitialization,
+            processingJobsInitialization: processingJobsInitialization,
+            pendingJobsInitialization: pendingJobsInitialization,
+            revertMigrations: revertMigrations,
+            test: test
+        )
     }
 
     func testBasic() async throws {
@@ -368,14 +381,13 @@ final class JobsTests: XCTestCase {
             finished.store(true, ordering: .relaxed)
         }
         let jobQueue = try await createJobQueue(
-            numWorkers: 1,
-            configuration: .init(pendingJobsInitialization: .remove, failedJobsInitialization: .rerun)
+            numWorkers: 1
         )
         jobQueue.registerJob(job)
-        try await self.testJobQueue(jobQueue: jobQueue, revertMigrations: true) { jobQueue in
-            // stall to give onInit a chance to run, so it can remove any pendng jobs
-            try await Task.sleep(for: .milliseconds(100))
-
+        try await self.testJobQueue(
+            jobQueue: jobQueue,
+            revertMigrations: true
+        ) { jobQueue in
             try await jobQueue.push(id: jobIdentifer, parameters: 0)
 
             await self.wait(for: [failedExpectation], timeout: 10)
@@ -384,9 +396,9 @@ final class JobsTests: XCTestCase {
             XCTAssertFalse(finished.load(ordering: .relaxed))
         }
 
-        let jobQueue2 = try await createJobQueue(numWorkers: 1, configuration: .init(failedJobsInitialization: .rerun))
+        let jobQueue2 = try await createJobQueue(numWorkers: 1)
         jobQueue2.registerJob(job)
-        try await self.testJobQueue(jobQueue: jobQueue2) { _ in
+        try await self.testJobQueue(jobQueue: jobQueue2, failedJobsInitialization: .rerun) { _ in
             await self.wait(for: [succeededExpectation], timeout: 10)
             XCTAssertTrue(finished.load(ordering: .relaxed))
         }
@@ -414,7 +426,6 @@ final class JobsTests: XCTestCase {
             .postgres(
                 client: postgresClient,
                 migrations: postgresMigrations,
-                configuration: .init(failedJobsInitialization: .remove, processingJobsInitialization: .remove),
                 logger: logger
             ),
             numWorkers: 2,
@@ -425,7 +436,6 @@ final class JobsTests: XCTestCase {
             .postgres(
                 client: postgresClient,
                 migrations: postgresMigrations2,
-                configuration: .init(failedJobsInitialization: .remove, processingJobsInitialization: .remove),
                 logger: logger
             ),
             numWorkers: 2,
@@ -447,6 +457,8 @@ final class JobsTests: XCTestCase {
             }
             try await postgresMigrations.apply(client: postgresClient, groups: [.jobQueue], logger: logger, dryRun: false)
             try await postgresMigrations2.apply(client: postgresClient, groups: [.jobQueue], logger: logger, dryRun: false)
+            try await jobQueue.queue.cleanup(failedJobs: .remove, processingJobs: .remove)
+            try await jobQueue2.queue.cleanup(failedJobs: .remove, processingJobs: .remove)
             do {
                 for i in 0..<200 {
                     try await jobQueue.push(id: jobIdentifer, parameters: i)
@@ -475,7 +487,6 @@ final class JobsTests: XCTestCase {
             let jobQueue = await PostgresJobQueue(
                 client: postgresClient,
                 migrations: postgresMigrations,
-                configuration: .init(failedJobsInitialization: .remove, processingJobsInitialization: .remove),
                 logger: logger
             )
             try await postgresMigrations.apply(client: postgresClient, groups: [.jobQueue], logger: logger, dryRun: false)