From 663dc00fa821992772d4da92e27b9c018c95f25e Mon Sep 17 00:00:00 2001 From: Mia Date: Fri, 7 Feb 2025 18:54:09 +0000 Subject: [PATCH] bring dataloader-rs in, we need custom tweaks --- Cargo.lock | 255 ++++++++++++++++++++++++++++++- Cargo.toml | 1 + dataloader-rs/.gitignore | 2 + dataloader-rs/Cargo.toml | 32 ++++ dataloader-rs/LICENSE-APACHE | 201 +++++++++++++++++++++++++ dataloader-rs/LICENSE-MIT | 25 ++++ dataloader-rs/README.md | 93 ++++++++++++ dataloader-rs/src/batch_fn.rs | 5 + dataloader-rs/src/cached.rs | 256 ++++++++++++++++++++++++++++++++ dataloader-rs/src/lib.rs | 32 ++++ dataloader-rs/src/non_cached.rs | 231 ++++++++++++++++++++++++++++ dataloader-rs/src/runtime.rs | 19 +++ parakeet/Cargo.toml | 2 +- 13 files changed, 1150 insertions(+), 4 deletions(-) create mode 100644 dataloader-rs/.gitignore create mode 100644 dataloader-rs/Cargo.toml create mode 100644 dataloader-rs/LICENSE-APACHE create mode 100644 dataloader-rs/LICENSE-MIT create mode 100644 dataloader-rs/README.md create mode 100644 dataloader-rs/src/batch_fn.rs create mode 100644 dataloader-rs/src/cached.rs create mode 100644 dataloader-rs/src/lib.rs create mode 100644 dataloader-rs/src/non_cached.rs create mode 100644 dataloader-rs/src/runtime.rs diff --git a/Cargo.lock b/Cargo.lock index 2bb4ccb7..06d49c47 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,6 +88,119 @@ version = "1.0.95" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04" +[[package]] +name = "async-channel" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81953c529336010edd6d8e358f886d9581267795c61b19475b71314bffa46d35" +dependencies = [ + "concurrent-queue", + "event-listener 2.5.3", + "futures-core", +] + +[[package]] +name = "async-channel" +version = "2.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89b47800b0be77592da0afd425cc03468052844aff33b84e33cc696f64e77b6a" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-executor" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30ca9a001c1e8ba5149f91a74362376cc6bc5b919d92d988668657bd570bdcec" +dependencies = [ + "async-task", + "concurrent-queue", + "fastrand", + "futures-lite", + "slab", +] + +[[package]] +name = "async-global-executor" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05b1b633a2115cd122d73b955eadd9916c18c8f510ec9cd1686404c60ad1c29c" +dependencies = [ + "async-channel 2.3.1", + "async-executor", + "async-io", + "async-lock", + "blocking", + "futures-lite", + "once_cell", +] + +[[package]] +name = "async-io" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a2b323ccce0a1d90b449fd71f2a06ca7faa7c54c2751f06c9bd851fc061059" +dependencies = [ + "async-lock", + "cfg-if", + "concurrent-queue", + "futures-io", + "futures-lite", + "parking", + "polling", + "rustix", + "slab", + "tracing", + "windows-sys 0.59.0", +] + +[[package]] +name = "async-lock" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff6e472cdea888a4bd64f342f09b3f50e1886d32afe8df3d663c01140b811b18" +dependencies = [ + "event-listener 5.4.0", + "event-listener-strategy", + "pin-project-lite", +] + +[[package]] +name = "async-std" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c634475f29802fde2b8f0b505b1bd00dfe4df7d4a000f0b36f7671197d5c3615" +dependencies = [ + "async-channel 1.9.0", + "async-global-executor", + "async-io", + "async-lock", + "crossbeam-utils", + "futures-channel", + "futures-core", + "futures-io", + "futures-lite", + "gloo-timers", + "kv-log-macro", + "log", + "memchr", + "once_cell", + "pin-project-lite", + "pin-utils", + "slab", + "wasm-bindgen-futures", +] + +[[package]] +name = "async-task" +version = "4.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de" + [[package]] name = "async-trait" version = "0.1.85" @@ -216,6 +329,19 @@ dependencies = [ "generic-array", ] +[[package]] +name = "blocking" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703f41c54fc768e63e091340b424302bb1c29ef4aa0c7f10fe849dfb114d29ea" +dependencies = [ + "async-channel 2.3.1", + "async-task", + "futures-io", + "futures-lite", + "piper", +] + [[package]] name = "bumpalo" version = "3.16.0" @@ -366,6 +492,15 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990" +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "consumer" version = "0.1.0" @@ -425,6 +560,12 @@ dependencies = [ "libc", ] +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + [[package]] name = "crunchy" version = "0.2.3" @@ -505,9 +646,9 @@ dependencies = [ [[package]] name = "dataloader" version = "0.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43533eb41c886a84d11ed5c72c17315dcdff46a5eb08d3bbfac6a9ef8faa4085" dependencies = [ + "async-std", + "futures", "tokio", ] @@ -669,6 +810,33 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "event-listener" +version = "2.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0" + +[[package]] +name = "event-listener" +version = "5.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3492acde4c3fc54c845eaab3eed8bd00c7a7d881f78bfc801e43a93dec1331ae" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c3e4e0dd3673c1139bf041f3008816d9cf2946bbfac2945c09e523b8d7b05b2" +dependencies = [ + "event-listener 5.4.0", + "pin-project-lite", +] + [[package]] name = "eyre" version = "0.6.12" @@ -783,6 +951,19 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +[[package]] +name = "futures-lite" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5edaec856126859abb19ed65f39e90fea3a9574b9707f13539acf4abf7eb532" +dependencies = [ + "fastrand", + "futures-core", + "futures-io", + "parking", + "pin-project-lite", +] + [[package]] name = "futures-macro" version = "0.3.31" @@ -851,6 +1032,18 @@ version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +[[package]] +name = "gloo-timers" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbb143cf96099802033e0d4f4963b19fd2e0b728bcf076cd9cf7f6634f092994" +dependencies = [ + "futures-channel", + "futures-core", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "h2" version = "0.4.7" @@ -898,6 +1091,12 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +[[package]] +name = "hermit-abi" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc" + [[package]] name = "hickory-proto" version = "0.24.2" @@ -1339,6 +1538,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "kv-log-macro" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de8b303297635ad57c9f5059fd9cee7a47f8e8daa09df0fcd07dd39fb22977f" +dependencies = [ + "log", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -1393,6 +1601,9 @@ name = "log" version = "0.4.25" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f" +dependencies = [ + "value-bag", +] [[package]] name = "lru-cache" @@ -1521,7 +1732,7 @@ version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" dependencies = [ - "hermit-abi", + "hermit-abi 0.3.9", "libc", ] @@ -1631,6 +1842,12 @@ dependencies = [ "walkdir", ] +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + [[package]] name = "parking_lot" version = "0.12.3" @@ -1713,12 +1930,38 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "piper" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c8c490f422ef9a4efd2cb5b42b76c8613d7e7dfc1caf667b8a3350a5acc066" +dependencies = [ + "atomic-waker", + "fastrand", + "futures-io", +] + [[package]] name = "pkg-config" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" +[[package]] +name = "polling" +version = "3.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a604568c3202727d1507653cb121dbd627a58684eb09a820fd746bee38b4442f" +dependencies = [ + "cfg-if", + "concurrent-queue", + "hermit-abi 0.4.0", + "pin-project-lite", + "rustix", + "tracing", + "windows-sys 0.59.0", +] + [[package]] name = "postgres-protocol" version = "0.6.7" @@ -2729,6 +2972,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "value-bag" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ef4c4aa54d5d05a279399bfa921ec387b7aba77caf7a682ae8d86785b8fdad2" + [[package]] name = "vcpkg" version = "0.2.15" diff --git a/Cargo.toml b/Cargo.toml index 397668e5..4b970c57 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ resolver = "2" members = [ "consumer", + "dataloader-rs", "did-resolver", "lexica", "parakeet", diff --git a/dataloader-rs/.gitignore b/dataloader-rs/.gitignore new file mode 100644 index 00000000..a9d37c56 --- /dev/null +++ b/dataloader-rs/.gitignore @@ -0,0 +1,2 @@ +target +Cargo.lock diff --git a/dataloader-rs/Cargo.toml b/dataloader-rs/Cargo.toml new file mode 100644 index 00000000..d42c86db --- /dev/null +++ b/dataloader-rs/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "dataloader" +version = "0.18.0" +edition = "2021" +authors = ["cksac ", "Lily"] +description = "Rust implementation of Facebook's DataLoader using async-await." +keywords = ["batcher", "dataloader", "cache"] +categories = ["asynchronous", "caching"] +license = "MIT/Apache-2.0" +readme = "README.md" +repository = "https://github.com/cksac/dataloader-rs" +homepage = "https://github.com/cksac/dataloader-rs" +documentation = "https://docs.rs/dataloader" + +[badges] +travis-ci = { repository = "/cksac/dataloader-rs" } + +[features] +default = ["runtime-async-std"] +runtime-async-std = [ + "async-std", +] +runtime-tokio = [ + "tokio" +] + +[dependencies] +async-std = { version = "1", optional = true } +tokio = { version = "1", features = [ "sync", "rt" ], optional = true } + +[dev-dependencies] +futures = "0.3" diff --git a/dataloader-rs/LICENSE-APACHE b/dataloader-rs/LICENSE-APACHE new file mode 100644 index 00000000..cfc7d73b --- /dev/null +++ b/dataloader-rs/LICENSE-APACHE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright (c) 2017 cksac + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/dataloader-rs/LICENSE-MIT b/dataloader-rs/LICENSE-MIT new file mode 100644 index 00000000..cb12b455 --- /dev/null +++ b/dataloader-rs/LICENSE-MIT @@ -0,0 +1,25 @@ +Copyright (c) 2017 cksac + +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. \ No newline at end of file diff --git a/dataloader-rs/README.md b/dataloader-rs/README.md new file mode 100644 index 00000000..04db77eb --- /dev/null +++ b/dataloader-rs/README.md @@ -0,0 +1,93 @@ +# Dataloader + +![Rust](https://github.com/cksac/dataloader-rs/workflows/Rust/badge.svg) +[![Crates.io](https://img.shields.io/crates/v/dataloader.svg)](https://crates.io/crates/dataloader) + +Rust implementation of [Facebook's DataLoader](https://github.com/facebook/dataloader) using async-await. + +[Documentation](https://docs.rs/dataloader) + +## Features +* [x] Batching load requests with caching +* [x] Batching load requests without caching + +## Usage +### Switching runtime, by using cargo features +- `runtime-async-std` (default), to use the [async-std](https://async.rs) runtime + - dataloader = "0.18" +- `runtime-tokio` to use the [Tokio](https://tokio.rs) runtime + - dataloader = { version = "0.18", default-features = false, features = ["runtime-tokio"]} + + +### Add to your `Cargo.toml`: +```toml +[dependencies] +dataloader = "0.18" +futures = "0.3" +``` + +### Example: +```rust +use dataloader::cached::Loader; +use dataloader::BatchFn; +use futures::executor::block_on; +use futures::future::ready; +use std::collections::HashMap; +use std::thread; + +struct MyLoadFn; + +impl BatchFn for MyLoadFn { + async fn load(&mut self, keys: &[usize]) -> HashMap { + println!("BatchFn load keys {:?}", keys); + let ret = keys.iter() + .map(|v| (v.clone(), v.clone())) + .collect::>(); + ready(ret).await + } +} + +fn main() { + let mut i = 0; + while i < 2 { + let a = MyLoadFn; + let loader = Loader::new(a).with_max_batch_size(4); + + let l1 = loader.clone(); + let h1 = thread::spawn(move || { + let r1 = l1.load(1); + let r2 = l1.load(2); + let r3 = l1.load(3); + + let r4 = l1.load_many(vec![2, 3, 4, 5, 6, 7, 8]); + let f = futures::future::join4(r1, r2, r3, r4); + println!("{:?}", block_on(f)); + }); + + let l2 = loader.clone(); + let h2 = thread::spawn(move || { + let r1 = l2.load(1); + let r2 = l2.load(2); + let r3 = l2.load(3); + let r4 = l2.load(4); + let f = futures::future::join4(r1, r2, r3, r4); + println!("{:?}", block_on(f)); + }); + + h1.join().unwrap(); + h2.join().unwrap(); + i += 1; + } +} +``` + +# LICENSE + +This project is licensed under either of + + * Apache License, Version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or + http://www.apache.org/licenses/LICENSE-2.0) + * MIT license ([LICENSE-MIT](LICENSE-MIT) or + http://opensource.org/licenses/MIT) + +at your option. \ No newline at end of file diff --git a/dataloader-rs/src/batch_fn.rs b/dataloader-rs/src/batch_fn.rs new file mode 100644 index 00000000..97ec9ffc --- /dev/null +++ b/dataloader-rs/src/batch_fn.rs @@ -0,0 +1,5 @@ +use std::collections::HashMap; + +pub trait BatchFn { + async fn load(&mut self, keys: &[K]) -> HashMap; +} diff --git a/dataloader-rs/src/cached.rs b/dataloader-rs/src/cached.rs new file mode 100644 index 00000000..40cc577a --- /dev/null +++ b/dataloader-rs/src/cached.rs @@ -0,0 +1,256 @@ +use crate::runtime::{Arc, Mutex}; +use crate::{yield_fn, BatchFn, WaitForWorkFn}; +use std::collections::{HashMap, HashSet}; +use std::fmt::Debug; +use std::hash::{BuildHasher, Hash}; +use std::iter::IntoIterator; + +pub trait Cache { + type Key; + type Val; + fn get(&mut self, key: &Self::Key) -> Option<&Self::Val>; + fn insert(&mut self, key: Self::Key, val: Self::Val); + fn remove(&mut self, key: &Self::Key) -> Option; + fn clear(&mut self); +} + +impl Cache for HashMap +where + K: Eq + Hash, +{ + type Key = K; + type Val = V; + + #[inline] + fn get(&mut self, key: &K) -> Option<&V> { + HashMap::get(self, key) + } + + #[inline] + fn insert(&mut self, key: K, val: V) { + HashMap::insert(self, key, val); + } + + #[inline] + fn remove(&mut self, key: &K) -> Option { + HashMap::remove(self, key) + } + + #[inline] + fn clear(&mut self) { + HashMap::clear(self) + } +} + +struct State> +where + C: Cache, +{ + completed: C, + pending: HashSet, +} + +impl State +where + C: Cache, +{ + fn with_cache(cache: C) -> Self { + State { + completed: cache, + pending: HashSet::new(), + } + } +} + +pub struct Loader> +where + K: Eq + Hash + Clone, + V: Clone, + F: BatchFn, + C: Cache, +{ + state: Arc>>, + load_fn: Arc>, + wait_for_work_fn: Arc, + max_batch_size: usize, +} + +impl Clone for Loader +where + K: Eq + Hash + Clone, + V: Clone, + F: BatchFn, + C: Cache, +{ + fn clone(&self) -> Self { + Loader { + state: self.state.clone(), + max_batch_size: self.max_batch_size, + load_fn: self.load_fn.clone(), + wait_for_work_fn: self.wait_for_work_fn.clone(), + } + } +} + +#[allow(clippy::implicit_hasher)] +impl Loader> +where + K: Eq + Hash + Clone + Debug, + V: Clone, + F: BatchFn, +{ + pub fn new(load_fn: F) -> Loader> { + Loader::with_cache(load_fn, HashMap::new()) + } +} + +impl Loader +where + K: Eq + Hash + Clone + Debug, + V: Clone, + F: BatchFn, + C: Cache, +{ + pub fn with_cache(load_fn: F, cache: C) -> Loader { + Loader { + state: Arc::new(Mutex::new(State::with_cache(cache))), + load_fn: Arc::new(Mutex::new(load_fn)), + max_batch_size: 200, + wait_for_work_fn: Arc::new(yield_fn(10)), + } + } + + pub fn with_max_batch_size(mut self, max_batch_size: usize) -> Self { + self.max_batch_size = max_batch_size; + self + } + + pub fn with_yield_count(mut self, yield_count: usize) -> Self { + self.wait_for_work_fn = Arc::new(yield_fn(yield_count)); + self + } + + /// Replaces the yielding for work behavior with an arbitrary future. Rather than yielding + /// the runtime repeatedly this will generate and `.await` a future of your choice. + /// ***This is incompatible with*** [`Self::with_yield_count()`]. + pub fn with_custom_wait_for_work(mut self, wait_for_work_fn: impl WaitForWorkFn) -> Self { + self.wait_for_work_fn = Arc::new(wait_for_work_fn); + self + } + + pub fn max_batch_size(&self) -> usize { + self.max_batch_size + } + + pub async fn load(&self, key: K) -> Option { + let mut state = self.state.lock().await; + if let Some(v) = state.completed.get(&key) { + return Some((*v).clone()); + } + + if !state.pending.contains(&key) { + state.pending.insert(key.clone()); + if state.pending.len() >= self.max_batch_size { + let keys = state.pending.drain().collect::>(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (k, v) in load_ret.into_iter() { + state.completed.insert(k, v); + } + return state.completed.get(&key).cloned(); + } + } + drop(state); + + (self.wait_for_work_fn)().await; + + let mut state = self.state.lock().await; + if let Some(v) = state.completed.get(&key) { + return Some((*v).clone()); + } + + if !state.pending.is_empty() { + let keys = state.pending.drain().collect::>(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (k, v) in load_ret.into_iter() { + state.completed.insert(k, v); + } + } + + state.completed.get(&key).cloned() + } + + pub async fn load_many(&self, keys: Vec) -> HashMap { + let mut state = self.state.lock().await; + let mut ret = HashMap::new(); + let mut rest = Vec::new(); + for key in keys.into_iter() { + if let Some(v) = state.completed.get(&key).cloned() { + ret.insert(key, v); + continue; + } + if !state.pending.contains(&key) { + state.pending.insert(key.clone()); + if state.pending.len() >= self.max_batch_size { + let keys = state.pending.drain().collect::>(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (k, v) in load_ret.into_iter() { + state.completed.insert(k, v); + } + } + } + rest.push(key); + } + drop(state); + + (self.wait_for_work_fn)().await; + + if !rest.is_empty() { + let mut state = self.state.lock().await; + if !state.pending.is_empty() { + let keys = state.pending.drain().collect::>(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (k, v) in load_ret.into_iter() { + state.completed.insert(k, v); + } + } + + for key in rest.into_iter() { + if let Some(v) = state.completed.get(&key).cloned() { + ret.insert(key, v); + } + } + } + + ret + } + + pub async fn prime(&self, key: K, val: V) { + let mut state = self.state.lock().await; + state.completed.insert(key, val); + } + + pub async fn prime_many(&self, values: impl IntoIterator) { + let mut state = self.state.lock().await; + for (k, v) in values.into_iter() { + state.completed.insert(k, v); + } + } + + pub async fn clear(&self, key: K) { + let mut state = self.state.lock().await; + state.completed.remove(&key); + } + + pub async fn clear_all(&self) { + let mut state = self.state.lock().await; + state.completed.clear() + } +} diff --git a/dataloader-rs/src/lib.rs b/dataloader-rs/src/lib.rs new file mode 100644 index 00000000..b06de808 --- /dev/null +++ b/dataloader-rs/src/lib.rs @@ -0,0 +1,32 @@ +#![allow(async_fn_in_trait)] + +mod batch_fn; +pub mod cached; +pub mod non_cached; +mod runtime; + +pub use batch_fn::BatchFn; + +use std::{future::Future, pin::Pin}; + +/// A trait alias. Read as "a function which returns a pinned box containing a future" +pub trait WaitForWorkFn: + Fn() -> Pin + Send + Sync>> + Send + Sync + 'static +{ +} + +impl WaitForWorkFn for T where + T: Fn() -> Pin + Send + Sync>> + Send + Sync + 'static +{ +} + +pub(crate) fn yield_fn(count: usize) -> impl WaitForWorkFn { + move || { + Box::pin(async move { + // yield for other load to append request + for _ in 0..count { + runtime::yield_now().await; + } + }) + } +} diff --git a/dataloader-rs/src/non_cached.rs b/dataloader-rs/src/non_cached.rs new file mode 100644 index 00000000..1d7956e1 --- /dev/null +++ b/dataloader-rs/src/non_cached.rs @@ -0,0 +1,231 @@ +use crate::runtime::{Arc, Mutex}; +use crate::{yield_fn, BatchFn, WaitForWorkFn}; +use std::collections::{HashMap, HashSet}; +use std::fmt::Debug; +use std::hash::Hash; + +type RequestId = usize; + +struct State { + completed: HashMap, + failed: HashMap, + pending: HashMap, + id_seq: RequestId, +} + +impl State { + fn new() -> Self { + State { + completed: HashMap::new(), + failed: HashMap::new(), + pending: HashMap::new(), + id_seq: 0, + } + } + fn next_request_id(&mut self) -> RequestId { + self.id_seq = self.id_seq.wrapping_add(1); + self.id_seq + } +} + +pub struct Loader +where + K: Eq + Hash + Clone, + V: Clone, + F: BatchFn, +{ + state: Arc>>, + load_fn: Arc>, + wait_for_work_fn: Arc, + max_batch_size: usize, +} + +impl Clone for Loader +where + K: Eq + Hash + Clone, + V: Clone, + F: BatchFn, +{ + fn clone(&self) -> Self { + Loader { + state: self.state.clone(), + load_fn: self.load_fn.clone(), + max_batch_size: self.max_batch_size, + wait_for_work_fn: self.wait_for_work_fn.clone(), + } + } +} + +impl Loader +where + K: Eq + Hash + Clone + Debug, + V: Clone, + F: BatchFn, +{ + pub fn new(load_fn: F) -> Loader { + Loader { + state: Arc::new(Mutex::new(State::new())), + load_fn: Arc::new(Mutex::new(load_fn)), + max_batch_size: 200, + wait_for_work_fn: Arc::new(yield_fn(10)), + } + } + + pub fn with_max_batch_size(mut self, max_batch_size: usize) -> Self { + self.max_batch_size = max_batch_size; + self + } + + pub fn with_yield_count(mut self, yield_count: usize) -> Self { + self.wait_for_work_fn = Arc::new(yield_fn(yield_count)); + self + } + + /// Replaces the yielding for work behavior with an arbitrary future. Rather than yielding + /// the runtime repeatedly this will generate and `.await` a future of your choice. + /// ***This is incompatible with*** [`Self::with_yield_count()`]. + pub fn with_custom_wait_for_work(mut self, wait_for_work_fn: impl WaitForWorkFn) -> Self { + self.wait_for_work_fn = Arc::new(wait_for_work_fn); + self + } + + pub fn max_batch_size(&self) -> usize { + self.max_batch_size + } + + pub async fn load(&self, key: K) -> Option { + let mut state = self.state.lock().await; + let request_id = state.next_request_id(); + state.pending.insert(request_id, key); + if state.pending.len() >= self.max_batch_size { + let batch = state.pending.drain().collect::>(); + let keys: Vec = batch + .values() + .cloned() + .collect::>() + .into_iter() + .collect(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (request_id, key) in batch.into_iter() { + if load_ret + .get(&key) + .and_then(|v| state.completed.insert(request_id, v.clone())) + .is_none() + { + state.failed.insert(request_id, key); + } + } + return state.completed.remove(&request_id); + } + drop(state); + + (self.wait_for_work_fn)().await; + + let mut state = self.state.lock().await; + + if !state.completed.contains_key(&request_id) { + let batch = state.pending.drain().collect::>(); + if !batch.is_empty() { + let keys: Vec = batch + .values() + .cloned() + .collect::>() + .into_iter() + .collect(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (request_id, key) in batch.into_iter() { + if load_ret + .get(&key) + .and_then(|v| state.completed.insert(request_id, v.clone())) + .is_none() + { + state.failed.insert(request_id, key); + } + } + } + } + state.completed.remove(&request_id) + } + + pub async fn load_many(&self, keys: Vec) -> HashMap { + let mut state = self.state.lock().await; + let mut ret = HashMap::new(); + let mut requests = Vec::new(); + for key in keys.into_iter() { + let request_id = state.next_request_id(); + requests.push((request_id, key.clone())); + state.pending.insert(request_id, key); + if state.pending.len() >= self.max_batch_size { + let batch = state.pending.drain().collect::>(); + let keys: Vec = batch + .values() + .cloned() + .collect::>() + .into_iter() + .collect(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (request_id, key) in batch.into_iter() { + if load_ret + .get(&key) + .and_then(|v| state.completed.insert(request_id, v.clone())) + .is_none() + { + state.failed.insert(request_id, key); + } + } + } + } + + drop(state); + + (self.wait_for_work_fn)().await; + + let mut state = self.state.lock().await; + + let mut rest = Vec::new(); + for (request_id, key) in requests.into_iter() { + if let Some(v) = state.completed.remove(&request_id) { + ret.insert(key, v); + } else { + rest.push((request_id, key)); + } + } + + if !rest.is_empty() { + let batch = state.pending.drain().collect::>(); + if !batch.is_empty() { + let keys: Vec = batch + .values() + .cloned() + .collect::>() + .into_iter() + .collect(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (request_id, key) in batch.into_iter() { + if load_ret + .get(&key) + .and_then(|v| state.completed.insert(request_id, v.clone())) + .is_none() + { + state.failed.insert(request_id, key); + } + } + } + for (request_id, key) in rest.into_iter() { + if let Some(v) = state.completed.remove(&request_id) { + ret.insert(key, v); + } + } + } + + ret + } +} diff --git a/dataloader-rs/src/runtime.rs b/dataloader-rs/src/runtime.rs new file mode 100644 index 00000000..e92f4e4f --- /dev/null +++ b/dataloader-rs/src/runtime.rs @@ -0,0 +1,19 @@ +// runtime-async-std +#[cfg(feature = "runtime-async-std")] +pub type Arc = async_std::sync::Arc; + +#[cfg(feature = "runtime-async-std")] +pub type Mutex = async_std::sync::Mutex; + +#[cfg(feature = "runtime-async-std")] +pub use async_std::task::yield_now; + +// runtime-tokio +#[cfg(feature = "runtime-tokio")] +pub type Arc = std::sync::Arc; + +#[cfg(feature = "runtime-tokio")] +pub type Mutex = tokio::sync::Mutex; + +#[cfg(feature = "runtime-tokio")] +pub use tokio::task::yield_now; diff --git a/parakeet/Cargo.toml b/parakeet/Cargo.toml index 10779f70..9c9269c4 100644 --- a/parakeet/Cargo.toml +++ b/parakeet/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" [dependencies] axum = { version = "0.8", features = ["json"] } -dataloader = { version = "0.18", default-features = false, features = ["runtime-tokio"] } +dataloader = { path = "../dataloader-rs", default-features = false, features = ["runtime-tokio"] } deadpool = { version = "0.12.1", features = ["managed"] } diesel = { version = "2.2.6", features = ["chrono", "serde_json"] } diesel-async = { version = "0.5.2", features = ["deadpool", "postgres"] } -- 2.51.2