Implementing cancellable and resumable jobs

Author: thoven87

Package.swift

@@ -10,8 +10,8 @@ let package = Package(
     ],
     dependencies: [
         // TODO: use a released version of swift-jobs
-        .package(url: "https://github.com/hummingbird-project/swift-jobs.git", branch: "main"),
-        .package(url: "https://github.com/hummingbird-project/hummingbird-postgres.git", from: "0.5.0"),
+        .package(url: "https://github.com/thoven87/swift-jobs.git", branch: "resumable-cancellable-jobs"),
+        .package(url: "https://github.com/hummingbird-project/hummingbird-postgres.git", branch: "main"),
         .package(url: "https://github.com/vapor/postgres-nio.git", from: "1.25.0"),
     ],
     targets: [

Sources/JobsPostgres/PostgresJobsQueue.swift

@@ -42,7 +42,8 @@ import PostgresNIO
 ///     try await migrations.apply(client: postgresClient, logger: logger, dryRun: applyMigrations)
 /// }
 /// ```
-public final class PostgresJobQueue: JobQueueDriver {
+public final class PostgresJobQueue: JobQueueDriver, CancellableJobQueueProtocol, ResumeableJobQueueProtocol {
+    
     public typealias JobID = UUID
     /// what to do with failed/processing jobs from last time queue was handled
     public enum JobCleanup: Sendable {
@@ -131,6 +132,8 @@ public final class PostgresJobQueue: JobQueueDriver {
         case pending = 0
         case processing = 1
         case failed = 2
+        case cancelled = 3
+        case paused = 4
     }
 
     /// Queue configuration
@@ -171,7 +174,7 @@ public final class PostgresJobQueue: JobQueueDriver {
         self.logger = logger
         self.isStopped = .init(false)
         self.migrations = migrations
-        await migrations.add(CreateSwiftJobsMigrations())
+        await migrations.add(CreateSwiftJobsMigrations(), checkForDuplicates: true)
     }
 
     public func onInit() async throws {
@@ -180,6 +183,56 @@ public final class PostgresJobQueue: JobQueueDriver {
         try await self.migrations.waitUntilCompleted()
     }
 
+    ///  Cancel job
+    ///
+    /// This function is used to cancel a job. Job cancellation is not gaurenteed howerever.
+    /// Cancellable jobs are jobs with a delayed greather than when the cancellation request was made
+    ///
+    /// - Parameters:
+    ///   - jobID: an existing job
+    /// - Throws:
+    public func cancel(jobID: JobID) async throws {
+        try await self.client.withTransaction(logger: logger) { connection in
+            try await deleteFromQueue(jobID: jobID, connection: connection)
+            try await setStatus(jobID: jobID, status: .cancelled, connection: connection)
+        }
+    }
+
+    ///  Pause job
+    ///
+    /// This function is used to pause a job. Job paus is not gaurenteed howerever.
+    /// Pausable jobs are jobs with a delayed greather than when the pause request was made
+    ///
+    /// - Parameters:
+    ///   - jobID: an existing job
+    /// - Throws:
+    public func pause(jobID: UUID) async throws {
+        try await self.client.withTransaction(logger: logger) { connection in
+            try await deleteFromQueue(jobID: jobID, connection: connection)
+            try await setStatus(jobID: jobID, status: .paused, connection: connection)
+        }
+    }
+
+    ///  Resume job
+    ///
+    /// This function is used to resume jobs. Job  is not gaurenteed howerever.
+    /// Cancellable jobs are jobs with a delayed greather than when the cancellation request was made
+    ///
+    /// - Parameters:
+    ///   - jobID: an existing job
+    /// - Throws:
+    public func resume(jobID: JobID) async throws {
+        try await self.client.withTransaction(logger: logger) { connection in
+            try await setStatus(jobID: jobID, status: .pending, connection: connection)
+            try await addToQueue(
+                jobID: jobID,
+                queueName: configuration.queueName,
+                options: .init(),
+                connection: connection
+            )
+        }
+    }
+
     ///  Cleanup job queues
     ///
     /// This function is used to re-run or delete jobs in a certain state. Failed jobs can be
@@ -431,6 +484,16 @@ public final class PostgresJobQueue: JobQueueDriver {
             logger: self.logger
         )
     }
+    
+    func deleteFromQueue(jobID: JobID, connection: PostgresConnection) async throws {
+        try await connection.query(
+            """
+            DELETE FROM swift_jobs.queues
+            WHERE job_id = \(jobID) AND queue_name = \(configuration.queueName)
+            """,
+            logger: self.logger
+        )
+    }
 
     func addToQueue(jobID: JobID, queueName: String, options: JobOptions, connection: PostgresConnection) async throws {
         try await connection.query(

Tests/JobsPostgresTests/JobsTests.swift

@@ -683,4 +683,74 @@ final class JobsTests: XCTestCase {
             group.cancelAll()
         }
     }
+    
+    func testJobActions() async throws {
+        struct TestParameters: JobParameters {
+            static let jobName = "testJobActions"
+            let value: Int
+        }
+        let expectation = XCTestExpectation(description: "TestJob.execute was called", expectedFulfillmentCount: 3)
+        let jobExecutionSequence: NIOLockedValueBox<[Int]> = .init([])
+
+        let jobQueue = try await self.createJobQueue(numWorkers: 1, configuration: .init(), function: #function)
+
+        try await testPriorityJobQueue(jobQueue: jobQueue) { queue in
+            queue.registerJob(parameters: TestParameters.self) { parameters, context in
+                context.logger.info("Parameters=\(parameters.value)")
+                jobExecutionSequence.withLockedValue {
+                    $0.append(parameters.value)
+                }
+                try await Task.sleep(for: .milliseconds(Int.random(in: 10..<50)))
+                expectation.fulfill()
+            }
+
+            let firstJob = try await queue.push(
+                TestParameters(value: 20),
+                options: .init(
+                    priority: .lowest()
+                )
+            )
+
+            try await jobQueue.pauseJob(jobID: firstJob)
+
+            try await queue.push(
+                TestParameters(value: 2025),
+                options: .init(
+                    priority: .highest()
+                )
+            )
+
+            let cancelID = try await queue.push(
+                TestParameters(value: 42),
+                options: .init(
+                    priority: .lower()
+                )
+            )
+
+            try await jobQueue.cancelJob(jobID: cancelID)
+
+            try await withThrowingTaskGroup(of: Void.self) { group in
+                let serviceGroup = ServiceGroup(services: [queue], logger: queue.logger)
+
+                let processingJobs = try await jobQueue.queue.getJobs(withStatus: .pending)
+                XCTAssertEqual(processingJobs.count, 1)
+
+                group.addTask {
+                    try await serviceGroup.run()
+                }
+
+                let pausedJobs = try await jobQueue.queue.getJobs(withStatus: .paused)
+                XCTAssertEqual(pausedJobs.count, 1)
+                let cancelledJobs = try await jobQueue.queue.getJobs(withStatus: .cancelled)
+                XCTAssertEqual(cancelledJobs.count, 1)
+
+                try await jobQueue.resumeJob(jobID: firstJob)
+                try await jobQueue.resumeJob(jobID: cancelID)
+
+                await fulfillment(of: [expectation], timeout: 10)
+                await serviceGroup.triggerGracefulShutdown()
+            }
+        }
+        XCTAssertEqual(jobExecutionSequence.withLockedValue { $0 }, [2025, 20, 42])
+    }
 }