Merge pull request #24 from thoven87/job-priority

Adding Job priority

Author: thoven87

Sources/JobsPostgres/Migrations/CreateSwiftJobsMigrations.swift

@@ -41,16 +41,17 @@ struct CreateSwiftJobsMigrations: DatabaseMigration {
                 job_id uuid PRIMARY KEY,
                 created_at TIMESTAMPTZ NOT NULL,
                 delayed_until TIMESTAMPTZ NOT NULL DEFAULT now(),
-                queue_name TEXT NOT NULL DEFAULT 'default'
+                queue_name TEXT NOT NULL DEFAULT 'default',
+                priority SMALLINT NOT NULL DEFAULT 0
             );
             """,
             logger: logger
         )
 
         try await connection.query(
             """
-            CREATE INDEX CONCURRENTLY IF NOT EXISTS queues_delayed_until_queue_name_idx 
-            ON swift_jobs.queues(delayed_until, queue_name)
+            CREATE INDEX CONCURRENTLY IF NOT EXISTS queues_delayed_until_priority_queue_name_idx 
+            ON swift_jobs.queues(priority DESC, delayed_until ASC, queue_name ASC)
             """,
             logger: logger
         )

Sources/JobsPostgres/PostgresJobsQueue.swift

@@ -51,20 +51,66 @@ public final class PostgresJobQueue: JobQueueDriver {
         case remove
     }
 
+    /// Job priority from lowest to highest
+    public struct JobPriority: Equatable, Sendable {
+        let rawValue: Priority
+
+        // Job priority
+        enum Priority: Int16, Sendable, PostgresCodable {
+            case lowest = 0
+            case lower = 1
+            case normal = 2
+            case higher = 3
+            case highest = 4
+        }
+        /// Lowest priority
+        public static func lowest() -> JobPriority {
+            JobPriority(rawValue: .lowest)
+        }
+        /// Lower priority
+        public static func lower() -> JobPriority {
+            JobPriority(rawValue: .lower)
+        }
+        /// Normal is the default priority
+        public static func normal() -> JobPriority {
+            JobPriority(rawValue: .normal)
+        }
+        /// Higher priority
+        public static func higher() -> JobPriority {
+            JobPriority(rawValue: .higher)
+        }
+        /// Higgest priority
+        public static func highest() -> JobPriority {
+            JobPriority(rawValue: .highest)
+        }
+    }
+
     /// Options for job pushed to queue
     public struct JobOptions: JobOptionsProtocol {
         /// Delay running job until
         public var delayUntil: Date
+        /// Priority for this job
+        public var priority: JobPriority
 
         /// Default initializer for JobOptions
         public init() {
             self.delayUntil = .now
+            self.priority = .normal()
         }
 
         ///  Initializer for JobOptions
         /// - Parameter delayUntil: Whether job execution should be delayed until a later date
         public init(delayUntil: Date?) {
             self.delayUntil = delayUntil ?? .now
+            self.priority = .normal()
+        }
+
+        ///  Initializer for JobOptions
+        /// - Parameter delayUntil: Whether job execution should be delayed until a later date
+        /// - Parameter priority: The priority for a job
+        public init(delayUntil: Date = .now, priority: JobPriority = .normal()) {
+            self.delayUntil = delayUntil
+            self.priority = priority
         }
     }
 
@@ -194,7 +240,7 @@ public final class PostgresJobQueue: JobQueueDriver {
         let jobID = JobID()
         try await self.client.withTransaction(logger: self.logger) { connection in
             try await self.add(jobID: jobID, jobRequest: jobRequest, queueName: configuration.queueName, connection: connection)
-            try await self.addToQueue(jobID: jobID, queueName: configuration.queueName, delayUntil: options.delayUntil, connection: connection)
+            try await self.addToQueue(jobID: jobID, queueName: configuration.queueName, options: options, connection: connection)
         }
         return jobID
     }
@@ -208,7 +254,12 @@ public final class PostgresJobQueue: JobQueueDriver {
         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.addToQueue(jobID: id, queueName: configuration.queueName, delayUntil: options.delayUntil, connection: connection)
+            try await self.addToQueue(
+                jobID: id,
+                queueName: configuration.queueName,
+                options: .init(delayUntil: options.delayUntil),
+                connection: connection
+            )
         }
     }
 
@@ -273,7 +324,7 @@ public final class PostgresJobQueue: JobQueueDriver {
                         FROM swift_jobs.queues
                         WHERE delayed_until <= NOW()
                         AND queue_name = \(configuration.queueName)
-                        ORDER BY created_at, delayed_until ASC
+                        ORDER BY priority DESC, delayed_until ASC, created_at ASC 
                         FOR UPDATE SKIP LOCKED
                         LIMIT 1
                     )
@@ -381,11 +432,11 @@ public final class PostgresJobQueue: JobQueueDriver {
         )
     }
 
-    func addToQueue(jobID: JobID, queueName: String, delayUntil: Date, connection: PostgresConnection) async throws {
+    func addToQueue(jobID: JobID, queueName: String, options: JobOptions, connection: PostgresConnection) async throws {
         try await connection.query(
             """
