Adding Job priority

* 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(delayed_until, queue_name, priority ASC)
             """,
             logger: logger
         )

Sources/JobsPostgres/PostgresJobsQueue.swift

@@ -55,16 +55,36 @@ public final class PostgresJobQueue: JobQueueDriver {
     public struct JobOptions: JobOptionsProtocol {
         /// Delay running job until
         public var delayUntil: Date
+        /// Priority for this jobs highest priority 0 to 9, lowest priority
+        public var priority: Int16
 
         /// Default initializer for JobOptions
         public init() {
             self.delayUntil = .now
+            self.priority = 0
         }
-
+        
         ///  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 = 0
+        }
+        
+        ///  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: Int16 = 0) {
+            self.delayUntil = delayUntil
+            self.priority = priority
+        }
+
+        ///  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?, priority: Int16 = 0) {
+            self.delayUntil = delayUntil ?? .now
+            self.priority = priority
         }
     }
 
@@ -194,7 +214,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 +228,7 @@ 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(), connection: connection)
         }
     }
 
@@ -273,7 +293,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 created_at ASC, delayed_until ASC, priority ASC
                         FOR UPDATE SKIP LOCKED
                         LIMIT 1
                     )
@@ -381,11 +401,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))
             -- We have found an existing job with the same id, SKIP this INSERT 
             ON CONFLICT (job_id) DO NOTHING
             """,
@@ -449,7 +469,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

@@ -206,6 +206,49 @@ 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([])
+
+        try await self.testJobQueue(numWorkers: 1) { jobQueue in
+            jobQueue.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()
+            }
+
+            await withThrowingTaskGroup(of: Void.self) { group in
+                for i in 0..<2 {
+                    group.addTask {
+                        try await jobQueue.push(
+                            TestParameters(value: 20 + i),
+                            options: .init(
+                                priority: Int16.random(in: 0..<9)
+                            )
+                        )
+                    }
+                }
+            }
+
+            let processingJobs = try await jobQueue.queue.getJobs(withStatus: .pending)
+            XCTAssertEqual(processingJobs.count, 2)
+
+            await fulfillment(of: [expectation], timeout: 10)
+
+            let pendingJobs = try await jobQueue.queue.getJobs(withStatus: .pending)
+            XCTAssertEqual(pendingJobs.count, 0)
+        }
+        // TODO: need to figure out ordering here
+        //XCTAssertEqual(jobExecutionSequence.withLockedValue { $0 }, [20, 21])
+    }
 
     func testMultipleWorkers() async throws {
         struct TestParameters: JobParameters {