├── .env ├── .github ├── actions │ └── postgres │ │ └── action.yml └── workflows │ ├── publish.yml │ └── toolchain.yml ├── .gitignore ├── .vscode └── settings.json ├── Cargo.toml ├── LICENSE-APACHE.txt ├── LICENSE-MIT.txt ├── README.md ├── examples └── graceful-shutdown.rs ├── migrations ├── 20210316025847_setup.down.sql ├── 20210316025847_setup.up.sql ├── 20210921115907_clear.down.sql ├── 20210921115907_clear.up.sql ├── 20211013151757_fix_mq_latest_message.down.sql ├── 20211013151757_fix_mq_latest_message.up.sql ├── 20220208120856_fix_concurrent_poll.down.sql ├── 20220208120856_fix_concurrent_poll.up.sql ├── 20220713122907_fix-clear_all-keep-nil-message.down.sql └── 20220713122907_fix-clear_all-keep-nil-message.up.sql ├── sqlxmq_macros ├── Cargo.toml └── src │ └── lib.rs ├── sqlxmq_stress ├── .env ├── Cargo.toml ├── migrations │ ├── 20210316025847_setup.down.sql │ ├── 20210316025847_setup.up.sql │ ├── 20210921115907_clear.down.sql │ ├── 20210921115907_clear.up.sql │ ├── 20211013151757_fix_mq_latest_message.down.sql │ └── 20211013151757_fix_mq_latest_message.up.sql └── src │ └── main.rs └── src ├── hidden.rs ├── lib.rs ├── registry.rs ├── runner.rs ├── spawn.rs └── utils.rs /.env: -------------------------------------------------------------------------------- 1 | DATABASE_URL=postgres://postgres:password@localhost/sqlxmq 2 | -------------------------------------------------------------------------------- /.github/actions/postgres/action.yml: -------------------------------------------------------------------------------- 1 | name: "Setup PostgreSQL database" 2 | runs: 3 | using: "composite" 4 | steps: 5 | - name: Set environment variable 6 | shell: bash 7 | run: echo "DATABASE_URL=postgres://postgres:password@localhost/sqlxmq" >> $GITHUB_ENV 8 | 9 | - name: Start PostgreSQL on Ubuntu 10 | shell: bash 11 | run: | 12 | sudo systemctl start postgresql.service 13 | pg_isready 14 | sudo -u postgres psql -c "ALTER USER postgres PASSWORD 'password'" 15 | 16 | - name: Setup database 17 | shell: bash 18 | run: cargo sqlx database setup 19 | -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | on: 2 | release: 3 | types: [published] 4 | 5 | name: Publish 6 | 7 | jobs: 8 | release: 9 | name: Release 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v4 13 | - uses: dtolnay/rust-toolchain@stable 14 | - run: cargo login ${{secrets.CARGO_TOKEN}} 15 | - run: cargo publish --manifest-path sqlxmq_macros/Cargo.toml 16 | - name: Wait for crates.io to update 17 | run: sleep 30 18 | - run: cargo publish --manifest-path Cargo.toml 19 | -------------------------------------------------------------------------------- /.github/workflows/toolchain.yml: -------------------------------------------------------------------------------- 1 | on: [push, pull_request] 2 | 3 | name: CI 4 | 5 | jobs: 6 | check: 7 | name: Check 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v4 11 | - uses: dtolnay/rust-toolchain@stable 12 | - uses: Swatinem/rust-cache@v2 13 | - run: cargo check 14 | 15 | fmt: 16 | name: Rustfmt 17 | runs-on: ubuntu-latest 18 | steps: 19 | - uses: actions/checkout@v4 20 | - uses: dtolnay/rust-toolchain@stable 21 | with: 22 | components: rustfmt 23 | - run: cargo fmt -- --check 24 | 25 | clippy: 26 | name: Clippy 27 | runs-on: ubuntu-latest 28 | steps: 29 | - uses: actions/checkout@v4 30 | - uses: dtolnay/rust-toolchain@stable 31 | with: 32 | components: clippy 33 | - uses: Swatinem/rust-cache@v2 34 | - run: cargo clippy --all-targets -- -D warnings 35 | 36 | test: 37 | name: Test 38 | runs-on: ubuntu-latest 39 | env: 40 | RUST_BACKTRACE: "1" 41 | steps: 42 | - uses: actions/checkout@v4 43 | - uses: dtolnay/rust-toolchain@stable 44 | - uses: Swatinem/rust-cache@v2 45 | - run: cargo install sqlx-cli --locked 46 | - uses: ./.github/actions/postgres 47 | - run: cargo test -- --nocapture 48 | 49 | test_nightly: 50 | name: Test (Nightly) 51 | runs-on: ubuntu-latest 52 | env: 53 | RUST_BACKTRACE: "1" 54 | steps: 55 | - uses: actions/checkout@v4 56 | - uses: dtolnay/rust-toolchain@nightly 57 | - uses: Swatinem/rust-cache@v2 58 | - run: cargo install sqlx-cli --locked 59 | - uses: ./.github/actions/postgres 60 | - run: cargo test -- --nocapture 61 | 62 | readme: 63 | name: Readme 64 | runs-on: ubuntu-latest 65 | steps: 66 | - uses: actions/checkout@v4 67 | - uses: dtolnay/rust-toolchain@stable 68 | - uses: Swatinem/rust-cache@v2 69 | - run: cargo install cargo-sync-readme --locked 70 | - name: Sync readme 71 | run: cargo sync-readme -c 72 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | Cargo.lock 3 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "rust-analyzer.cargo.features": [ 3 | "runtime-tokio-native-tls" 4 | ] 5 | } -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "sqlxmq" 3 | version = "0.6.0" 4 | authors = ["Diggory Blake "] 5 | edition = "2018" 6 | license = "MIT OR Apache-2.0" 7 | repository = "https://github.com/Diggsey/sqlxmq" 8 | description = "A reliable job queue using PostgreSQL as a backing store" 9 | readme = "README.md" 10 | documentation = "https://docs.rs/sqlxmq" 11 | 12 | [workspace] 13 | members = ["sqlxmq_macros", "sqlxmq_stress"] 14 | 15 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 16 | 17 | [dependencies] 18 | sqlx = { version = "0.8", features = ["postgres", "chrono", "uuid"] } 19 | tokio = { version = "1.8.3", features = ["full"] } 20 | dotenvy = "0.15.3" 21 | chrono = "0.4.19" 22 | uuid = { version = "1.1.2", features = ["v4"] } 23 | log = "0.4.14" 24 | serde_json = "1.0.64" 25 | serde = "1.0.124" 26 | sqlxmq_macros = { version = "0.6.0", path = "sqlxmq_macros" } 27 | anymap2 = "0.13.0" 28 | 29 | [features] 30 | default = ["runtime-tokio-native-tls"] 31 | runtime-tokio-native-tls = ["sqlx/runtime-tokio-native-tls"] 32 | runtime-tokio-rustls = ["sqlx/runtime-tokio-rustls"] 33 | 34 | [dev-dependencies] 35 | dotenvy = "0.15.3" 36 | pretty_env_logger = "0.4.0" 37 | futures = "0.3.13" 38 | tokio = { version = "1", features = ["full"] } 39 | -------------------------------------------------------------------------------- /LICENSE-APACHE.txt: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /LICENSE-MIT.txt: -------------------------------------------------------------------------------- 1 | Permission is hereby granted, free of charge, to any 2 | person obtaining a copy of this software and associated 3 | documentation files (the "Software"), to deal in the 4 | Software without restriction, including without 5 | limitation the rights to use, copy, modify, merge, 6 | publish, distribute, sublicense, and/or sell copies of 7 | the Software, and to permit persons to whom the Software 8 | is furnished to do so, subject to the following 9 | conditions: 10 | 11 | The above copyright notice and this permission notice 12 | shall be included in all copies or substantial portions 13 | of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF 16 | ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED 17 | TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A 18 | PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT 19 | SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY 20 | CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 21 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR 22 | IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 23 | DEALINGS IN THE SOFTWARE. 24 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![CI Status](https://github.com/Diggsey/sqlxmq/workflows/CI/badge.svg)](https://github.com/Diggsey/sqlxmq/actions?query=workflow%3ACI) 2 | [![Documentation](https://docs.rs/sqlxmq/badge.svg)](https://docs.rs/sqlxmq) 3 | [![crates.io](https://img.shields.io/crates/v/sqlxmq.svg)](https://crates.io/crates/sqlxmq) 4 | 5 | 6 | 7 | # sqlxmq 8 | 9 | A job queue built on `sqlx` and `PostgreSQL`. 10 | 11 | This library allows a CRUD application to run background jobs without complicating its 12 | deployment. The only runtime dependency is `PostgreSQL`, so this is ideal for applications 13 | already using a `PostgreSQL` database. 14 | 15 | Although using a SQL database as a job queue means compromising on latency of 16 | delivered jobs, there are several show-stopping issues present in ordinary job 17 | queues which are avoided altogether. 18 | 19 | With most other job queues, in-flight jobs are state that is not covered by normal 20 | database backups. Even if jobs _are_ backed up, there is no way to restore both 21 | a database and a job queue to a consistent point-in-time without manually 22 | resolving conflicts. 23 | 24 | By storing jobs in the database, existing backup procedures will store a perfectly 25 | consistent state of both in-flight jobs and persistent data. Additionally, jobs can 26 | be spawned and completed as part of other transactions, making it easy to write correct 27 | application code. 28 | 29 | Leveraging the power of `PostgreSQL`, this job queue offers several features not 30 | present in other job queues. 31 | 32 | # Features 33 | 34 | - **Send/receive multiple jobs at once.** 35 | 36 | This reduces the number of queries to the database. 37 | 38 | - **Send jobs to be executed at a future date and time.** 39 | 40 | Avoids the need for a separate scheduling system. 41 | 42 | - **Reliable delivery of jobs.** 43 | 44 | - **Automatic retries with exponential backoff.** 45 | 46 | Number of retries and initial backoff parameters are configurable. 47 | 48 | - **Transactional sending of jobs.** 49 | 50 | Avoids sending spurious jobs if a transaction is rolled back. 51 | 52 | - **Transactional completion of jobs.** 53 | 54 | If all side-effects of a job are updates to the database, this provides 55 | true exactly-once execution of jobs. 56 | 57 | - **Transactional check-pointing of jobs.** 58 | 59 | Long-running jobs can check-point their state to avoid having to restart 60 | from the beginning if there is a failure: the next retry can continue 61 | from the last check-point. 62 | 63 | - **Opt-in strictly ordered job delivery.** 64 | 65 | Jobs within the same channel will be processed strictly in-order 66 | if this option is enabled for the job. 67 | 68 | - **Fair job delivery.** 69 | 70 | A channel with a lot of jobs ready to run will not starve a channel with fewer 71 | jobs. 72 | 73 | - **Opt-in two-phase commit.** 74 | 75 | This is particularly useful on an ordered channel where a position can be "reserved" 76 | in the job order, but not committed until later. 77 | 78 | - **JSON and/or binary payloads.** 79 | 80 | Jobs can use whichever is most convenient. 81 | 82 | - **Automatic keep-alive of jobs.** 83 | 84 | Long-running jobs will automatically be "kept alive" to prevent them being 85 | retried whilst they're still ongoing. 86 | 87 | - **Concurrency limits.** 88 | 89 | Specify the minimum and maximum number of concurrent jobs each runner should 90 | handle. 91 | 92 | - **Built-in job registry via an attribute macro.** 93 | 94 | Jobs can be easily registered with a runner, and default configuration specified 95 | on a per-job basis. 96 | 97 | - **Implicit channels.** 98 | 99 | Channels are implicitly created and destroyed when jobs are sent and processed, 100 | so no setup is required. 101 | 102 | - **Channel groups.** 103 | 104 | Easily subscribe to multiple channels at once, thanks to the separation of 105 | channel name and channel arguments. 106 | 107 | - **NOTIFY-based polling.** 108 | 109 | This saves resources when few jobs are being processed. 110 | 111 | # Getting started 112 | 113 | ## Database schema 114 | 115 | This crate expects certain database tables and stored procedures to exist. 116 | You can copy the migration files from this crate into your own migrations 117 | folder. 118 | 119 | All database items created by this crate are prefixed with `mq`, so as not 120 | to conflict with your own schema. 121 | 122 | ## Defining jobs 123 | 124 | The first step is to define a function to be run on the job queue. 125 | 126 | ```rust 127 | use std::error::Error; 128 | 129 | use sqlxmq::{job, CurrentJob}; 130 | 131 | // Arguments to the `#[job]` attribute allow setting default job options. 132 | #[job(channel_name = "foo")] 133 | async fn example_job( 134 | // The first argument should always be the current job. 135 | mut current_job: CurrentJob, 136 | // Additional arguments are optional, but can be used to access context 137 | // provided via [`JobRegistry::set_context`]. 138 | message: &'static str, 139 | ) -> Result<(), Box> { 140 | // Decode a JSON payload 141 | let who: Option = current_job.json()?; 142 | 143 | // Do some work 144 | println!("{}, {}!", message, who.as_deref().unwrap_or("world")); 145 | 146 | // Mark the job as complete 147 | current_job.complete().await?; 148 | 149 | Ok(()) 150 | } 151 | ``` 152 | 153 | ## Listening for jobs 154 | 155 | Next we need to create a job runner: this is what listens for new jobs 156 | and executes them. 157 | 158 | ```rust,no_run 159 | use std::error::Error; 160 | 161 | use sqlxmq::JobRegistry; 162 | 163 | 164 | #[tokio::main] 165 | async fn main() -> Result<(), Box> { 166 | // You'll need to provide a Postgres connection pool. 167 | let pool = connect_to_db().await?; 168 | 169 | // Construct a job registry from our single job. 170 | let mut registry = JobRegistry::new(&[example_job]); 171 | // Here is where you can configure the registry 172 | // registry.set_error_handler(...) 173 | 174 | // And add context 175 | registry.set_context("Hello"); 176 | 177 | let runner = registry 178 | // Create a job runner using the connection pool. 179 | .runner(&pool) 180 | // Here is where you can configure the job runner 181 | // Aim to keep 10-20 jobs running at a time. 182 | .set_concurrency(10, 20) 183 | // Start the job runner in the background. 184 | .run() 185 | .await?; 186 | 187 | // The job runner will continue listening and running 188 | // jobs until `runner` is dropped. 189 | Ok(()) 190 | } 191 | ``` 192 | 193 | ## Spawning a job 194 | 195 | The final step is to actually run a job. 196 | 197 | ```rust 198 | example_job.builder() 199 | // This is where we can override job configuration 200 | .set_channel_name("bar") 201 | .set_json("John")? 202 | .spawn(&pool) 203 | .await?; 204 | ``` 205 | 206 | 207 | 208 | ## Note on README 209 | 210 | Most of the readme is automatically copied from the crate documentation by [cargo-readme-sync][]. 211 | This way the readme is always in sync with the docs and examples are tested. 212 | 213 | So if you find a part of the readme you'd like to change between `` 214 | and `` markers, don't edit `README.md` directly, but rather change 215 | the documentation on top of `src/lib.rs` and then synchronize the readme with: 216 | ```bash 217 | cargo sync-readme 218 | ``` 219 | (make sure the cargo command is installed): 220 | ```bash 221 | cargo install cargo-sync-readme 222 | 223 | -------------------------------------------------------------------------------- /examples/graceful-shutdown.rs: -------------------------------------------------------------------------------- 1 | use sqlxmq::{job, CurrentJob, JobRegistry}; 2 | use std::time::Duration; 3 | use tokio::time; 4 | 5 | #[tokio::main] 6 | async fn main() -> Result<(), Box> { 7 | dotenvy::dotenv().ok(); 8 | let db = sqlx::PgPool::connect(&std::env::var("DATABASE_URL").unwrap()).await?; 9 | 10 | sleep.builder().set_json(&5u64)?.spawn(&db).await?; 11 | 12 | let mut handle = JobRegistry::new(&[sleep]).runner(&db).run().await?; 13 | 14 | // Let's emulate a stop signal in a couple of seconts after running the job 15 | time::sleep(Duration::from_secs(2)).await; 16 | println!("A stop signal received"); 17 | 18 | // Stop listening for new jobs 19 | handle.stop().await; 20 | 21 | // Wait for the running jobs to stop for maximum 10 seconds 22 | handle.wait_jobs_finish(Duration::from_secs(10)).await; 23 | 24 | Ok(()) 25 | } 26 | 27 | #[job] 28 | pub async fn sleep(mut job: CurrentJob) -> sqlx::Result<()> { 29 | let second = Duration::from_secs(1); 30 | let mut to_sleep: u64 = job.json().unwrap().unwrap(); 31 | while to_sleep > 0 { 32 | println!("job#{} {to_sleep} more seconds to sleep ...", job.id()); 33 | time::sleep(second).await; 34 | to_sleep -= 1; 35 | } 36 | job.complete().await 37 | } 38 | -------------------------------------------------------------------------------- /migrations/20210316025847_setup.down.sql: -------------------------------------------------------------------------------- 1 | DROP FUNCTION mq_checkpoint; 2 | DROP FUNCTION mq_keep_alive; 3 | DROP FUNCTION mq_delete; 4 | DROP FUNCTION mq_commit; 5 | DROP FUNCTION mq_insert; 6 | DROP FUNCTION mq_poll; 7 | DROP FUNCTION mq_active_channels; 8 | DROP FUNCTION mq_latest_message; 9 | DROP TABLE mq_payloads; 10 | DROP TABLE mq_msgs; 11 | DROP FUNCTION mq_uuid_exists; 12 | DROP TYPE mq_new_t; 13 | -------------------------------------------------------------------------------- /migrations/20210316025847_setup.up.sql: -------------------------------------------------------------------------------- 1 | CREATE EXTENSION IF NOT EXISTS "uuid-ossp"; 2 | 3 | -- The UDT for creating messages 4 | CREATE TYPE mq_new_t AS ( 5 | -- Unique message ID 6 | id UUID, 7 | -- Delay before message is processed 8 | delay INTERVAL, 9 | -- Number of retries if initial processing fails 10 | retries INT, 11 | -- Initial backoff between retries 12 | retry_backoff INTERVAL, 13 | -- Name of channel 14 | channel_name TEXT, 15 | -- Arguments to channel 16 | channel_args TEXT, 17 | -- Interval for two-phase commit (or NULL to disable two-phase commit) 18 | commit_interval INTERVAL, 19 | -- Whether this message should be processed in order with respect to other 20 | -- ordered messages. 21 | ordered BOOLEAN, 22 | -- Name of message 23 | name TEXT, 24 | -- JSON payload 25 | payload_json TEXT, 26 | -- Binary payload 27 | payload_bytes BYTEA 28 | ); 29 | 30 | -- Small, frequently updated table of messages 31 | CREATE TABLE mq_msgs ( 32 | id UUID PRIMARY KEY, 33 | created_at TIMESTAMPTZ DEFAULT NOW(), 34 | attempt_at TIMESTAMPTZ DEFAULT NOW(), 35 | attempts INT NOT NULL DEFAULT 5, 36 | retry_backoff INTERVAL NOT NULL DEFAULT INTERVAL '1 second', 37 | channel_name TEXT NOT NULL, 38 | channel_args TEXT NOT NULL, 39 | commit_interval INTERVAL, 40 | after_message_id UUID DEFAULT uuid_nil() REFERENCES mq_msgs(id) ON DELETE SET DEFAULT 41 | ); 42 | 43 | -- Insert dummy message so that the 'nil' UUID can be referenced 44 | INSERT INTO mq_msgs (id, channel_name, channel_args, after_message_id) VALUES (uuid_nil(), '', '', NULL); 45 | 46 | -- Internal helper function to check that a UUID is neither NULL nor NIL 47 | CREATE FUNCTION mq_uuid_exists( 48 | id UUID 49 | ) RETURNS BOOLEAN AS $$ 50 | SELECT id IS NOT NULL AND id != uuid_nil() 51 | $$ LANGUAGE SQL IMMUTABLE; 52 | 53 | -- Index for polling 54 | CREATE INDEX ON mq_msgs(channel_name, channel_args, attempt_at) WHERE id != uuid_nil() AND NOT mq_uuid_exists(after_message_id); 55 | -- Index for adding messages 56 | CREATE INDEX ON mq_msgs(channel_name, channel_args, created_at, id) WHERE id != uuid_nil() AND after_message_id IS NOT NULL; 57 | 58 | -- Index for ensuring strict message order 59 | CREATE UNIQUE INDEX mq_msgs_channel_name_channel_args_after_message_id_idx ON mq_msgs(channel_name, channel_args, after_message_id); 60 | 61 | 62 | -- Large, less frequently updated table of message payloads 63 | CREATE TABLE mq_payloads( 64 | id UUID PRIMARY KEY, 65 | name TEXT NOT NULL, 66 | payload_json JSONB, 67 | payload_bytes BYTEA 68 | ); 69 | 70 | -- Internal helper function to return the most recently added message in a queue. 71 | CREATE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) 72 | RETURNS UUID AS $$ 73 | SELECT COALESCE( 74 | ( 75 | SELECT id FROM mq_msgs 76 | WHERE channel_name = from_channel_name 77 | AND channel_args = from_channel_args 78 | AND after_message_id IS NOT NULL 79 | AND id != uuid_nil() 80 | ORDER BY created_at DESC, id DESC 81 | LIMIT 1 82 | ), 83 | uuid_nil() 84 | ) 85 | $$ LANGUAGE SQL STABLE; 86 | 87 | -- Internal helper function to randomly select a set of channels with "ready" messages. 88 | CREATE FUNCTION mq_active_channels(channel_names TEXT[], batch_size INT) 89 | RETURNS TABLE(name TEXT, args TEXT) AS $$ 90 | SELECT channel_name, channel_args 91 | FROM mq_msgs 92 | WHERE id != uuid_nil() 93 | AND attempt_at <= NOW() 94 | AND (channel_names IS NULL OR channel_name = ANY(channel_names)) 95 | AND NOT mq_uuid_exists(after_message_id) 96 | GROUP BY channel_name, channel_args 97 | ORDER BY RANDOM() 98 | LIMIT batch_size 99 | $$ LANGUAGE SQL STABLE; 100 | 101 | -- Main entry-point for job runner: pulls a batch of messages from the queue. 102 | CREATE FUNCTION mq_poll(channel_names TEXT[], batch_size INT DEFAULT 1) 103 | RETURNS TABLE( 104 | id UUID, 105 | is_committed BOOLEAN, 106 | name TEXT, 107 | payload_json TEXT, 108 | payload_bytes BYTEA, 109 | retry_backoff INTERVAL, 110 | wait_time INTERVAL 111 | ) AS $$ 112 | BEGIN 113 | RETURN QUERY UPDATE mq_msgs 114 | SET 115 | attempt_at = CASE WHEN mq_msgs.attempts = 1 THEN NULL ELSE NOW() + mq_msgs.retry_backoff END, 116 | attempts = mq_msgs.attempts - 1, 117 | retry_backoff = mq_msgs.retry_backoff * 2 118 | FROM ( 119 | SELECT 120 | msgs.id 121 | FROM mq_active_channels(channel_names, batch_size) AS active_channels 122 | INNER JOIN LATERAL ( 123 | SELECT * FROM mq_msgs 124 | WHERE mq_msgs.id != uuid_nil() 125 | AND mq_msgs.attempt_at <= NOW() 126 | AND mq_msgs.channel_name = active_channels.name 127 | AND mq_msgs.channel_args = active_channels.args 128 | AND NOT mq_uuid_exists(mq_msgs.after_message_id) 129 | ORDER BY mq_msgs.attempt_at ASC 130 | LIMIT batch_size 131 | ) AS msgs ON TRUE 132 | LIMIT batch_size 133 | ) AS messages_to_update 134 | LEFT JOIN mq_payloads ON mq_payloads.id = messages_to_update.id 135 | WHERE mq_msgs.id = messages_to_update.id 136 | RETURNING 137 | mq_msgs.id, 138 | mq_msgs.commit_interval IS NULL, 139 | mq_payloads.name, 140 | mq_payloads.payload_json::TEXT, 141 | mq_payloads.payload_bytes, 142 | mq_msgs.retry_backoff / 2, 143 | interval '0' AS wait_time; 144 | 145 | IF NOT FOUND THEN 146 | RETURN QUERY SELECT 147 | NULL::UUID, 148 | NULL::BOOLEAN, 149 | NULL::TEXT, 150 | NULL::TEXT, 151 | NULL::BYTEA, 152 | NULL::INTERVAL, 153 | MIN(mq_msgs.attempt_at) - NOW() 154 | FROM mq_msgs 155 | WHERE mq_msgs.id != uuid_nil() 156 | AND NOT mq_uuid_exists(mq_msgs.after_message_id) 157 | AND (channel_names IS NULL OR mq_msgs.channel_name = ANY(channel_names)); 158 | END IF; 159 | END; 160 | $$ LANGUAGE plpgsql; 161 | 162 | -- Creates new messages 163 | CREATE FUNCTION mq_insert(new_messages mq_new_t[]) 164 | RETURNS VOID AS $$ 165 | BEGIN 166 | PERFORM pg_notify(CONCAT('mq_', channel_name), '') 167 | FROM unnest(new_messages) AS new_msgs 168 | GROUP BY channel_name; 169 | 170 | IF FOUND THEN 171 | PERFORM pg_notify('mq', ''); 172 | END IF; 173 | 174 | INSERT INTO mq_payloads ( 175 | id, 176 | name, 177 | payload_json, 178 | payload_bytes 179 | ) SELECT 180 | id, 181 | name, 182 | payload_json::JSONB, 183 | payload_bytes 184 | FROM UNNEST(new_messages); 185 | 186 | INSERT INTO mq_msgs ( 187 | id, 188 | attempt_at, 189 | attempts, 190 | retry_backoff, 191 | channel_name, 192 | channel_args, 193 | commit_interval, 194 | after_message_id 195 | ) 196 | SELECT 197 | id, 198 | NOW() + delay + COALESCE(commit_interval, INTERVAL '0'), 199 | retries + 1, 200 | retry_backoff, 201 | channel_name, 202 | channel_args, 203 | commit_interval, 204 | CASE WHEN ordered 205 | THEN 206 | LAG(id, 1, mq_latest_message(channel_name, channel_args)) 207 | OVER (PARTITION BY channel_name, channel_args, ordered ORDER BY id) 208 | ELSE 209 | NULL 210 | END 211 | FROM UNNEST(new_messages); 212 | END; 213 | $$ LANGUAGE plpgsql; 214 | 215 | -- Commits messages previously created with a non-NULL commit interval. 216 | CREATE FUNCTION mq_commit(msg_ids UUID[]) 217 | RETURNS VOID AS $$ 218 | BEGIN 219 | UPDATE mq_msgs 220 | SET 221 | attempt_at = attempt_at - commit_interval, 222 | commit_interval = NULL 223 | WHERE id = ANY(msg_ids) 224 | AND commit_interval IS NOT NULL; 225 | END; 226 | $$ LANGUAGE plpgsql; 227 | 228 | 229 | -- Deletes messages from the queue. This occurs when a message has been 230 | -- processed, or when it expires without being processed. 231 | CREATE FUNCTION mq_delete(msg_ids UUID[]) 232 | RETURNS VOID AS $$ 233 | BEGIN 234 | PERFORM pg_notify(CONCAT('mq_', channel_name), '') 235 | FROM mq_msgs 236 | WHERE id = ANY(msg_ids) 237 | AND after_message_id = uuid_nil() 238 | GROUP BY channel_name; 239 | 240 | IF FOUND THEN 241 | PERFORM pg_notify('mq', ''); 242 | END IF; 243 | 244 | DELETE FROM mq_msgs WHERE id = ANY(msg_ids); 245 | DELETE FROM mq_payloads WHERE id = ANY(msg_ids); 246 | END; 247 | $$ LANGUAGE plpgsql; 248 | 249 | 250 | -- Can be called during the initial commit interval, or when processing 251 | -- a message. Indicates that the caller is still active and will prevent either 252 | -- the commit interval elapsing or the message being retried for the specified 253 | -- interval. 254 | CREATE FUNCTION mq_keep_alive(msg_ids UUID[], duration INTERVAL) 255 | RETURNS VOID AS $$ 256 | UPDATE mq_msgs 257 | SET 258 | attempt_at = NOW() + duration, 259 | commit_interval = commit_interval + ((NOW() + duration) - attempt_at) 260 | WHERE id = ANY(msg_ids) 261 | AND attempt_at < NOW() + duration; 262 | $$ LANGUAGE SQL; 263 | 264 | 265 | -- Called during lengthy processing of a message to checkpoint the progress. 266 | -- As well as behaving like `mq_keep_alive`, the message payload can be 267 | -- updated. 268 | CREATE FUNCTION mq_checkpoint( 269 | msg_id UUID, 270 | duration INTERVAL, 271 | new_payload_json TEXT, 272 | new_payload_bytes BYTEA, 273 | extra_retries INT 274 | ) 275 | RETURNS VOID AS $$ 276 | UPDATE mq_msgs 277 | SET 278 | attempt_at = GREATEST(attempt_at, NOW() + duration), 279 | attempts = attempts + COALESCE(extra_retries, 0) 280 | WHERE id = msg_id; 281 | 282 | UPDATE mq_payloads 283 | SET 284 | payload_json = COALESCE(new_payload_json::JSONB, payload_json), 285 | payload_bytes = COALESCE(new_payload_bytes, payload_bytes) 286 | WHERE 287 | id = msg_id; 288 | $$ LANGUAGE SQL; 289 | 290 | -------------------------------------------------------------------------------- /migrations/20210921115907_clear.down.sql: -------------------------------------------------------------------------------- 1 | DROP FUNCTION mq_clear; 2 | DROP FUNCTION mq_clear_all; 3 | -------------------------------------------------------------------------------- /migrations/20210921115907_clear.up.sql: -------------------------------------------------------------------------------- 1 | -- Deletes all messages from a list of channel names. 2 | CREATE FUNCTION mq_clear(channel_names TEXT[]) 3 | RETURNS VOID AS $$ 4 | BEGIN 5 | WITH deleted_ids AS ( 6 | DELETE FROM mq_msgs WHERE channel_name = ANY(channel_names) RETURNING id 7 | ) 8 | DELETE FROM mq_payloads WHERE id IN (SELECT id FROM deleted_ids); 9 | END; 10 | $$ LANGUAGE plpgsql; 11 | 12 | -- Deletes all messages. 13 | CREATE FUNCTION mq_clear_all() 14 | RETURNS VOID AS $$ 15 | BEGIN 16 | WITH deleted_ids AS ( 17 | DELETE FROM mq_msgs RETURNING id 18 | ) 19 | DELETE FROM mq_payloads WHERE id IN (SELECT id FROM deleted_ids); 20 | END; 21 | $$ LANGUAGE plpgsql; 22 | -------------------------------------------------------------------------------- /migrations/20211013151757_fix_mq_latest_message.down.sql: -------------------------------------------------------------------------------- 1 | CREATE OR REPLACE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) 2 | RETURNS UUID AS $$ 3 | SELECT COALESCE( 4 | ( 5 | SELECT id FROM mq_msgs 6 | WHERE channel_name = from_channel_name 7 | AND channel_args = from_channel_args 8 | AND after_message_id IS NOT NULL 9 | AND id != uuid_nil() 10 | ORDER BY created_at DESC, id DESC 11 | LIMIT 1 12 | ), 13 | uuid_nil() 14 | ) 15 | $$ LANGUAGE SQL STABLE; 16 | -------------------------------------------------------------------------------- /migrations/20211013151757_fix_mq_latest_message.up.sql: -------------------------------------------------------------------------------- 1 | CREATE OR REPLACE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) 2 | RETURNS UUID AS $$ 3 | SELECT COALESCE( 4 | ( 5 | SELECT id FROM mq_msgs 6 | WHERE channel_name = from_channel_name 7 | AND channel_args = from_channel_args 8 | AND after_message_id IS NOT NULL 9 | AND id != uuid_nil() 10 | AND NOT EXISTS( 11 | SELECT * FROM mq_msgs AS mq_msgs2 12 | WHERE mq_msgs2.after_message_id = mq_msgs.id 13 | ) 14 | ORDER BY created_at DESC 15 | LIMIT 1 16 | ), 17 | uuid_nil() 18 | ) 19 | $$ LANGUAGE SQL STABLE; -------------------------------------------------------------------------------- /migrations/20220208120856_fix_concurrent_poll.down.sql: -------------------------------------------------------------------------------- 1 | -- Main entry-point for job runner: pulls a batch of messages from the queue. 2 | CREATE OR REPLACE FUNCTION mq_poll(channel_names TEXT[], batch_size INT DEFAULT 1) 3 | RETURNS TABLE( 4 | id UUID, 5 | is_committed BOOLEAN, 6 | name TEXT, 7 | payload_json TEXT, 8 | payload_bytes BYTEA, 9 | retry_backoff INTERVAL, 10 | wait_time INTERVAL 11 | ) AS $$ 12 | BEGIN 13 | RETURN QUERY UPDATE mq_msgs 14 | SET 15 | attempt_at = CASE WHEN mq_msgs.attempts = 1 THEN NULL ELSE NOW() + mq_msgs.retry_backoff END, 16 | attempts = mq_msgs.attempts - 1, 17 | retry_backoff = mq_msgs.retry_backoff * 2 18 | FROM ( 19 | SELECT 20 | msgs.id 21 | FROM mq_active_channels(channel_names, batch_size) AS active_channels 22 | INNER JOIN LATERAL ( 23 | SELECT * FROM mq_msgs 24 | WHERE mq_msgs.id != uuid_nil() 25 | AND mq_msgs.attempt_at <= NOW() 26 | AND mq_msgs.channel_name = active_channels.name 27 | AND mq_msgs.channel_args = active_channels.args 28 | AND NOT mq_uuid_exists(mq_msgs.after_message_id) 29 | ORDER BY mq_msgs.attempt_at ASC 30 | LIMIT batch_size 31 | ) AS msgs ON TRUE 32 | LIMIT batch_size 33 | ) AS messages_to_update 34 | LEFT JOIN mq_payloads ON mq_payloads.id = messages_to_update.id 35 | WHERE mq_msgs.id = messages_to_update.id 36 | RETURNING 37 | mq_msgs.id, 38 | mq_msgs.commit_interval IS NULL, 39 | mq_payloads.name, 40 | mq_payloads.payload_json::TEXT, 41 | mq_payloads.payload_bytes, 42 | mq_msgs.retry_backoff / 2, 43 | interval '0' AS wait_time; 44 | 45 | IF NOT FOUND THEN 46 | RETURN QUERY SELECT 47 | NULL::UUID, 48 | NULL::BOOLEAN, 49 | NULL::TEXT, 50 | NULL::TEXT, 51 | NULL::BYTEA, 52 | NULL::INTERVAL, 53 | MIN(mq_msgs.attempt_at) - NOW() 54 | FROM mq_msgs 55 | WHERE mq_msgs.id != uuid_nil() 56 | AND NOT mq_uuid_exists(mq_msgs.after_message_id) 57 | AND (channel_names IS NULL OR mq_msgs.channel_name = ANY(channel_names)); 58 | END IF; 59 | END; 60 | $$ LANGUAGE plpgsql; -------------------------------------------------------------------------------- /migrations/20220208120856_fix_concurrent_poll.up.sql: -------------------------------------------------------------------------------- 1 | 2 | -- Main entry-point for job runner: pulls a batch of messages from the queue. 3 | CREATE OR REPLACE FUNCTION mq_poll(channel_names TEXT[], batch_size INT DEFAULT 1) 4 | RETURNS TABLE( 5 | id UUID, 6 | is_committed BOOLEAN, 7 | name TEXT, 8 | payload_json TEXT, 9 | payload_bytes BYTEA, 10 | retry_backoff INTERVAL, 11 | wait_time INTERVAL 12 | ) AS $$ 13 | BEGIN 14 | RETURN QUERY UPDATE mq_msgs 15 | SET 16 | attempt_at = CASE WHEN mq_msgs.attempts = 1 THEN NULL ELSE NOW() + mq_msgs.retry_backoff END, 17 | attempts = mq_msgs.attempts - 1, 18 | retry_backoff = mq_msgs.retry_backoff * 2 19 | FROM ( 20 | SELECT 21 | msgs.id 22 | FROM mq_active_channels(channel_names, batch_size) AS active_channels 23 | INNER JOIN LATERAL ( 24 | SELECT mq_msgs.id FROM mq_msgs 25 | WHERE mq_msgs.id != uuid_nil() 26 | AND mq_msgs.attempt_at <= NOW() 27 | AND mq_msgs.channel_name = active_channels.name 28 | AND mq_msgs.channel_args = active_channels.args 29 | AND NOT mq_uuid_exists(mq_msgs.after_message_id) 30 | ORDER BY mq_msgs.attempt_at ASC 31 | LIMIT batch_size 32 | ) AS msgs ON TRUE 33 | LIMIT batch_size 34 | ) AS messages_to_update 35 | LEFT JOIN mq_payloads ON mq_payloads.id = messages_to_update.id 36 | WHERE mq_msgs.id = messages_to_update.id 37 | AND mq_msgs.attempt_at <= NOW() 38 | RETURNING 39 | mq_msgs.id, 40 | mq_msgs.commit_interval IS NULL, 41 | mq_payloads.name, 42 | mq_payloads.payload_json::TEXT, 43 | mq_payloads.payload_bytes, 44 | mq_msgs.retry_backoff / 2, 45 | interval '0' AS wait_time; 46 | 47 | IF NOT FOUND THEN 48 | RETURN QUERY SELECT 49 | NULL::UUID, 50 | NULL::BOOLEAN, 51 | NULL::TEXT, 52 | NULL::TEXT, 53 | NULL::BYTEA, 54 | NULL::INTERVAL, 55 | MIN(mq_msgs.attempt_at) - NOW() 56 | FROM mq_msgs 57 | WHERE mq_msgs.id != uuid_nil() 58 | AND NOT mq_uuid_exists(mq_msgs.after_message_id) 59 | AND (channel_names IS NULL OR mq_msgs.channel_name = ANY(channel_names)); 60 | END IF; 61 | END; 62 | $$ LANGUAGE plpgsql; 63 | -------------------------------------------------------------------------------- /migrations/20220713122907_fix-clear_all-keep-nil-message.down.sql: -------------------------------------------------------------------------------- 1 | -- Add down migration script here 2 | -------------------------------------------------------------------------------- /migrations/20220713122907_fix-clear_all-keep-nil-message.up.sql: -------------------------------------------------------------------------------- 1 | CREATE OR REPLACE FUNCTION mq_clear(channel_names TEXT[]) 2 | RETURNS VOID AS $$ 3 | BEGIN 4 | WITH deleted_ids AS ( 5 | DELETE FROM mq_msgs 6 | WHERE channel_name = ANY(channel_names) 7 | AND id != uuid_nil() 8 | RETURNING id 9 | ) 10 | DELETE FROM mq_payloads WHERE id IN (SELECT id FROM deleted_ids); 11 | END; 12 | $$ LANGUAGE plpgsql; 13 | COMMENT ON FUNCTION mq_clear IS 14 | 'Deletes all messages with corresponding payloads from a list of channel names'; 15 | 16 | 17 | CREATE OR REPLACE FUNCTION mq_clear_all() 18 | RETURNS VOID AS $$ 19 | BEGIN 20 | WITH deleted_ids AS ( 21 | DELETE FROM mq_msgs 22 | WHERE id != uuid_nil() 23 | RETURNING id 24 | ) 25 | DELETE FROM mq_payloads WHERE id IN (SELECT id FROM deleted_ids); 26 | END; 27 | $$ LANGUAGE plpgsql; 28 | COMMENT ON FUNCTION mq_clear_all IS 29 | 'Deletes all messages with corresponding payloads'; 30 | -------------------------------------------------------------------------------- /sqlxmq_macros/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "sqlxmq_macros" 3 | version = "0.6.0" 4 | authors = ["Diggory Blake "] 5 | edition = "2018" 6 | license = "MIT OR Apache-2.0" 7 | repository = "https://github.com/Diggsey/sqlxmq" 8 | description = "Procedural macros for sqlxmq" 9 | readme = "../README.md" 10 | documentation = "https://docs.rs/sqlxmq" 11 | 12 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 13 | [lib] 14 | proc-macro = true 15 | 16 | [dependencies] 17 | syn = { version = "1.0.80", features = ["full"] } 18 | quote = "1.0.10" 19 | proc-macro2 = "1.0.30" 20 | -------------------------------------------------------------------------------- /sqlxmq_macros/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![deny(missing_docs, unsafe_code)] 2 | //! # sqlxmq_macros 3 | //! 4 | //! Provides procedural macros for the `sqlxmq` crate. 5 | 6 | use std::mem; 7 | 8 | use proc_macro::TokenStream; 9 | use proc_macro2::TokenStream as TokenStream2; 10 | use quote::{quote, ToTokens, TokenStreamExt}; 11 | use syn::{ 12 | parse::{Parse, ParseStream}, 13 | parse_macro_input, parse_quote, AttrStyle, Attribute, AttributeArgs, Error, Lit, Meta, 14 | NestedMeta, Path, Result, Signature, Visibility, 15 | }; 16 | 17 | #[derive(Default)] 18 | struct JobOptions { 19 | proto: Option, 20 | name: Option, 21 | channel_name: Option, 22 | retries: Option, 23 | backoff_secs: Option, 24 | ordered: Option, 25 | } 26 | 27 | enum OptionValue<'a> { 28 | None, 29 | Lit(&'a Lit), 30 | Path(&'a Path), 31 | } 32 | 33 | fn interpret_job_arg(options: &mut JobOptions, arg: NestedMeta) -> Result<()> { 34 | fn error(arg: NestedMeta) -> Result<()> { 35 | Err(Error::new_spanned(arg, "Unexpected attribute argument")) 36 | } 37 | match &arg { 38 | NestedMeta::Lit(Lit::Str(s)) if options.name.is_none() => { 39 | options.name = Some(s.value()); 40 | } 41 | NestedMeta::Meta(m) => { 42 | if let Some(ident) = m.path().get_ident() { 43 | let name = ident.to_string(); 44 | let value = match &m { 45 | Meta::List(l) => { 46 | if let NestedMeta::Meta(Meta::Path(p)) = &l.nested[0] { 47 | OptionValue::Path(p) 48 | } else { 49 | return error(arg); 50 | } 51 | } 52 | Meta::Path(_) => OptionValue::None, 53 | Meta::NameValue(nvp) => OptionValue::Lit(&nvp.lit), 54 | }; 55 | match (name.as_str(), value) { 56 | ("proto", OptionValue::Path(p)) if options.proto.is_none() => { 57 | options.proto = Some(p.clone()); 58 | } 59 | ("name", OptionValue::Lit(Lit::Str(s))) if options.name.is_none() => { 60 | options.name = Some(s.value()); 61 | } 62 | ("channel_name", OptionValue::Lit(Lit::Str(s))) 63 | if options.channel_name.is_none() => 64 | { 65 | options.channel_name = Some(s.value()); 66 | } 67 | ("retries", OptionValue::Lit(Lit::Int(n))) if options.retries.is_none() => { 68 | options.retries = Some(n.base10_parse()?); 69 | } 70 | ("backoff_secs", OptionValue::Lit(Lit::Float(n))) 71 | if options.backoff_secs.is_none() => 72 | { 73 | options.backoff_secs = Some(n.base10_parse()?); 74 | } 75 | ("backoff_secs", OptionValue::Lit(Lit::Int(n))) 76 | if options.backoff_secs.is_none() => 77 | { 78 | options.backoff_secs = Some(n.base10_parse()?); 79 | } 80 | ("ordered", OptionValue::None) if options.ordered.is_none() => { 81 | options.ordered = Some(true); 82 | } 83 | ("ordered", OptionValue::Lit(Lit::Bool(b))) if options.ordered.is_none() => { 84 | options.ordered = Some(b.value); 85 | } 86 | _ => return error(arg), 87 | } 88 | } 89 | } 90 | _ => return error(arg), 91 | } 92 | Ok(()) 93 | } 94 | 95 | #[derive(Clone)] 96 | struct MaybeItemFn { 97 | attrs: Vec, 98 | vis: Visibility, 99 | sig: Signature, 100 | block: TokenStream2, 101 | } 102 | 103 | /// This parses a `TokenStream` into a `MaybeItemFn` 104 | /// (just like `ItemFn`, but skips parsing the body). 105 | impl Parse for MaybeItemFn { 106 | fn parse(input: ParseStream<'_>) -> syn::Result { 107 | let attrs = input.call(syn::Attribute::parse_outer)?; 108 | let vis: Visibility = input.parse()?; 109 | let sig: Signature = input.parse()?; 110 | let block: TokenStream2 = input.parse()?; 111 | Ok(Self { 112 | attrs, 113 | vis, 114 | sig, 115 | block, 116 | }) 117 | } 118 | } 119 | 120 | impl ToTokens for MaybeItemFn { 121 | fn to_tokens(&self, tokens: &mut TokenStream2) { 122 | tokens.append_all( 123 | self.attrs 124 | .iter() 125 | .filter(|attr| matches!(attr.style, AttrStyle::Outer)), 126 | ); 127 | self.vis.to_tokens(tokens); 128 | self.sig.to_tokens(tokens); 129 | self.block.to_tokens(tokens); 130 | } 131 | } 132 | 133 | /// Marks a function as being a background job. 134 | /// 135 | /// The first argument to the function must have type `CurrentJob`. 136 | /// Additional arguments can be used to access context from the job 137 | /// registry. Context is accessed based on the type of the argument. 138 | /// Context arguments must be `Send + Sync + Clone + 'static`. 139 | /// 140 | /// The function should be async or return a future. 141 | /// 142 | /// The async result must be a `Result<(), E>` type, where `E` is convertible 143 | /// to a `Box`, which is the case for most 144 | /// error types. 145 | /// 146 | /// Several options can be provided to the `#[job]` attribute: 147 | /// 148 | /// # Name 149 | /// 150 | /// ```ignore 151 | /// #[job("example")] 152 | /// #[job(name="example")] 153 | /// ``` 154 | /// 155 | /// This overrides the name for this job. If unspecified, the fully-qualified 156 | /// name of the function is used. If you move a job to a new module or rename 157 | /// the function, you may which to override the job name to prevent it from 158 | /// changing. 159 | /// 160 | /// # Channel name 161 | /// 162 | /// ```ignore 163 | /// #[job(channel_name="foo")] 164 | /// ``` 165 | /// 166 | /// This sets the default channel name on which the job will be spawned. 167 | /// 168 | /// # Retries 169 | /// 170 | /// ```ignore 171 | /// #[job(retries = 3)] 172 | /// ``` 173 | /// 174 | /// This sets the default number of retries for the job. 175 | /// 176 | /// # Retry backoff 177 | /// 178 | /// ```ignore 179 | /// #[job(backoff_secs=1.5)] 180 | /// #[job(backoff_secs=2)] 181 | /// ``` 182 | /// 183 | /// This sets the default initial retry backoff for the job in seconds. 184 | /// 185 | /// # Ordered 186 | /// 187 | /// ```ignore 188 | /// #[job(ordered)] 189 | /// #[job(ordered=true)] 190 | /// #[job(ordered=false)] 191 | /// ``` 192 | /// 193 | /// This sets whether the job will be strictly ordered by default. 194 | /// 195 | /// # Prototype 196 | /// 197 | /// ```ignore 198 | /// fn my_proto<'a, 'b>( 199 | /// builder: &'a mut JobBuilder<'b> 200 | /// ) -> &'a mut JobBuilder<'b> { 201 | /// builder.set_channel_name("bar") 202 | /// } 203 | /// 204 | /// #[job(proto(my_proto))] 205 | /// ``` 206 | /// 207 | /// This allows setting several job options at once using the specified function, 208 | /// and can be convient if you have several jobs which should have similar 209 | /// defaults. 210 | /// 211 | /// # Combinations 212 | /// 213 | /// Multiple job options can be combined. The order is not important, but the 214 | /// prototype will always be applied first so that explicit options can override it. 215 | /// Each option can only be provided once in the attribute. 216 | /// 217 | /// ```ignore 218 | /// #[job("my_job", proto(my_proto), retries=0, ordered)] 219 | /// ``` 220 | /// 221 | #[proc_macro_attribute] 222 | pub fn job(attr: TokenStream, item: TokenStream) -> TokenStream { 223 | let args = parse_macro_input!(attr as AttributeArgs); 224 | let mut inner_fn = parse_macro_input!(item as MaybeItemFn); 225 | 226 | let mut options = JobOptions::default(); 227 | let mut errors = Vec::new(); 228 | for arg in args { 229 | if let Err(e) = interpret_job_arg(&mut options, arg) { 230 | errors.push(e.into_compile_error()); 231 | } 232 | } 233 | 234 | let outer_docs = inner_fn 235 | .attrs 236 | .iter() 237 | .filter(|attr| attr.path.is_ident("doc")); 238 | 239 | let vis = mem::replace(&mut inner_fn.vis, Visibility::Inherited); 240 | let name = mem::replace(&mut inner_fn.sig.ident, parse_quote! {inner}); 241 | let fq_name = if let Some(name) = options.name { 242 | quote! { #name } 243 | } else { 244 | let name_str = name.to_string(); 245 | quote! { concat!(module_path!(), "::", #name_str) } 246 | }; 247 | 248 | let mut chain = Vec::new(); 249 | if let Some(proto) = &options.proto { 250 | chain.push(quote! { 251 | .set_proto(#proto) 252 | }); 253 | } 254 | if let Some(channel_name) = &options.channel_name { 255 | chain.push(quote! { 256 | .set_channel_name(#channel_name) 257 | }); 258 | } 259 | if let Some(retries) = &options.retries { 260 | chain.push(quote! { 261 | .set_retries(#retries) 262 | }); 263 | } 264 | if let Some(backoff_secs) = &options.backoff_secs { 265 | chain.push(quote! { 266 | .set_retry_backoff(::std::time::Duration::from_secs_f64(#backoff_secs)) 267 | }); 268 | } 269 | if let Some(ordered) = options.ordered { 270 | chain.push(quote! { 271 | .set_ordered(#ordered) 272 | }); 273 | } 274 | 275 | let extract_ctx: Vec<_> = inner_fn 276 | .sig 277 | .inputs 278 | .iter() 279 | .skip(1) 280 | .map(|_| { 281 | quote! { 282 | registry.context() 283 | } 284 | }) 285 | .collect(); 286 | 287 | let expanded = quote! { 288 | #(#errors)* 289 | #(#outer_docs)* 290 | #[allow(non_upper_case_globals)] 291 | #vis static #name: &'static sqlxmq::NamedJob = &{ 292 | #inner_fn 293 | sqlxmq::NamedJob::new_internal( 294 | #fq_name, 295 | sqlxmq::hidden::BuildFn(|builder| { 296 | builder #(#chain)* 297 | }), 298 | sqlxmq::hidden::RunFn(|registry, current_job| { 299 | registry.spawn_internal(#fq_name, inner(current_job #(, #extract_ctx)*)); 300 | }), 301 | ) 302 | }; 303 | }; 304 | // Hand the output tokens back to the compiler. 305 | TokenStream::from(expanded) 306 | } 307 | -------------------------------------------------------------------------------- /sqlxmq_stress/.env: -------------------------------------------------------------------------------- 1 | DATABASE_URL=postgres://postgres:password@localhost/sqlxmq_stress 2 | -------------------------------------------------------------------------------- /sqlxmq_stress/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "sqlxmq_stress" 3 | version = "0.1.0" 4 | authors = ["Diggory Blake "] 5 | edition = "2018" 6 | 7 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 8 | 9 | [dependencies] 10 | sqlxmq = { path = ".." } 11 | tokio = { version = "1.4.0", features = ["full"] } 12 | dotenvy = "0.15" 13 | sqlx = "0.8" 14 | serde = "1.0.125" 15 | lazy_static = "1.4.0" 16 | futures = "0.3.13" 17 | -------------------------------------------------------------------------------- /sqlxmq_stress/migrations/20210316025847_setup.down.sql: -------------------------------------------------------------------------------- 1 | DROP FUNCTION mq_checkpoint; 2 | DROP FUNCTION mq_keep_alive; 3 | DROP FUNCTION mq_delete; 4 | DROP FUNCTION mq_commit; 5 | DROP FUNCTION mq_insert; 6 | DROP FUNCTION mq_poll; 7 | DROP FUNCTION mq_active_channels; 8 | DROP FUNCTION mq_latest_message; 9 | DROP TABLE mq_payloads; 10 | DROP TABLE mq_msgs; 11 | DROP FUNCTION mq_uuid_exists; 12 | DROP TYPE mq_new_t; 13 | -------------------------------------------------------------------------------- /sqlxmq_stress/migrations/20210316025847_setup.up.sql: -------------------------------------------------------------------------------- 1 | CREATE EXTENSION IF NOT EXISTS "uuid-ossp"; 2 | 3 | -- The UDT for creating messages 4 | CREATE TYPE mq_new_t AS ( 5 | -- Unique message ID 6 | id UUID, 7 | -- Delay before message is processed 8 | delay INTERVAL, 9 | -- Number of retries if initial processing fails 10 | retries INT, 11 | -- Initial backoff between retries 12 | retry_backoff INTERVAL, 13 | -- Name of channel 14 | channel_name TEXT, 15 | -- Arguments to channel 16 | channel_args TEXT, 17 | -- Interval for two-phase commit (or NULL to disable two-phase commit) 18 | commit_interval INTERVAL, 19 | -- Whether this message should be processed in order with respect to other 20 | -- ordered messages. 21 | ordered BOOLEAN, 22 | -- Name of message 23 | name TEXT, 24 | -- JSON payload 25 | payload_json TEXT, 26 | -- Binary payload 27 | payload_bytes BYTEA 28 | ); 29 | 30 | -- Small, frequently updated table of messages 31 | CREATE TABLE mq_msgs ( 32 | id UUID PRIMARY KEY, 33 | created_at TIMESTAMPTZ DEFAULT NOW(), 34 | attempt_at TIMESTAMPTZ DEFAULT NOW(), 35 | attempts INT NOT NULL DEFAULT 5, 36 | retry_backoff INTERVAL NOT NULL DEFAULT INTERVAL '1 second', 37 | channel_name TEXT NOT NULL, 38 | channel_args TEXT NOT NULL, 39 | commit_interval INTERVAL, 40 | after_message_id UUID DEFAULT uuid_nil() REFERENCES mq_msgs(id) ON DELETE SET DEFAULT 41 | ); 42 | 43 | -- Insert dummy message so that the 'nil' UUID can be referenced 44 | INSERT INTO mq_msgs (id, channel_name, channel_args, after_message_id) VALUES (uuid_nil(), '', '', NULL); 45 | 46 | -- Internal helper function to check that a UUID is neither NULL nor NIL 47 | CREATE FUNCTION mq_uuid_exists( 48 | id UUID 49 | ) RETURNS BOOLEAN AS $$ 50 | SELECT id IS NOT NULL AND id != uuid_nil() 51 | $$ LANGUAGE SQL IMMUTABLE; 52 | 53 | -- Index for polling 54 | CREATE INDEX ON mq_msgs(channel_name, channel_args, attempt_at) WHERE id != uuid_nil() AND NOT mq_uuid_exists(after_message_id); 55 | -- Index for adding messages 56 | CREATE INDEX ON mq_msgs(channel_name, channel_args, created_at, id) WHERE id != uuid_nil() AND after_message_id IS NOT NULL; 57 | 58 | -- Index for ensuring strict message order 59 | CREATE UNIQUE INDEX mq_msgs_channel_name_channel_args_after_message_id_idx ON mq_msgs(channel_name, channel_args, after_message_id); 60 | 61 | 62 | -- Large, less frequently updated table of message payloads 63 | CREATE TABLE mq_payloads( 64 | id UUID PRIMARY KEY, 65 | name TEXT NOT NULL, 66 | payload_json JSONB, 67 | payload_bytes BYTEA 68 | ); 69 | 70 | -- Internal helper function to return the most recently added message in a queue. 71 | CREATE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) 72 | RETURNS UUID AS $$ 73 | SELECT COALESCE( 74 | ( 75 | SELECT id FROM mq_msgs 76 | WHERE channel_name = from_channel_name 77 | AND channel_args = from_channel_args 78 | AND after_message_id IS NOT NULL 79 | AND id != uuid_nil() 80 | ORDER BY created_at DESC, id DESC 81 | LIMIT 1 82 | ), 83 | uuid_nil() 84 | ) 85 | $$ LANGUAGE SQL STABLE; 86 | 87 | -- Internal helper function to randomly select a set of channels with "ready" messages. 88 | CREATE FUNCTION mq_active_channels(channel_names TEXT[], batch_size INT) 89 | RETURNS TABLE(name TEXT, args TEXT) AS $$ 90 | SELECT channel_name, channel_args 91 | FROM mq_msgs 92 | WHERE id != uuid_nil() 93 | AND attempt_at <= NOW() 94 | AND (channel_names IS NULL OR channel_name = ANY(channel_names)) 95 | AND NOT mq_uuid_exists(after_message_id) 96 | GROUP BY channel_name, channel_args 97 | ORDER BY RANDOM() 98 | LIMIT batch_size 99 | $$ LANGUAGE SQL STABLE; 100 | 101 | -- Main entry-point for job runner: pulls a batch of messages from the queue. 102 | CREATE FUNCTION mq_poll(channel_names TEXT[], batch_size INT DEFAULT 1) 103 | RETURNS TABLE( 104 | id UUID, 105 | is_committed BOOLEAN, 106 | name TEXT, 107 | payload_json TEXT, 108 | payload_bytes BYTEA, 109 | retry_backoff INTERVAL, 110 | wait_time INTERVAL 111 | ) AS $$ 112 | BEGIN 113 | RETURN QUERY UPDATE mq_msgs 114 | SET 115 | attempt_at = CASE WHEN mq_msgs.attempts = 1 THEN NULL ELSE NOW() + mq_msgs.retry_backoff END, 116 | attempts = mq_msgs.attempts - 1, 117 | retry_backoff = mq_msgs.retry_backoff * 2 118 | FROM ( 119 | SELECT 120 | msgs.id 121 | FROM mq_active_channels(channel_names, batch_size) AS active_channels 122 | INNER JOIN LATERAL ( 123 | SELECT * FROM mq_msgs 124 | WHERE mq_msgs.id != uuid_nil() 125 | AND mq_msgs.attempt_at <= NOW() 126 | AND mq_msgs.channel_name = active_channels.name 127 | AND mq_msgs.channel_args = active_channels.args 128 | AND NOT mq_uuid_exists(mq_msgs.after_message_id) 129 | ORDER BY mq_msgs.attempt_at ASC 130 | LIMIT batch_size 131 | ) AS msgs ON TRUE 132 | LIMIT batch_size 133 | ) AS messages_to_update 134 | LEFT JOIN mq_payloads ON mq_payloads.id = messages_to_update.id 135 | WHERE mq_msgs.id = messages_to_update.id 136 | RETURNING 137 | mq_msgs.id, 138 | mq_msgs.commit_interval IS NULL, 139 | mq_payloads.name, 140 | mq_payloads.payload_json::TEXT, 141 | mq_payloads.payload_bytes, 142 | mq_msgs.retry_backoff / 2, 143 | interval '0' AS wait_time; 144 | 145 | IF NOT FOUND THEN 146 | RETURN QUERY SELECT 147 | NULL::UUID, 148 | NULL::BOOLEAN, 149 | NULL::TEXT, 150 | NULL::TEXT, 151 | NULL::BYTEA, 152 | NULL::INTERVAL, 153 | MIN(mq_msgs.attempt_at) - NOW() 154 | FROM mq_msgs 155 | WHERE mq_msgs.id != uuid_nil() 156 | AND NOT mq_uuid_exists(mq_msgs.after_message_id) 157 | AND (channel_names IS NULL OR mq_msgs.channel_name = ANY(channel_names)); 158 | END IF; 159 | END; 160 | $$ LANGUAGE plpgsql; 161 | 162 | -- Creates new messages 163 | CREATE FUNCTION mq_insert(new_messages mq_new_t[]) 164 | RETURNS VOID AS $$ 165 | BEGIN 166 | PERFORM pg_notify(CONCAT('mq_', channel_name), '') 167 | FROM unnest(new_messages) AS new_msgs 168 | GROUP BY channel_name; 169 | 170 | IF FOUND THEN 171 | PERFORM pg_notify('mq', ''); 172 | END IF; 173 | 174 | INSERT INTO mq_payloads ( 175 | id, 176 | name, 177 | payload_json, 178 | payload_bytes 179 | ) SELECT 180 | id, 181 | name, 182 | payload_json::JSONB, 183 | payload_bytes 184 | FROM UNNEST(new_messages); 185 | 186 | INSERT INTO mq_msgs ( 187 | id, 188 | attempt_at, 189 | attempts, 190 | retry_backoff, 191 | channel_name, 192 | channel_args, 193 | commit_interval, 194 | after_message_id 195 | ) 196 | SELECT 197 | id, 198 | NOW() + delay + COALESCE(commit_interval, INTERVAL '0'), 199 | retries + 1, 200 | retry_backoff, 201 | channel_name, 202 | channel_args, 203 | commit_interval, 204 | CASE WHEN ordered 205 | THEN 206 | LAG(id, 1, mq_latest_message(channel_name, channel_args)) 207 | OVER (PARTITION BY channel_name, channel_args, ordered ORDER BY id) 208 | ELSE 209 | NULL 210 | END 211 | FROM UNNEST(new_messages); 212 | END; 213 | $$ LANGUAGE plpgsql; 214 | 215 | -- Commits messages previously created with a non-NULL commit interval. 216 | CREATE FUNCTION mq_commit(msg_ids UUID[]) 217 | RETURNS VOID AS $$ 218 | BEGIN 219 | UPDATE mq_msgs 220 | SET 221 | attempt_at = attempt_at - commit_interval, 222 | commit_interval = NULL 223 | WHERE id = ANY(msg_ids) 224 | AND commit_interval IS NOT NULL; 225 | END; 226 | $$ LANGUAGE plpgsql; 227 | 228 | 229 | -- Deletes messages from the queue. This occurs when a message has been 230 | -- processed, or when it expires without being processed. 231 | CREATE FUNCTION mq_delete(msg_ids UUID[]) 232 | RETURNS VOID AS $$ 233 | BEGIN 234 | PERFORM pg_notify(CONCAT('mq_', channel_name), '') 235 | FROM mq_msgs 236 | WHERE id = ANY(msg_ids) 237 | AND after_message_id = uuid_nil() 238 | GROUP BY channel_name; 239 | 240 | IF FOUND THEN 241 | PERFORM pg_notify('mq', ''); 242 | END IF; 243 | 244 | DELETE FROM mq_msgs WHERE id = ANY(msg_ids); 245 | DELETE FROM mq_payloads WHERE id = ANY(msg_ids); 246 | END; 247 | $$ LANGUAGE plpgsql; 248 | 249 | 250 | -- Can be called during the initial commit interval, or when processing 251 | -- a message. Indicates that the caller is still active and will prevent either 252 | -- the commit interval elapsing or the message being retried for the specified 253 | -- interval. 254 | CREATE FUNCTION mq_keep_alive(msg_ids UUID[], duration INTERVAL) 255 | RETURNS VOID AS $$ 256 | UPDATE mq_msgs 257 | SET 258 | attempt_at = NOW() + duration, 259 | commit_interval = commit_interval + ((NOW() + duration) - attempt_at) 260 | WHERE id = ANY(msg_ids) 261 | AND attempt_at < NOW() + duration; 262 | $$ LANGUAGE SQL; 263 | 264 | 265 | -- Called during lengthy processing of a message to checkpoint the progress. 266 | -- As well as behaving like `mq_keep_alive`, the message payload can be 267 | -- updated. 268 | CREATE FUNCTION mq_checkpoint( 269 | msg_id UUID, 270 | duration INTERVAL, 271 | new_payload_json TEXT, 272 | new_payload_bytes BYTEA, 273 | extra_retries INT 274 | ) 275 | RETURNS VOID AS $$ 276 | UPDATE mq_msgs 277 | SET 278 | attempt_at = GREATEST(attempt_at, NOW() + duration), 279 | attempts = attempts + COALESCE(extra_retries, 0) 280 | WHERE id = msg_id; 281 | 282 | UPDATE mq_payloads 283 | SET 284 | payload_json = COALESCE(new_payload_json::JSONB, payload_json), 285 | payload_bytes = COALESCE(new_payload_bytes, payload_bytes) 286 | WHERE 287 | id = msg_id; 288 | $$ LANGUAGE SQL; 289 | 290 | -------------------------------------------------------------------------------- /sqlxmq_stress/migrations/20210921115907_clear.down.sql: -------------------------------------------------------------------------------- 1 | DROP FUNCTION mq_clear; 2 | DROP FUNCTION mq_clear_all; 3 | -------------------------------------------------------------------------------- /sqlxmq_stress/migrations/20210921115907_clear.up.sql: -------------------------------------------------------------------------------- 1 | -- Deletes all messages from a list of channel names. 2 | CREATE FUNCTION mq_clear(channel_names TEXT[]) 3 | RETURNS VOID AS $$ 4 | BEGIN 5 | WITH deleted_ids AS ( 6 | DELETE FROM mq_msgs WHERE channel_name = ANY(channel_names) RETURNING id 7 | ) 8 | DELETE FROM mq_payloads WHERE id IN (SELECT id FROM deleted_ids); 9 | END; 10 | $$ LANGUAGE plpgsql; 11 | 12 | -- Deletes all messages. 13 | CREATE FUNCTION mq_clear_all() 14 | RETURNS VOID AS $$ 15 | BEGIN 16 | WITH deleted_ids AS ( 17 | DELETE FROM mq_msgs RETURNING id 18 | ) 19 | DELETE FROM mq_payloads WHERE id IN (SELECT id FROM deleted_ids); 20 | END; 21 | $$ LANGUAGE plpgsql; 22 | -------------------------------------------------------------------------------- /sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.down.sql: -------------------------------------------------------------------------------- 1 | CREATE OR REPLACE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) 2 | RETURNS UUID AS $$ 3 | SELECT COALESCE( 4 | ( 5 | SELECT id FROM mq_msgs 6 | WHERE channel_name = from_channel_name 7 | AND channel_args = from_channel_args 8 | AND after_message_id IS NOT NULL 9 | AND id != uuid_nil() 10 | ORDER BY created_at DESC, id DESC 11 | LIMIT 1 12 | ), 13 | uuid_nil() 14 | ) 15 | $$ LANGUAGE SQL STABLE; 16 | -------------------------------------------------------------------------------- /sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.up.sql: -------------------------------------------------------------------------------- 1 | CREATE OR REPLACE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) 2 | RETURNS UUID AS $$ 3 | SELECT COALESCE( 4 | ( 5 | SELECT id FROM mq_msgs 6 | WHERE channel_name = from_channel_name 7 | AND channel_args = from_channel_args 8 | AND after_message_id IS NOT NULL 9 | AND id != uuid_nil() 10 | AND NOT EXISTS( 11 | SELECT * FROM mq_msgs AS mq_msgs2 12 | WHERE mq_msgs2.after_message_id = mq_msgs.id 13 | ) 14 | ORDER BY created_at DESC 15 | LIMIT 1 16 | ), 17 | uuid_nil() 18 | ) 19 | $$ LANGUAGE SQL STABLE; -------------------------------------------------------------------------------- /sqlxmq_stress/src/main.rs: -------------------------------------------------------------------------------- 1 | use std::env; 2 | use std::error::Error; 3 | use std::process::abort; 4 | use std::sync::RwLock; 5 | use std::time::{Duration, Instant}; 6 | 7 | use futures::channel::mpsc; 8 | use futures::StreamExt; 9 | use lazy_static::lazy_static; 10 | use serde::{Deserialize, Serialize}; 11 | use sqlx::{Pool, Postgres}; 12 | use sqlxmq::{job, CurrentJob, JobRegistry}; 13 | use tokio::task; 14 | 15 | lazy_static! { 16 | static ref INSTANT_EPOCH: Instant = Instant::now(); 17 | static ref CHANNEL: RwLock> = RwLock::new(mpsc::unbounded().0); 18 | } 19 | 20 | struct JobResult { 21 | duration: Duration, 22 | } 23 | 24 | #[derive(Serialize, Deserialize)] 25 | struct JobData { 26 | start_time: Duration, 27 | } 28 | 29 | // Arguments to the `#[job]` attribute allow setting default job options. 30 | #[job(channel_name = "foo")] 31 | async fn example_job( 32 | mut current_job: CurrentJob, 33 | ) -> Result<(), Box> { 34 | // Decode a JSON payload 35 | let data: JobData = current_job.json()?.unwrap(); 36 | 37 | // Mark the job as complete 38 | current_job.complete().await?; 39 | let end_time = INSTANT_EPOCH.elapsed(); 40 | 41 | CHANNEL.read().unwrap().unbounded_send(JobResult { 42 | duration: end_time - data.start_time, 43 | })?; 44 | 45 | Ok(()) 46 | } 47 | 48 | async fn start_job( 49 | pool: Pool, 50 | seed: usize, 51 | ) -> Result<(), Box> { 52 | let channel_name = if seed % 3 == 0 { "foo" } else { "bar" }; 53 | let channel_args = format!("{}", seed / 32); 54 | example_job 55 | .builder() 56 | // This is where we can override job configuration 57 | .set_channel_name(channel_name) 58 | .set_channel_args(&channel_args) 59 | .set_json(&JobData { 60 | start_time: INSTANT_EPOCH.elapsed(), 61 | })? 62 | .spawn(&pool) 63 | .await?; 64 | Ok(()) 65 | } 66 | 67 | async fn schedule_tasks(num_jobs: usize, interval: Duration, pool: Pool) { 68 | let mut stream = tokio::time::interval(interval); 69 | for i in 0..num_jobs { 70 | let pool = pool.clone(); 71 | task::spawn(async move { 72 | if let Err(e) = start_job(pool, i).await { 73 | eprintln!("Failed to start job: {:?}", e); 74 | abort(); 75 | } 76 | }); 77 | stream.tick().await; 78 | } 79 | } 80 | 81 | #[tokio::main] 82 | async fn main() -> Result<(), Box> { 83 | let _ = dotenvy::dotenv(); 84 | 85 | let pool = Pool::connect(&env::var("DATABASE_URL")?).await?; 86 | 87 | // Make sure the queues are empty 88 | sqlxmq::clear_all(&pool).await?; 89 | 90 | let registry = JobRegistry::new(&[example_job]); 91 | 92 | let _runner = registry 93 | .runner(&pool) 94 | .set_concurrency(50, 100) 95 | .run() 96 | .await?; 97 | let num_jobs = 10000; 98 | let interval = Duration::from_nanos(700_000); 99 | 100 | let (tx, rx) = mpsc::unbounded(); 101 | *CHANNEL.write()? = tx; 102 | 103 | let start_time = Instant::now(); 104 | task::spawn(schedule_tasks(num_jobs, interval, pool.clone())); 105 | 106 | let mut results: Vec<_> = rx.take(num_jobs).collect().await; 107 | let total_duration = start_time.elapsed(); 108 | 109 | assert_eq!(results.len(), num_jobs); 110 | 111 | results.sort_by_key(|r| r.duration); 112 | let (min, max, median, pct) = ( 113 | results[0].duration, 114 | results[num_jobs - 1].duration, 115 | results[num_jobs / 2].duration, 116 | results[(num_jobs * 19) / 20].duration, 117 | ); 118 | let throughput = num_jobs as f64 / total_duration.as_secs_f64(); 119 | 120 | println!("min: {}s", min.as_secs_f64()); 121 | println!("max: {}s", max.as_secs_f64()); 122 | println!("median: {}s", median.as_secs_f64()); 123 | println!("95th percentile: {}s", pct.as_secs_f64()); 124 | println!("throughput: {}/s", throughput); 125 | 126 | // The job runner will continue listening and running 127 | // jobs until `runner` is dropped. 128 | Ok(()) 129 | } 130 | -------------------------------------------------------------------------------- /src/hidden.rs: -------------------------------------------------------------------------------- 1 | use crate::{CurrentJob, JobBuilder, JobRegistry}; 2 | 3 | #[doc(hidden)] 4 | pub struct BuildFn(pub for<'a> fn(&'a mut JobBuilder<'static>) -> &'a mut JobBuilder<'static>); 5 | #[doc(hidden)] 6 | pub struct RunFn(pub fn(&JobRegistry, CurrentJob)); 7 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | #![deny(missing_docs, unsafe_code)] 2 | //! # sqlxmq 3 | //! 4 | //! A job queue built on `sqlx` and `PostgreSQL`. 5 | //! 6 | //! This library allows a CRUD application to run background jobs without complicating its 7 | //! deployment. The only runtime dependency is `PostgreSQL`, so this is ideal for applications 8 | //! already using a `PostgreSQL` database. 9 | //! 10 | //! Although using a SQL database as a job queue means compromising on latency of 11 | //! delivered jobs, there are several show-stopping issues present in ordinary job 12 | //! queues which are avoided altogether. 13 | //! 14 | //! With most other job queues, in-flight jobs are state that is not covered by normal 15 | //! database backups. Even if jobs _are_ backed up, there is no way to restore both 16 | //! a database and a job queue to a consistent point-in-time without manually 17 | //! resolving conflicts. 18 | //! 19 | //! By storing jobs in the database, existing backup procedures will store a perfectly 20 | //! consistent state of both in-flight jobs and persistent data. Additionally, jobs can 21 | //! be spawned and completed as part of other transactions, making it easy to write correct 22 | //! application code. 23 | //! 24 | //! Leveraging the power of `PostgreSQL`, this job queue offers several features not 25 | //! present in other job queues. 26 | //! 27 | //! # Features 28 | //! 29 | //! - **Send/receive multiple jobs at once.** 30 | //! 31 | //! This reduces the number of queries to the database. 32 | //! 33 | //! - **Send jobs to be executed at a future date and time.** 34 | //! 35 | //! Avoids the need for a separate scheduling system. 36 | //! 37 | //! - **Reliable delivery of jobs.** 38 | //! 39 | //! - **Automatic retries with exponential backoff.** 40 | //! 41 | //! Number of retries and initial backoff parameters are configurable. 42 | //! 43 | //! - **Transactional sending of jobs.** 44 | //! 45 | //! Avoids sending spurious jobs if a transaction is rolled back. 46 | //! 47 | //! - **Transactional completion of jobs.** 48 | //! 49 | //! If all side-effects of a job are updates to the database, this provides 50 | //! true exactly-once execution of jobs. 51 | //! 52 | //! - **Transactional check-pointing of jobs.** 53 | //! 54 | //! Long-running jobs can check-point their state to avoid having to restart 55 | //! from the beginning if there is a failure: the next retry can continue 56 | //! from the last check-point. 57 | //! 58 | //! - **Opt-in strictly ordered job delivery.** 59 | //! 60 | //! Jobs within the same channel will be processed strictly in-order 61 | //! if this option is enabled for the job. 62 | //! 63 | //! - **Fair job delivery.** 64 | //! 65 | //! A channel with a lot of jobs ready to run will not starve a channel with fewer 66 | //! jobs. 67 | //! 68 | //! - **Opt-in two-phase commit.** 69 | //! 70 | //! This is particularly useful on an ordered channel where a position can be "reserved" 71 | //! in the job order, but not committed until later. 72 | //! 73 | //! - **JSON and/or binary payloads.** 74 | //! 75 | //! Jobs can use whichever is most convenient. 76 | //! 77 | //! - **Automatic keep-alive of jobs.** 78 | //! 79 | //! Long-running jobs will automatically be "kept alive" to prevent them being 80 | //! retried whilst they're still ongoing. 81 | //! 82 | //! - **Concurrency limits.** 83 | //! 84 | //! Specify the minimum and maximum number of concurrent jobs each runner should 85 | //! handle. 86 | //! 87 | //! - **Built-in job registry via an attribute macro.** 88 | //! 89 | //! Jobs can be easily registered with a runner, and default configuration specified 90 | //! on a per-job basis. 91 | //! 92 | //! - **Implicit channels.** 93 | //! 94 | //! Channels are implicitly created and destroyed when jobs are sent and processed, 95 | //! so no setup is required. 96 | //! 97 | //! - **Channel groups.** 98 | //! 99 | //! Easily subscribe to multiple channels at once, thanks to the separation of 100 | //! channel name and channel arguments. 101 | //! 102 | //! - **NOTIFY-based polling.** 103 | //! 104 | //! This saves resources when few jobs are being processed. 105 | //! 106 | //! # Getting started 107 | //! 108 | //! ## Database schema 109 | //! 110 | //! This crate expects certain database tables and stored procedures to exist. 111 | //! You can copy the migration files from this crate into your own migrations 112 | //! folder. 113 | //! 114 | //! All database items created by this crate are prefixed with `mq`, so as not 115 | //! to conflict with your own schema. 116 | //! 117 | //! ## Defining jobs 118 | //! 119 | //! The first step is to define a function to be run on the job queue. 120 | //! 121 | //! ```rust 122 | //! use std::error::Error; 123 | //! 124 | //! use sqlxmq::{job, CurrentJob}; 125 | //! 126 | //! // Arguments to the `#[job]` attribute allow setting default job options. 127 | //! #[job(channel_name = "foo")] 128 | //! async fn example_job( 129 | //! // The first argument should always be the current job. 130 | //! mut current_job: CurrentJob, 131 | //! // Additional arguments are optional, but can be used to access context 132 | //! // provided via [`JobRegistry::set_context`]. 133 | //! message: &'static str, 134 | //! ) -> Result<(), Box> { 135 | //! // Decode a JSON payload 136 | //! let who: Option = current_job.json()?; 137 | //! 138 | //! // Do some work 139 | //! println!("{}, {}!", message, who.as_deref().unwrap_or("world")); 140 | //! 141 | //! // Mark the job as complete 142 | //! current_job.complete().await?; 143 | //! 144 | //! Ok(()) 145 | //! } 146 | //! ``` 147 | //! 148 | //! ## Listening for jobs 149 | //! 150 | //! Next we need to create a job runner: this is what listens for new jobs 151 | //! and executes them. 152 | //! 153 | //! ```rust,no_run 154 | //! use std::error::Error; 155 | //! 156 | //! use sqlxmq::JobRegistry; 157 | //! 158 | //! # use sqlxmq::{job, CurrentJob}; 159 | //! # 160 | //! # #[job] 161 | //! # async fn example_job( 162 | //! # current_job: CurrentJob, 163 | //! # ) -> Result<(), Box> { Ok(()) } 164 | //! # 165 | //! # async fn connect_to_db() -> sqlx::Result> { 166 | //! # unimplemented!() 167 | //! # } 168 | //! 169 | //! #[tokio::main] 170 | //! async fn main() -> Result<(), Box> { 171 | //! // You'll need to provide a Postgres connection pool. 172 | //! let pool = connect_to_db().await?; 173 | //! 174 | //! // Construct a job registry from our single job. 175 | //! let mut registry = JobRegistry::new(&[example_job]); 176 | //! // Here is where you can configure the registry 177 | //! // registry.set_error_handler(...) 178 | //! 179 | //! // And add context 180 | //! registry.set_context("Hello"); 181 | //! 182 | //! let runner = registry 183 | //! // Create a job runner using the connection pool. 184 | //! .runner(&pool) 185 | //! // Here is where you can configure the job runner 186 | //! // Aim to keep 10-20 jobs running at a time. 187 | //! .set_concurrency(10, 20) 188 | //! // Start the job runner in the background. 189 | //! .run() 190 | //! .await?; 191 | //! 192 | //! // The job runner will continue listening and running 193 | //! // jobs until `runner` is dropped. 194 | //! Ok(()) 195 | //! } 196 | //! ``` 197 | //! 198 | //! ## Spawning a job 199 | //! 200 | //! The final step is to actually run a job. 201 | //! 202 | //! ```rust 203 | //! # use std::error::Error; 204 | //! # use sqlxmq::{job, CurrentJob}; 205 | //! # 206 | //! # #[job] 207 | //! # async fn example_job( 208 | //! # current_job: CurrentJob, 209 | //! # ) -> Result<(), Box> { Ok(()) } 210 | //! # 211 | //! # async fn example( 212 | //! # pool: sqlx::Pool 213 | //! # ) -> Result<(), Box> { 214 | //! example_job.builder() 215 | //! // This is where we can override job configuration 216 | //! .set_channel_name("bar") 217 | //! .set_json("John")? 218 | //! .spawn(&pool) 219 | //! .await?; 220 | //! # Ok(()) 221 | //! # } 222 | //! ``` 223 | 224 | #[doc(hidden)] 225 | pub mod hidden; 226 | mod registry; 227 | mod runner; 228 | mod spawn; 229 | mod utils; 230 | 231 | pub use registry::*; 232 | pub use runner::*; 233 | pub use spawn::*; 234 | pub use sqlxmq_macros::job; 235 | pub use utils::OwnedHandle; 236 | 237 | /// Helper function to determine if a particular error condition is retryable. 238 | /// 239 | /// For best results, database operations should be automatically retried if one 240 | /// of these errors is returned. 241 | pub fn should_retry(error: &sqlx::Error) -> bool { 242 | if let Some(db_error) = error.as_database_error() { 243 | // It's more readable as a match 244 | #[allow(clippy::match_like_matches_macro)] 245 | match (db_error.code().as_deref(), db_error.constraint()) { 246 | // Foreign key constraint violation on ordered channel 247 | (Some("23503"), Some("mq_msgs_after_message_id_fkey")) => true, 248 | // Unique constraint violation on ordered channel 249 | (Some("23505"), Some("mq_msgs_channel_name_channel_args_after_message_id_idx")) => true, 250 | // Serialization failure 251 | (Some("40001"), _) => true, 252 | // Deadlock detected 253 | (Some("40P01"), _) => true, 254 | // Other 255 | _ => false, 256 | } 257 | } else { 258 | false 259 | } 260 | } 261 | 262 | #[cfg(test)] 263 | mod tests { 264 | use super::*; 265 | use crate as sqlxmq; 266 | 267 | use std::env; 268 | use std::error::Error; 269 | use std::future::Future; 270 | use std::ops::Deref; 271 | use std::sync::atomic::{AtomicUsize, Ordering}; 272 | use std::sync::{Arc, Once}; 273 | use std::time::Duration; 274 | 275 | use futures::channel::mpsc; 276 | use futures::StreamExt; 277 | use sqlx::{Pool, Postgres}; 278 | use tokio::sync::{Mutex, MutexGuard}; 279 | use tokio::task; 280 | 281 | // field 0 is never read, but its drop is important 282 | #[allow(dead_code)] 283 | struct TestGuard(MutexGuard<'static, ()>, T); 284 | 285 | impl Deref for TestGuard { 286 | type Target = T; 287 | 288 | fn deref(&self) -> &T { 289 | &self.1 290 | } 291 | } 292 | 293 | async fn test_pool() -> TestGuard> { 294 | static INIT_LOGGER: Once = Once::new(); 295 | static TEST_MUTEX: Mutex<()> = Mutex::const_new(()); 296 | 297 | let guard = TEST_MUTEX.lock().await; 298 | 299 | let _ = dotenvy::dotenv(); 300 | 301 | INIT_LOGGER.call_once(pretty_env_logger::init); 302 | 303 | let pool = Pool::connect(&env::var("DATABASE_URL").unwrap()) 304 | .await 305 | .unwrap(); 306 | 307 | sqlx::query("TRUNCATE TABLE mq_payloads") 308 | .execute(&pool) 309 | .await 310 | .unwrap(); 311 | sqlx::query("DELETE FROM mq_msgs WHERE id != uuid_nil()") 312 | .execute(&pool) 313 | .await 314 | .unwrap(); 315 | 316 | TestGuard(guard, pool) 317 | } 318 | 319 | async fn test_job_runner( 320 | pool: &Pool, 321 | f: impl (Fn(CurrentJob) -> F) + Send + Sync + 'static, 322 | ) -> (JobRunnerHandle, Arc) 323 | where 324 | F::Output: Send + 'static, 325 | { 326 | let counter = Arc::new(AtomicUsize::new(0)); 327 | let counter2 = counter.clone(); 328 | let runner = JobRunnerOptions::new(pool, move |job| { 329 | counter2.fetch_add(1, Ordering::SeqCst); 330 | task::spawn(f(job)); 331 | }) 332 | .run() 333 | .await 334 | .unwrap(); 335 | (runner, counter) 336 | } 337 | 338 | fn job_proto<'a, 'b>(builder: &'a mut JobBuilder<'b>) -> &'a mut JobBuilder<'b> { 339 | builder.set_channel_name("bar") 340 | } 341 | 342 | #[job(channel_name = "foo", ordered, retries = 3, backoff_secs = 2.0)] 343 | async fn example_job1( 344 | mut current_job: CurrentJob, 345 | ) -> Result<(), Box> { 346 | current_job.complete().await?; 347 | Ok(()) 348 | } 349 | 350 | #[job(proto(job_proto))] 351 | async fn example_job2( 352 | mut current_job: CurrentJob, 353 | ) -> Result<(), Box> { 354 | current_job.complete().await?; 355 | Ok(()) 356 | } 357 | 358 | #[job] 359 | async fn example_job_with_ctx( 360 | mut current_job: CurrentJob, 361 | ctx1: i32, 362 | ctx2: &'static str, 363 | ) -> Result<(), Box> { 364 | assert_eq!(ctx1, 42); 365 | assert_eq!(ctx2, "Hello, world!"); 366 | current_job.complete().await?; 367 | Ok(()) 368 | } 369 | 370 | async fn named_job_runner(pool: &Pool) -> JobRunnerHandle { 371 | let mut registry = JobRegistry::new(&[example_job1, example_job2, example_job_with_ctx]); 372 | registry.set_context(42).set_context("Hello, world!"); 373 | registry.runner(pool).run().await.unwrap() 374 | } 375 | 376 | fn is_ci() -> bool { 377 | std::env::var("CI").ok().is_some() 378 | } 379 | 380 | fn default_pause() -> u64 { 381 | if is_ci() { 382 | 1000 383 | } else { 384 | 200 385 | } 386 | } 387 | 388 | async fn pause() { 389 | pause_ms(default_pause()).await; 390 | } 391 | 392 | async fn pause_ms(ms: u64) { 393 | tokio::time::sleep(Duration::from_millis(ms)).await; 394 | } 395 | 396 | #[tokio::test] 397 | async fn it_can_spawn_job() { 398 | { 399 | let pool = &*test_pool().await; 400 | let (_runner, counter) = 401 | test_job_runner(pool, |mut job| async move { job.complete().await }).await; 402 | 403 | assert_eq!(counter.load(Ordering::SeqCst), 0); 404 | JobBuilder::new("foo").spawn(pool).await.unwrap(); 405 | pause().await; 406 | assert_eq!(counter.load(Ordering::SeqCst), 1); 407 | } 408 | pause().await; 409 | } 410 | 411 | #[tokio::test] 412 | async fn it_can_clear_jobs() { 413 | { 414 | let pool = &*test_pool().await; 415 | JobBuilder::new("foo") 416 | .set_channel_name("foo") 417 | .spawn(pool) 418 | .await 419 | .unwrap(); 420 | JobBuilder::new("foo") 421 | .set_channel_name("foo") 422 | .spawn(pool) 423 | .await 424 | .unwrap(); 425 | JobBuilder::new("foo") 426 | .set_channel_name("bar") 427 | .spawn(pool) 428 | .await 429 | .unwrap(); 430 | JobBuilder::new("foo") 431 | .set_channel_name("bar") 432 | .spawn(pool) 433 | .await 434 | .unwrap(); 435 | JobBuilder::new("foo") 436 | .set_channel_name("baz") 437 | .spawn(pool) 438 | .await 439 | .unwrap(); 440 | JobBuilder::new("foo") 441 | .set_channel_name("baz") 442 | .spawn(pool) 443 | .await 444 | .unwrap(); 445 | 446 | sqlxmq::clear(pool, &["foo", "baz"]).await.unwrap(); 447 | 448 | let (_runner, counter) = 449 | test_job_runner(pool, |mut job| async move { job.complete().await }).await; 450 | 451 | pause().await; 452 | assert_eq!(counter.load(Ordering::SeqCst), 2); 453 | } 454 | pause().await; 455 | } 456 | 457 | #[tokio::test] 458 | async fn it_runs_jobs_in_order() { 459 | { 460 | let pool = &*test_pool().await; 461 | let (tx, mut rx) = mpsc::unbounded(); 462 | 463 | let (_runner, counter) = test_job_runner(pool, move |job| { 464 | let tx = tx.clone(); 465 | async move { 466 | tx.unbounded_send(job).unwrap(); 467 | } 468 | }) 469 | .await; 470 | 471 | assert_eq!(counter.load(Ordering::SeqCst), 0); 472 | JobBuilder::new("foo") 473 | .set_ordered(true) 474 | .spawn(pool) 475 | .await 476 | .unwrap(); 477 | JobBuilder::new("bar") 478 | .set_ordered(true) 479 | .spawn(pool) 480 | .await 481 | .unwrap(); 482 | 483 | pause().await; 484 | assert_eq!(counter.load(Ordering::SeqCst), 1); 485 | 486 | let mut job = rx.next().await.unwrap(); 487 | job.complete().await.unwrap(); 488 | 489 | pause().await; 490 | assert_eq!(counter.load(Ordering::SeqCst), 2); 491 | } 492 | pause().await; 493 | } 494 | 495 | #[tokio::test] 496 | async fn it_runs_jobs_in_parallel() { 497 | { 498 | let pool = &*test_pool().await; 499 | let (tx, mut rx) = mpsc::unbounded(); 500 | 501 | let (_runner, counter) = test_job_runner(pool, move |job| { 502 | let tx = tx.clone(); 503 | async move { 504 | tx.unbounded_send(job).unwrap(); 505 | } 506 | }) 507 | .await; 508 | 509 | assert_eq!(counter.load(Ordering::SeqCst), 0); 510 | JobBuilder::new("foo").spawn(pool).await.unwrap(); 511 | JobBuilder::new("bar").spawn(pool).await.unwrap(); 512 | 513 | pause().await; 514 | assert_eq!(counter.load(Ordering::SeqCst), 2); 515 | 516 | for _ in 0..2 { 517 | let mut job = rx.next().await.unwrap(); 518 | job.complete().await.unwrap(); 519 | } 520 | } 521 | pause().await; 522 | } 523 | 524 | #[tokio::test] 525 | async fn it_retries_failed_jobs() { 526 | { 527 | let pool = &*test_pool().await; 528 | let (_runner, counter) = test_job_runner(pool, move |_| async {}).await; 529 | 530 | let backoff = default_pause() + 300; 531 | 532 | assert_eq!(counter.load(Ordering::SeqCst), 0); 533 | JobBuilder::new("foo") 534 | .set_retry_backoff(Duration::from_millis(backoff)) 535 | .set_retries(2) 536 | .spawn(pool) 537 | .await 538 | .unwrap(); 539 | 540 | // First attempt 541 | pause().await; 542 | assert_eq!(counter.load(Ordering::SeqCst), 1); 543 | 544 | // Second attempt 545 | pause_ms(backoff).await; 546 | pause().await; 547 | assert_eq!(counter.load(Ordering::SeqCst), 2); 548 | 549 | // Third attempt 550 | pause_ms(backoff * 2).await; 551 | pause().await; 552 | assert_eq!(counter.load(Ordering::SeqCst), 3); 553 | 554 | // No more attempts 555 | pause_ms(backoff * 5).await; 556 | assert_eq!(counter.load(Ordering::SeqCst), 3); 557 | } 558 | pause().await; 559 | } 560 | 561 | #[tokio::test] 562 | async fn it_can_checkpoint_jobs() { 563 | { 564 | let pool = &*test_pool().await; 565 | let (_runner, counter) = test_job_runner(pool, move |mut current_job| async move { 566 | let state: bool = current_job.json().unwrap().unwrap(); 567 | if state { 568 | current_job.complete().await.unwrap(); 569 | } else { 570 | current_job 571 | .checkpoint(Checkpoint::new().set_json(&true).unwrap()) 572 | .await 573 | .unwrap(); 574 | } 575 | }) 576 | .await; 577 | 578 | let backoff = default_pause(); 579 | 580 | assert_eq!(counter.load(Ordering::SeqCst), 0); 581 | JobBuilder::new("foo") 582 | .set_retry_backoff(Duration::from_millis(backoff)) 583 | .set_retries(5) 584 | .set_json(&false) 585 | .unwrap() 586 | .spawn(pool) 587 | .await 588 | .unwrap(); 589 | 590 | // First attempt 591 | pause().await; 592 | assert_eq!(counter.load(Ordering::SeqCst), 1); 593 | 594 | // Second attempt 595 | pause_ms(backoff).await; 596 | assert_eq!(counter.load(Ordering::SeqCst), 2); 597 | 598 | // No more attempts 599 | pause_ms(backoff * 3).await; 600 | assert_eq!(counter.load(Ordering::SeqCst), 2); 601 | } 602 | pause().await; 603 | } 604 | 605 | #[tokio::test] 606 | async fn it_can_use_registry() { 607 | { 608 | let pool = &*test_pool().await; 609 | let _runner = named_job_runner(pool).await; 610 | 611 | example_job1.builder().spawn(pool).await.unwrap(); 612 | example_job2.builder().spawn(pool).await.unwrap(); 613 | example_job_with_ctx.builder().spawn(pool).await.unwrap(); 614 | pause().await; 615 | } 616 | pause().await; 617 | } 618 | } 619 | -------------------------------------------------------------------------------- /src/registry.rs: -------------------------------------------------------------------------------- 1 | use std::any::type_name; 2 | use std::collections::HashMap; 3 | use std::error::Error; 4 | use std::fmt::Display; 5 | use std::future::Future; 6 | use std::sync::Arc; 7 | use std::time::Instant; 8 | 9 | use anymap2::any::CloneAnySendSync; 10 | use anymap2::Map; 11 | use sqlx::{Pool, Postgres}; 12 | use uuid::Uuid; 13 | 14 | use crate::hidden::{BuildFn, RunFn}; 15 | use crate::utils::Opaque; 16 | use crate::{JobBuilder, JobRunnerOptions}; 17 | 18 | type BoxedError = Box; 19 | 20 | /// Stores a mapping from job name to job. Can be used to construct 21 | /// a job runner. 22 | pub struct JobRegistry { 23 | #[allow(clippy::type_complexity)] 24 | error_handler: Arc, 25 | job_map: HashMap<&'static str, &'static NamedJob>, 26 | context: Map, 27 | } 28 | 29 | /// Error returned when a job is received whose name is not in the registry. 30 | #[derive(Debug)] 31 | pub struct UnknownJobError; 32 | 33 | impl Error for UnknownJobError {} 34 | impl Display for UnknownJobError { 35 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 36 | f.write_str("Unknown job") 37 | } 38 | } 39 | 40 | impl JobRegistry { 41 | /// Construct a new job registry from the provided job list. 42 | pub fn new(jobs: &[&'static NamedJob]) -> Self { 43 | let mut job_map = HashMap::new(); 44 | for &job in jobs { 45 | if job_map.insert(job.name(), job).is_some() { 46 | panic!("Duplicate job registered: {}", job.name()); 47 | } 48 | } 49 | Self { 50 | error_handler: Arc::new(Self::default_error_handler), 51 | job_map, 52 | context: Map::new(), 53 | } 54 | } 55 | 56 | /// Set a function to be called whenever a job returns an error. 57 | pub fn set_error_handler( 58 | &mut self, 59 | error_handler: impl Fn(&str, BoxedError) + Send + Sync + 'static, 60 | ) -> &mut Self { 61 | self.error_handler = Arc::new(error_handler); 62 | self 63 | } 64 | 65 | /// Provide context for the jobs. 66 | pub fn set_context(&mut self, context: C) -> &mut Self { 67 | self.context.insert(context); 68 | self 69 | } 70 | 71 | /// Access job context. Will panic if context with this type has not been provided. 72 | pub fn context(&self) -> C { 73 | if let Some(c) = self.context.get::() { 74 | c.clone() 75 | } else { 76 | panic!( 77 | "No context of type `{}` has been provided.", 78 | type_name::() 79 | ); 80 | } 81 | } 82 | 83 | /// Look-up a job by name. 84 | pub fn resolve_job(&self, name: &str) -> Option<&'static NamedJob> { 85 | self.job_map.get(name).copied() 86 | } 87 | 88 | /// The default error handler implementation, which simply logs the error. 89 | pub fn default_error_handler(name: &str, error: BoxedError) { 90 | log::error!("Job `{}` failed: {:?}", name, error); 91 | } 92 | 93 | #[doc(hidden)] 94 | pub fn spawn_internal>>( 95 | &self, 96 | name: &'static str, 97 | f: impl Future> + Send + 'static, 98 | ) { 99 | let error_handler = self.error_handler.clone(); 100 | tokio::spawn(async move { 101 | let start_time = Instant::now(); 102 | log::info!("Job `{}` started.", name); 103 | if let Err(e) = f.await { 104 | error_handler(name, e.into()); 105 | } else { 106 | log::info!( 107 | "Job `{}` completed in {}s.", 108 | name, 109 | start_time.elapsed().as_secs_f64() 110 | ); 111 | } 112 | }); 113 | } 114 | 115 | /// Construct a job runner from this registry and the provided connection 116 | /// pool. 117 | pub fn runner(self, pool: &Pool) -> JobRunnerOptions { 118 | JobRunnerOptions::new(pool, move |current_job| { 119 | if let Some(job) = self.resolve_job(current_job.name()) { 120 | (job.run_fn.0 .0)(&self, current_job); 121 | } else { 122 | (self.error_handler)(current_job.name(), Box::new(UnknownJobError)) 123 | } 124 | }) 125 | } 126 | } 127 | 128 | /// Type for a named job. Functions annotated with `#[job]` are 129 | /// transformed into static variables whose type is `&'static NamedJob`. 130 | #[derive(Debug)] 131 | pub struct NamedJob { 132 | name: &'static str, 133 | build_fn: Opaque, 134 | run_fn: Opaque, 135 | } 136 | 137 | impl NamedJob { 138 | #[doc(hidden)] 139 | pub const fn new_internal(name: &'static str, build_fn: BuildFn, run_fn: RunFn) -> Self { 140 | Self { 141 | name, 142 | build_fn: Opaque(build_fn), 143 | run_fn: Opaque(run_fn), 144 | } 145 | } 146 | /// Initialize a job builder with the name and defaults of this job. 147 | pub fn builder(&self) -> JobBuilder<'static> { 148 | let mut builder = JobBuilder::new(self.name); 149 | (self.build_fn.0 .0)(&mut builder); 150 | builder 151 | } 152 | /// Initialize a job builder with the name and defaults of this job, 153 | /// using the provided job ID. 154 | pub fn builder_with_id(&self, id: Uuid) -> JobBuilder<'static> { 155 | let mut builder = JobBuilder::new_with_id(id, self.name); 156 | (self.build_fn.0 .0)(&mut builder); 157 | builder 158 | } 159 | 160 | /// Returns the name of this job. 161 | pub const fn name(&self) -> &'static str { 162 | self.name 163 | } 164 | } 165 | -------------------------------------------------------------------------------- /src/runner.rs: -------------------------------------------------------------------------------- 1 | use std::borrow::Cow; 2 | use std::fmt::Debug; 3 | use std::sync::atomic::{AtomicUsize, Ordering}; 4 | use std::sync::Arc; 5 | use std::time::{Duration, Instant}; 6 | 7 | use serde::{Deserialize, Serialize}; 8 | use sqlx::postgres::types::PgInterval; 9 | use sqlx::postgres::PgListener; 10 | use sqlx::{Pool, Postgres}; 11 | use tokio::sync::{oneshot, Notify}; 12 | use tokio::task; 13 | use uuid::Uuid; 14 | 15 | use crate::utils::{Opaque, OwnedHandle}; 16 | 17 | /// Type used to build a job runner. 18 | #[derive(Debug, Clone)] 19 | pub struct JobRunnerOptions { 20 | min_concurrency: usize, 21 | max_concurrency: usize, 22 | channel_names: Option>, 23 | dispatch: Opaque>, 24 | pool: Pool, 25 | keep_alive: bool, 26 | } 27 | 28 | #[derive(Debug)] 29 | struct JobRunner { 30 | options: JobRunnerOptions, 31 | running_jobs: AtomicUsize, 32 | notify: Notify, 33 | } 34 | 35 | /// Job runner handle 36 | pub struct JobRunnerHandle { 37 | runner: Arc, 38 | handle: Option, 39 | } 40 | 41 | /// Type used to checkpoint a running job. 42 | #[derive(Debug, Clone, Default)] 43 | pub struct Checkpoint<'a> { 44 | duration: Duration, 45 | extra_retries: usize, 46 | payload_json: Option>, 47 | payload_bytes: Option<&'a [u8]>, 48 | } 49 | 50 | impl<'a> Checkpoint<'a> { 51 | /// Construct a new checkpoint which also keeps the job alive 52 | /// for the specified interval. 53 | pub fn new_keep_alive(duration: Duration) -> Self { 54 | Self { 55 | duration, 56 | extra_retries: 0, 57 | payload_json: None, 58 | payload_bytes: None, 59 | } 60 | } 61 | /// Construct a new checkpoint. 62 | pub fn new() -> Self { 63 | Self::default() 64 | } 65 | /// Add extra retries to the current job. 66 | pub fn set_extra_retries(&mut self, extra_retries: usize) -> &mut Self { 67 | self.extra_retries = extra_retries; 68 | self 69 | } 70 | /// Specify a new raw JSON payload. 71 | pub fn set_raw_json(&mut self, raw_json: &'a str) -> &mut Self { 72 | self.payload_json = Some(Cow::Borrowed(raw_json)); 73 | self 74 | } 75 | /// Specify a new raw binary payload. 76 | pub fn set_raw_bytes(&mut self, raw_bytes: &'a [u8]) -> &mut Self { 77 | self.payload_bytes = Some(raw_bytes); 78 | self 79 | } 80 | /// Specify a new JSON payload. 81 | pub fn set_json(&mut self, value: &T) -> Result<&mut Self, serde_json::Error> { 82 | let value = serde_json::to_string(value)?; 83 | self.payload_json = Some(Cow::Owned(value)); 84 | Ok(self) 85 | } 86 | async fn execute<'b, E: sqlx::Executor<'b, Database = Postgres>>( 87 | &self, 88 | job_id: Uuid, 89 | executor: E, 90 | ) -> Result<(), sqlx::Error> { 91 | sqlx::query("SELECT mq_checkpoint($1, $2, $3, $4, $5)") 92 | .bind(job_id) 93 | .bind(self.duration) 94 | .bind(self.payload_json.as_deref()) 95 | .bind(self.payload_bytes) 96 | .bind(self.extra_retries as i32) 97 | .execute(executor) 98 | .await?; 99 | Ok(()) 100 | } 101 | } 102 | 103 | /// Handle to the currently executing job. 104 | /// When dropped, the job is assumed to no longer be running. 105 | /// To prevent the job being retried, it must be explicitly completed using 106 | /// one of the `.complete_` methods. 107 | #[derive(Debug)] 108 | pub struct CurrentJob { 109 | id: Uuid, 110 | name: String, 111 | payload_json: Option, 112 | payload_bytes: Option>, 113 | job_runner: Arc, 114 | keep_alive: Option, 115 | } 116 | 117 | impl CurrentJob { 118 | /// Returns the database pool used to receive this job. 119 | pub fn pool(&self) -> &Pool { 120 | &self.job_runner.options.pool 121 | } 122 | async fn delete( 123 | &self, 124 | executor: impl sqlx::Executor<'_, Database = Postgres>, 125 | ) -> Result<(), sqlx::Error> { 126 | sqlx::query("SELECT mq_delete(ARRAY[$1])") 127 | .bind(self.id) 128 | .execute(executor) 129 | .await?; 130 | Ok(()) 131 | } 132 | 133 | async fn stop_keep_alive(&mut self) { 134 | if let Some(keep_alive) = self.keep_alive.take() { 135 | keep_alive.stop().await; 136 | } 137 | } 138 | 139 | /// Complete this job and commit the provided transaction at the same time. 140 | /// If the transaction cannot be committed, the job will not be completed. 141 | pub async fn complete_with_transaction( 142 | &mut self, 143 | mut tx: sqlx::Transaction<'_, Postgres>, 144 | ) -> Result<(), sqlx::Error> { 145 | self.delete(&mut *tx).await?; 146 | tx.commit().await?; 147 | self.stop_keep_alive().await; 148 | Ok(()) 149 | } 150 | /// Complete this job. 151 | pub async fn complete(&mut self) -> Result<(), sqlx::Error> { 152 | self.delete(self.pool()).await?; 153 | self.stop_keep_alive().await; 154 | Ok(()) 155 | } 156 | /// Checkpoint this job and commit the provided transaction at the same time. 157 | /// If the transaction cannot be committed, the job will not be checkpointed. 158 | /// Checkpointing allows the job payload to be replaced for the next retry. 159 | pub async fn checkpoint_with_transaction( 160 | &mut self, 161 | mut tx: sqlx::Transaction<'_, Postgres>, 162 | checkpoint: &Checkpoint<'_>, 163 | ) -> Result<(), sqlx::Error> { 164 | checkpoint.execute(self.id, &mut *tx).await?; 165 | tx.commit().await?; 166 | Ok(()) 167 | } 168 | /// Checkpointing allows the job payload to be replaced for the next retry. 169 | pub async fn checkpoint(&mut self, checkpoint: &Checkpoint<'_>) -> Result<(), sqlx::Error> { 170 | checkpoint.execute(self.id, self.pool()).await?; 171 | Ok(()) 172 | } 173 | /// Prevent this job from being retried for the specified interval. 174 | pub async fn keep_alive(&mut self, duration: Duration) -> Result<(), sqlx::Error> { 175 | sqlx::query("SELECT mq_keep_alive(ARRAY[$1], $2)") 176 | .bind(self.id) 177 | .bind(duration) 178 | .execute(self.pool()) 179 | .await?; 180 | Ok(()) 181 | } 182 | /// Returns the ID of this job. 183 | pub fn id(&self) -> Uuid { 184 | self.id 185 | } 186 | /// Returns the name of this job. 187 | pub fn name(&self) -> &str { 188 | &self.name 189 | } 190 | /// Extracts the JSON payload belonging to this job (if present). 191 | pub fn json<'a, T: Deserialize<'a>>(&'a self) -> Result, serde_json::Error> { 192 | if let Some(payload_json) = &self.payload_json { 193 | serde_json::from_str(payload_json).map(Some) 194 | } else { 195 | Ok(None) 196 | } 197 | } 198 | /// Returns the raw JSON payload for this job. 199 | pub fn raw_json(&self) -> Option<&str> { 200 | self.payload_json.as_deref() 201 | } 202 | /// Returns the raw binary payload for this job. 203 | pub fn raw_bytes(&self) -> Option<&[u8]> { 204 | self.payload_bytes.as_deref() 205 | } 206 | } 207 | 208 | impl Drop for CurrentJob { 209 | fn drop(&mut self) { 210 | if self.job_runner.running_jobs.fetch_sub(1, Ordering::SeqCst) 211 | == self.job_runner.options.min_concurrency 212 | { 213 | self.job_runner.notify.notify_one(); 214 | } 215 | } 216 | } 217 | 218 | impl JobRunnerOptions { 219 | /// Begin constructing a new job runner using the specified connection pool, 220 | /// and the provided execution function. 221 | pub fn new(pool: &Pool, f: F) -> Self { 222 | Self { 223 | min_concurrency: 16, 224 | max_concurrency: 32, 225 | channel_names: None, 226 | keep_alive: true, 227 | dispatch: Opaque(Arc::new(f)), 228 | pool: pool.clone(), 229 | } 230 | } 231 | /// Set the concurrency limits for this job runner. When the number of active 232 | /// jobs falls below the minimum, the runner will poll for more, up to the maximum. 233 | /// 234 | /// The difference between the min and max will dictate the maximum batch size which 235 | /// can be received: larger batch sizes are more efficient. 236 | pub fn set_concurrency(&mut self, min_concurrency: usize, max_concurrency: usize) -> &mut Self { 237 | self.min_concurrency = min_concurrency; 238 | self.max_concurrency = max_concurrency; 239 | self 240 | } 241 | /// Set the channel names which this job runner will subscribe to. If unspecified, 242 | /// the job runner will subscribe to all channels. 243 | pub fn set_channel_names<'a>(&'a mut self, channel_names: &[&str]) -> &'a mut Self { 244 | self.channel_names = Some( 245 | channel_names 246 | .iter() 247 | .copied() 248 | .map(ToOwned::to_owned) 249 | .collect(), 250 | ); 251 | self 252 | } 253 | /// Choose whether to automatically keep jobs alive whilst they're still 254 | /// running. Defaults to `true`. 255 | pub fn set_keep_alive(&mut self, keep_alive: bool) -> &mut Self { 256 | self.keep_alive = keep_alive; 257 | self 258 | } 259 | 260 | /// Start the job runner in the background. The job runner will stop when the 261 | /// returned handle is dropped. 262 | pub async fn run(&self) -> Result { 263 | let options = self.clone(); 264 | let job_runner = Arc::new(JobRunner { 265 | options, 266 | running_jobs: AtomicUsize::new(0), 267 | notify: Notify::new(), 268 | }); 269 | let listener_task = start_listener(job_runner.clone()).await?; 270 | let handle = OwnedHandle::new(task::spawn(main_loop(job_runner.clone(), listener_task))); 271 | Ok(JobRunnerHandle { 272 | runner: job_runner, 273 | handle: Some(handle), 274 | }) 275 | } 276 | 277 | /// Run a single job and then return. Intended for use by tests. The job should 278 | /// have been spawned normally and be ready to run. 279 | pub async fn test_one(&self) -> Result<(), sqlx::Error> { 280 | let options = self.clone(); 281 | let job_runner = Arc::new(JobRunner { 282 | options, 283 | running_jobs: AtomicUsize::new(0), 284 | notify: Notify::new(), 285 | }); 286 | 287 | log::info!("Polling for single message"); 288 | let mut messages = sqlx::query_as::<_, PolledMessage>("SELECT * FROM mq_poll($1, 1)") 289 | .bind(&self.channel_names) 290 | .fetch_all(&self.pool) 291 | .await?; 292 | 293 | assert_eq!(messages.len(), 1, "Expected one message to be ready"); 294 | let msg = messages.pop().unwrap(); 295 | 296 | if let PolledMessage { 297 | id: Some(id), 298 | is_committed: Some(true), 299 | name: Some(name), 300 | payload_json, 301 | payload_bytes, 302 | .. 303 | } = msg 304 | { 305 | let (tx, rx) = oneshot::channel::<()>(); 306 | let keep_alive = Some(OwnedHandle::new(task::spawn(async move { 307 | let _tx = tx; 308 | loop { 309 | tokio::time::sleep(Duration::from_secs(1)).await; 310 | } 311 | }))); 312 | let current_job = CurrentJob { 313 | id, 314 | name, 315 | payload_json, 316 | payload_bytes, 317 | job_runner: job_runner.clone(), 318 | keep_alive, 319 | }; 320 | job_runner.running_jobs.fetch_add(1, Ordering::SeqCst); 321 | (self.dispatch)(current_job); 322 | 323 | // Wait for job to complete 324 | let _ = rx.await; 325 | } 326 | Ok(()) 327 | } 328 | } 329 | 330 | impl JobRunnerHandle { 331 | /// Return the number of still running jobs 332 | pub fn num_running_jobs(&self) -> usize { 333 | self.runner.running_jobs.load(Ordering::Relaxed) 334 | } 335 | 336 | /// Wait for the jobs to finish, but not more than `timeout` 337 | pub async fn wait_jobs_finish(&self, timeout: Duration) { 338 | let start = Instant::now(); 339 | let step = Duration::from_millis(10); 340 | while self.num_running_jobs() > 0 && start.elapsed() < timeout { 341 | tokio::time::sleep(step).await; 342 | } 343 | } 344 | 345 | /// Stop the inner task and wait for it to finish. 346 | pub async fn stop(&mut self) { 347 | if let Some(handle) = self.handle.take() { 348 | handle.stop().await 349 | } 350 | } 351 | } 352 | 353 | async fn start_listener(job_runner: Arc) -> Result { 354 | let mut listener = PgListener::connect_with(&job_runner.options.pool).await?; 355 | if let Some(channels) = &job_runner.options.channel_names { 356 | let names: Vec = channels.iter().map(|c| format!("mq_{}", c)).collect(); 357 | listener 358 | .listen_all(names.iter().map(|s| s.as_str())) 359 | .await?; 360 | } else { 361 | listener.listen("mq").await?; 362 | } 363 | Ok(OwnedHandle::new(task::spawn(async move { 364 | let mut num_errors = 0; 365 | loop { 366 | if num_errors > 0 || listener.recv().await.is_ok() { 367 | job_runner.notify.notify_one(); 368 | num_errors = 0; 369 | } else { 370 | tokio::time::sleep(Duration::from_secs(1 << num_errors)).await; 371 | num_errors += 1; 372 | } 373 | } 374 | }))) 375 | } 376 | 377 | #[derive(sqlx::FromRow)] 378 | struct PolledMessage { 379 | id: Option, 380 | is_committed: Option, 381 | name: Option, 382 | payload_json: Option, 383 | payload_bytes: Option>, 384 | retry_backoff: Option, 385 | wait_time: Option, 386 | } 387 | 388 | fn to_duration(interval: PgInterval) -> Duration { 389 | const SECONDS_PER_DAY: u64 = 24 * 60 * 60; 390 | if interval.microseconds < 0 || interval.days < 0 || interval.months < 0 { 391 | Duration::default() 392 | } else { 393 | let days = (interval.days as u64) + (interval.months as u64) * 30; 394 | Duration::from_micros(interval.microseconds as u64) 395 | + Duration::from_secs(days * SECONDS_PER_DAY) 396 | } 397 | } 398 | 399 | async fn poll_and_dispatch( 400 | job_runner: &Arc, 401 | batch_size: i32, 402 | ) -> Result { 403 | log::info!("Polling for messages"); 404 | 405 | let options = &job_runner.options; 406 | let messages = sqlx::query_as::<_, PolledMessage>("SELECT * FROM mq_poll($1, $2)") 407 | .bind(&options.channel_names) 408 | .bind(batch_size) 409 | .fetch_all(&options.pool) 410 | .await?; 411 | 412 | let ids_to_delete: Vec<_> = messages 413 | .iter() 414 | .filter(|msg| msg.is_committed == Some(false)) 415 | .filter_map(|msg| msg.id) 416 | .collect(); 417 | 418 | log::info!("Deleting {} messages", ids_to_delete.len()); 419 | if !ids_to_delete.is_empty() { 420 | sqlx::query("SELECT mq_delete($1)") 421 | .bind(ids_to_delete) 422 | .execute(&options.pool) 423 | .await?; 424 | } 425 | 426 | const MAX_WAIT: Duration = Duration::from_secs(60); 427 | 428 | let wait_time = messages 429 | .iter() 430 | .filter_map(|msg| msg.wait_time) 431 | .map(to_duration) 432 | .min() 433 | .unwrap_or(MAX_WAIT); 434 | 435 | for msg in messages { 436 | if let PolledMessage { 437 | id: Some(id), 438 | is_committed: Some(true), 439 | name: Some(name), 440 | payload_json, 441 | payload_bytes, 442 | retry_backoff: Some(retry_backoff), 443 | .. 444 | } = msg 445 | { 446 | let retry_backoff = to_duration(retry_backoff); 447 | let keep_alive = if options.keep_alive { 448 | Some(OwnedHandle::new(task::spawn(keep_job_alive( 449 | id, 450 | options.pool.clone(), 451 | retry_backoff, 452 | )))) 453 | } else { 454 | None 455 | }; 456 | let current_job = CurrentJob { 457 | id, 458 | name, 459 | payload_json, 460 | payload_bytes, 461 | job_runner: job_runner.clone(), 462 | keep_alive, 463 | }; 464 | job_runner.running_jobs.fetch_add(1, Ordering::SeqCst); 465 | (options.dispatch)(current_job); 466 | } 467 | } 468 | 469 | Ok(wait_time) 470 | } 471 | 472 | async fn main_loop(job_runner: Arc, _listener_task: OwnedHandle) { 473 | let options = &job_runner.options; 474 | let mut failures = 0; 475 | loop { 476 | let running_jobs = job_runner.running_jobs.load(Ordering::SeqCst); 477 | let duration = if running_jobs < options.min_concurrency { 478 | let batch_size = (options.max_concurrency - running_jobs) as i32; 479 | 480 | match poll_and_dispatch(&job_runner, batch_size).await { 481 | Ok(duration) => { 482 | failures = 0; 483 | duration 484 | } 485 | Err(e) => { 486 | failures += 1; 487 | log::error!("Failed to poll for messages: {}", e); 488 | Duration::from_millis(50 << failures) 489 | } 490 | } 491 | } else { 492 | Duration::from_secs(60) 493 | }; 494 | 495 | // Wait for us to be notified, or for the timeout to elapse 496 | let _ = tokio::time::timeout(duration, job_runner.notify.notified()).await; 497 | } 498 | } 499 | 500 | async fn keep_job_alive(id: Uuid, pool: Pool, mut interval: Duration) { 501 | loop { 502 | tokio::time::sleep(interval / 2).await; 503 | interval *= 2; 504 | if let Err(e) = sqlx::query("SELECT mq_keep_alive(ARRAY[$1], $2)") 505 | .bind(id) 506 | .bind(interval) 507 | .execute(&pool) 508 | .await 509 | { 510 | log::error!("Failed to keep job {} alive: {}", id, e); 511 | break; 512 | } 513 | } 514 | } 515 | -------------------------------------------------------------------------------- /src/spawn.rs: -------------------------------------------------------------------------------- 1 | use std::borrow::Cow; 2 | use std::fmt::Debug; 3 | use std::time::Duration; 4 | 5 | use serde::Serialize; 6 | use sqlx::Postgres; 7 | use uuid::Uuid; 8 | 9 | /// Type for building a job to send. 10 | #[derive(Debug, Clone)] 11 | pub struct JobBuilder<'a> { 12 | id: Uuid, 13 | delay: Duration, 14 | channel_name: &'a str, 15 | channel_args: &'a str, 16 | retries: u32, 17 | retry_backoff: Duration, 18 | commit_interval: Option, 19 | ordered: bool, 20 | name: &'a str, 21 | payload_json: Option>, 22 | payload_bytes: Option<&'a [u8]>, 23 | } 24 | 25 | impl<'a> JobBuilder<'a> { 26 | /// Prepare to send a job with the specified name. 27 | pub fn new(name: &'a str) -> Self { 28 | Self::new_with_id(Uuid::new_v4(), name) 29 | } 30 | /// Prepare to send a job with the specified name and ID. 31 | pub fn new_with_id(id: Uuid, name: &'a str) -> Self { 32 | Self { 33 | id, 34 | delay: Duration::from_secs(0), 35 | channel_name: "", 36 | channel_args: "", 37 | retries: 4, 38 | retry_backoff: Duration::from_secs(1), 39 | commit_interval: None, 40 | ordered: false, 41 | name, 42 | payload_json: None, 43 | payload_bytes: None, 44 | } 45 | } 46 | /// Use the provided function to set any number of configuration 47 | /// options at once. 48 | pub fn set_proto<'b>( 49 | &'b mut self, 50 | proto: impl FnOnce(&'b mut Self) -> &'b mut Self, 51 | ) -> &'b mut Self { 52 | proto(self) 53 | } 54 | /// Set the channel name (default ""). 55 | pub fn set_channel_name(&mut self, channel_name: &'a str) -> &mut Self { 56 | self.channel_name = channel_name; 57 | self 58 | } 59 | /// Set the channel arguments (default ""). 60 | pub fn set_channel_args(&mut self, channel_args: &'a str) -> &mut Self { 61 | self.channel_args = channel_args; 62 | self 63 | } 64 | /// Set the number of retries after the initial attempt (default 4). 65 | pub fn set_retries(&mut self, retries: u32) -> &mut Self { 66 | self.retries = retries; 67 | self 68 | } 69 | /// Set the initial backoff for retries (default 1s). 70 | pub fn set_retry_backoff(&mut self, retry_backoff: Duration) -> &mut Self { 71 | self.retry_backoff = retry_backoff; 72 | self 73 | } 74 | /// Set the commit interval for two-phase commit (default disabled). 75 | pub fn set_commit_interval(&mut self, commit_interval: Option) -> &mut Self { 76 | self.commit_interval = commit_interval; 77 | self 78 | } 79 | /// Set whether this job is strictly ordered with respect to other ordered 80 | /// job in the same channel (default false). 81 | pub fn set_ordered(&mut self, ordered: bool) -> &mut Self { 82 | self.ordered = ordered; 83 | self 84 | } 85 | 86 | /// Set a delay before this job is executed (default none). 87 | pub fn set_delay(&mut self, delay: Duration) -> &mut Self { 88 | self.delay = delay; 89 | self 90 | } 91 | 92 | /// Set a raw JSON payload for the job. 93 | pub fn set_raw_json(&mut self, raw_json: &'a str) -> &mut Self { 94 | self.payload_json = Some(Cow::Borrowed(raw_json)); 95 | self 96 | } 97 | 98 | /// Set a raw binary payload for the job. 99 | pub fn set_raw_bytes(&mut self, raw_bytes: &'a [u8]) -> &mut Self { 100 | self.payload_bytes = Some(raw_bytes); 101 | self 102 | } 103 | 104 | /// Set a JSON payload for the job. 105 | pub fn set_json( 106 | &mut self, 107 | value: &T, 108 | ) -> Result<&mut Self, serde_json::Error> { 109 | let value = serde_json::to_string(value)?; 110 | self.payload_json = Some(Cow::Owned(value)); 111 | Ok(self) 112 | } 113 | 114 | /// Spawn the job using the given executor. This might be a connection 115 | /// pool, a connection, or a transaction. 116 | pub async fn spawn<'b, E: sqlx::Executor<'b, Database = Postgres>>( 117 | &self, 118 | executor: E, 119 | ) -> Result { 120 | sqlx::query( 121 | "SELECT mq_insert(ARRAY[($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)::mq_new_t])", 122 | ) 123 | .bind(self.id) 124 | .bind(self.delay) 125 | .bind(self.retries as i32) 126 | .bind(self.retry_backoff) 127 | .bind(self.channel_name) 128 | .bind(self.channel_args) 129 | .bind(self.commit_interval) 130 | .bind(self.ordered) 131 | .bind(self.name) 132 | .bind(self.payload_json.as_deref()) 133 | .bind(self.payload_bytes) 134 | .execute(executor) 135 | .await?; 136 | Ok(self.id) 137 | } 138 | } 139 | 140 | /// Commit the specified jobs. The jobs should have been previously spawned 141 | /// with the two-phase commit option enabled. 142 | pub async fn commit<'b, E: sqlx::Executor<'b, Database = Postgres>>( 143 | executor: E, 144 | job_ids: &[Uuid], 145 | ) -> Result<(), sqlx::Error> { 146 | sqlx::query("SELECT mq_commit($1)") 147 | .bind(job_ids) 148 | .execute(executor) 149 | .await?; 150 | Ok(()) 151 | } 152 | 153 | /// Clear jobs from the specified channels. 154 | pub async fn clear<'b, E: sqlx::Executor<'b, Database = Postgres>>( 155 | executor: E, 156 | channel_names: &[&str], 157 | ) -> Result<(), sqlx::Error> { 158 | sqlx::query("SELECT mq_clear($1)") 159 | .bind(channel_names) 160 | .execute(executor) 161 | .await?; 162 | Ok(()) 163 | } 164 | 165 | /// Clear jobs from all channels. 166 | pub async fn clear_all<'b, E: sqlx::Executor<'b, Database = Postgres>>( 167 | executor: E, 168 | ) -> Result<(), sqlx::Error> { 169 | sqlx::query("SELECT mq_clear_all()") 170 | .execute(executor) 171 | .await?; 172 | Ok(()) 173 | } 174 | 175 | /// Check if a job with that ID exists 176 | pub async fn exists<'b, E: sqlx::Executor<'b, Database = Postgres>>( 177 | executor: E, 178 | id: Uuid, 179 | ) -> Result { 180 | let exists = sqlx::query_scalar("SELECT EXISTS(SELECT id FROM mq_msgs WHERE id = $1)") 181 | .bind(id) 182 | .fetch_one(executor) 183 | .await?; 184 | Ok(exists) 185 | } 186 | -------------------------------------------------------------------------------- /src/utils.rs: -------------------------------------------------------------------------------- 1 | use std::any::Any; 2 | use std::fmt::{self, Debug}; 3 | use std::ops::{Deref, DerefMut}; 4 | 5 | use tokio::task::JoinHandle; 6 | 7 | #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] 8 | pub struct Opaque(pub T); 9 | 10 | impl Debug for Opaque { 11 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 12 | ::fmt(&self.0, f) 13 | } 14 | } 15 | 16 | impl Deref for Opaque { 17 | type Target = T; 18 | 19 | fn deref(&self) -> &Self::Target { 20 | &self.0 21 | } 22 | } 23 | 24 | impl DerefMut for Opaque { 25 | fn deref_mut(&mut self) -> &mut Self::Target { 26 | &mut self.0 27 | } 28 | } 29 | 30 | /// A handle to a background task which will be automatically cancelled if 31 | /// the handle is dropped. Extract the inner join handle to prevent this 32 | /// behaviour. 33 | #[derive(Debug)] 34 | pub struct OwnedHandle(Option>); 35 | 36 | impl OwnedHandle { 37 | /// Construct a new `OwnedHandle` from the provided `JoinHandle` 38 | pub fn new(inner: JoinHandle<()>) -> Self { 39 | Self(Some(inner)) 40 | } 41 | /// Get back the original `JoinHandle` 42 | pub fn into_inner(mut self) -> JoinHandle<()> { 43 | self.0.take().expect("Only consumed once") 44 | } 45 | /// Stop the task and wait for it to finish. 46 | pub async fn stop(self) { 47 | let handle = self.into_inner(); 48 | handle.abort(); 49 | let _ = handle.await; 50 | } 51 | } 52 | 53 | impl Drop for OwnedHandle { 54 | fn drop(&mut self) { 55 | if let Some(handle) = self.0.take() { 56 | handle.abort(); 57 | } 58 | } 59 | } 60 | --------------------------------------------------------------------------------