-            INSERT INTO swift_jobs.queues (job_id, created_at, delayed_until, queue_name)
-            VALUES (\(jobID), \(Date.now), \(delayUntil), \(queueName))
+            INSERT INTO swift_jobs.queues (job_id, created_at, delayed_until, queue_name, priority)
+            VALUES (\(jobID), \(Date.now), \(options.delayUntil), \(queueName), \(options.priority.rawValue))
             -- We have found an existing job with the same id, SKIP this INSERT 
             ON CONFLICT (job_id) DO NOTHING
             """,
@@ -449,7 +500,7 @@ public final class PostgresJobQueue: JobQueueDriver {
             let jobs = try await getJobs(withStatus: status)
             self.logger.info("Moving \(jobs.count) jobs with status: \(status) to job queue")
             for jobID in jobs {
-                try await self.addToQueue(jobID: jobID, queueName: configuration.queueName, delayUntil: Date.now, connection: connection)
+                try await self.addToQueue(jobID: jobID, queueName: configuration.queueName, options: .init(), connection: connection)
             }
 
         case .doNothing:

Tests/JobsPostgresTests/JobsTests.swift

@@ -120,6 +120,54 @@ final class JobsTests: XCTestCase {
         }
     }
 
+    /// Helper for testing job priority
+    @discardableResult public func testPriorityJobQueue<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 {
+        do {
+            return try await withThrowingTaskGroup(of: Void.self) { group in
+                let serviceGroup = ServiceGroup(
+                    configuration: .init(
+                        services: [jobQueue.queue.client],
+                        gracefulShutdownSignals: [.sigterm, .sigint],
+                        logger: jobQueue.queue.logger
+                    )
+                )
+                group.addTask {
+                    try await serviceGroup.run()
+                }
+                do {
+                    let migrations = jobQueue.queue.migrations
+                    let client = jobQueue.queue.client
+                    let logger = jobQueue.queue.logger
+                    if revertMigrations {
+                        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
+                } catch let error as PSQLError {
+                    XCTFail("\(String(reflecting: error))")
+                    await serviceGroup.triggerGracefulShutdown()
+                    throw error
+                } catch {
+                    await serviceGroup.triggerGracefulShutdown()
+                    throw error
+                }
+            }
+        } catch let error as PSQLError {
+            XCTFail("\(String(reflecting: error))")
+            throw error
+        }
+    }
+
     /// Helper function for test a server
     ///
     /// Creates test client, runs test function abd ensures everything is
@@ -207,6 +255,115 @@ final class JobsTests: XCTestCase {
         XCTAssertEqual(jobExecutionSequence.withLockedValue { $0 }, [5, 1])
     }
 
+    func testJobPriorities() async throws {
+        struct TestParameters: JobParameters {
+            static let jobName = "testPriorityJobs"
+            let value: Int
+        }
+        let expectation = XCTestExpectation(description: "TestJob.execute was called", expectedFulfillmentCount: 2)
+        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()
+            }
+
+            try await queue.push(
+                TestParameters(value: 20),
+                options: .init(
+                    priority: .lowest()
+                )
+            )
+
+            try await queue.push(
+                TestParameters(value: 2025),
+                options: .init(
+                    priority: .highest()
+                )
+            )
+
+            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, 2)
+
+                group.addTask {
+                    try await serviceGroup.run()
+                }
+
+                await fulfillment(of: [expectation], timeout: 10)
+
+                let pendingJobs = try await jobQueue.queue.getJobs(withStatus: .pending)
+                XCTAssertEqual(pendingJobs.count, 0)
+                await serviceGroup.triggerGracefulShutdown()
+            }
+        }
+        XCTAssertEqual(jobExecutionSequence.withLockedValue { $0 }, [2025, 20])
+    }
+
+    func testJobPrioritiesWithDelay() async throws {
+        struct TestParameters: JobParameters {
+            static let jobName = "testPriorityJobsWithDelay"
+            let value: Int
+        }
+        let expectation = XCTestExpectation(description: "TestJob.execute was called", expectedFulfillmentCount: 2)
+        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()
+            }
+
+            try await queue.push(
+                TestParameters(value: 20),
+                options: .init(
+                    priority: .lower()
+                )
+            )
+
+            try await queue.push(
+                TestParameters(value: 2025),
+                options: .init(
+                    delayUntil: Date.now.addingTimeInterval(1),
+                    priority: .higher()
+                )
+            )
+
+            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, 2)
+
+                group.addTask {
+                    try await serviceGroup.run()
+                }
+
+                await fulfillment(of: [expectation], timeout: 10)
+
+                let pendingJobs = try await jobQueue.queue.getJobs(withStatus: .pending)
+                XCTAssertEqual(pendingJobs.count, 0)
+                await serviceGroup.triggerGracefulShutdown()
+            }
+        }
+        XCTAssertEqual(jobExecutionSequence.withLockedValue { $0 }, [20, 2025])
+    }
+
     func testMultipleWorkers() async throws {
         struct TestParameters: JobParameters {
             static let jobName = "testMultipleWorkers"