Adding Job Actions which allows for cancelling, pausing and resuming 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: "job-actions"),
+        .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

@@ -131,6 +131,8 @@ public final class PostgresJobQueue: JobQueueDriver {
         case pending = 0
         case processing = 1
         case failed = 2
+        case paused = 3
+        case cancelled = 4
     }
 
     /// Queue configuration
@@ -247,22 +249,51 @@ public final class PostgresJobQueue: JobQueueDriver {
 
     /// Retry an existing Job
     /// - Parameters
-    ///   - id: Job instance ID
+    ///   - jobID: Job instance ID
     ///   - jobRequest: Job Request
     ///   - options: Job retry options
-    public func retry<Parameters: JobParameters>(_ id: JobID, jobRequest: JobRequest<Parameters>, options: JobRetryOptions) async throws {
+    public func retry<Parameters: JobParameters>(_ jobID: JobID, jobRequest: JobRequest<Parameters>, options: JobRetryOptions) async throws {
         let buffer = try self.jobRegistry.encode(jobRequest: jobRequest)
         try await self.client.withTransaction(logger: self.logger) { connection in
-            try await self.updateJob(id: id, buffer: buffer, connection: connection)
+            try await self.updateJob(jobID: jobID, buffer: buffer, connection: connection)
             try await self.addToQueue(
-                jobID: id,
+                jobID: jobID,
                 queueName: configuration.queueName,
                 options: .init(delayUntil: options.delayUntil),
                 connection: connection
             )
         }
     }
 
+    /// Perform actions on job
+    /// - Parameters
+    ///   - jobID: Job instance ID
+    ///   - action Job Action
+    public func performAction(jobID: JobID, action: JobAction) async throws {
+        switch action {
+            case .cancel():
+                try await self.performJobAction(jobID: jobID, status: .cancelled)
+            case .pause():
+                try await self.performJobAction(jobID: jobID, status: .paused)
+            case .resume():
+                try await self.client.withTransaction(logger: logger) { connection in
+                    try await self.setStatus(jobID: jobID, status: .pending, connection: connection)
+                    try await self.addToQueue(
+                        jobID: jobID,
+                        queueName: configuration.queueName,
+                        options: .init(),
+                        connection: connection
+                    )
+                }
+            default:
+                break
+        }
+    }
+    
+    public func isEmpty() async throws -> Bool {
+        true
+    }
+
     /// This is called to say job has finished processing and it can be deleted
     public func finished(jobID: JobID) async throws {
         try await self.delete(jobID: jobID)
@@ -408,15 +439,15 @@ public final class PostgresJobQueue: JobQueueDriver {
             logger: self.logger
         )
     }
-    // TODO: maybe add a new column colum for attempt so far after PR https://github.com/hummingbird-project/swift-jobs/pull/63 is merged?
-    func updateJob(id: JobID, buffer: ByteBuffer, connection: PostgresConnection) async throws {
+
+    func updateJob(jobID: JobID, buffer: ByteBuffer, connection: PostgresConnection) async throws {
         try await connection.query(
             """
             UPDATE swift_jobs.jobs
             SET job = \(buffer),
                 last_modified = \(Date.now),
                 status = \(Status.failed)
-            WHERE id = \(id) AND queue_name = \(configuration.queueName)
+            WHERE id = \(jobID) AND queue_name = \(configuration.queueName)
             """,
             logger: self.logger
         )
@@ -485,6 +516,20 @@ public final class PostgresJobQueue: JobQueueDriver {
         return jobs
     }
 
+    func performJobAction(jobID: JobID, status: Status) async throws {
+        try await self.client.withTransaction(logger: logger) { connection in
+            
+            try await connection.query(
+                """
+                DELETE FROM swift_jobs.queues
+                WHERE job_id = \(jobID) AND queue_name = \(configuration.queueName)
+                """,
+                logger: self.logger
+            )
+            try await self.setStatus(jobID: jobID, status: status, connection: connection)
+        }
+    }
+
     func updateJobsOnInit(withStatus status: Status, onInit: JobCleanup, connection: PostgresConnection) async throws {
         switch onInit {
         case .remove:

Tests/JobsPostgresTests/JobsTests.swift

@@ -683,4 +683,72 @@ 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.performAction(jobID: firstJob, action: .pause())
+
+            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.performAction(jobID: cancelID, action: .cancel())
+
+            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.performAction(jobID: firstJob, action: .resume())
+                try await jobQueue.performAction(jobID: cancelID, action: .resume())
+                await fulfillment(of: [expectation], timeout: 10)
+                await serviceGroup.triggerGracefulShutdown()
+            }
+        }
+        XCTAssertEqual(jobExecutionSequence.withLockedValue { $0 }, [2025, 20, 42])
+    }
 }