├── .github └── workflows │ └── ci.yml ├── .gitignore ├── Cargo.toml ├── LICENSE ├── README.md ├── benches └── cache_benchmark.rs └── src ├── lib.rs └── lib_test.rs /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | on: [push] 2 | 3 | name: ci 4 | 5 | jobs: 6 | test: 7 | name: Test Suite 8 | runs-on: ubuntu-latest 9 | steps: 10 | - name: Checkout sources 11 | uses: actions/checkout@v2 12 | - name: Install stable toolchain 13 | uses: actions-rs/toolchain@v1 14 | with: 15 | profile: minimal 16 | toolchain: stable 17 | override: true 18 | - name: Run cargo test 19 | uses: actions-rs/cargo@v1 20 | with: 21 | command: test 22 | - name: Benchmark 23 | uses: actions-rs/cargo@v1 24 | with: 25 | command: bench 26 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | # will have compiled files and executables 3 | /target/ 4 | 5 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 6 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 7 | Cargo.lock 8 | 9 | # These are backup files generated by rustfmt 10 | **/*.rs.bk 11 | 12 | 13 | # Added by cargo 14 | 15 | /target 16 | /Cargo.lock 17 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "sqlite-cache" 3 | version = "0.1.4" 4 | edition = "2021" 5 | license = "Apache-2.0" 6 | authors = ["Heyang Zhou "] 7 | description = "SQLite-based on-disk cache." 8 | repository = "https://github.com/losfair/sqlite-cache" 9 | 10 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 11 | 12 | [dependencies] 13 | rusqlite = "0.32" 14 | data-encoding = "2.6.0" 15 | futures = "0.3.30" 16 | tracing = "0.1.40" 17 | 18 | [dev-dependencies] 19 | tracing-test = "0.2" 20 | tokio = { version = "1.39", features = ["macros", "rt", "time"] } 21 | rand = "0.8" 22 | criterion = "0.5" 23 | 24 | [[bench]] 25 | name = "cache_benchmark" 26 | harness = false 27 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # sqlite-cache 2 | 3 | [![crates.io](https://img.shields.io/crates/v/sqlite-cache.svg)](https://crates.io/crates/sqlite-cache) 4 | 5 | SQLite-based on-disk cache for Rust. 6 | 7 | ## Usage 8 | 9 | ```rust 10 | let cache = Cache::new( 11 | CacheConfig::default(), 12 | rusqlite::Connection::open_in_memory().unwrap(), 13 | ).unwrap(); 14 | let topic = cache.topic("test-topic").unwrap(); 15 | assert!(topic.get("hello").unwrap().is_none()); 16 | topic.set("hello", b"world", Duration::from_secs(60)) 17 | assert!(&topic.get("hello").unwrap().unwrap().data[..] == b"world"); 18 | ``` 19 | 20 | ### Locked updates 21 | 22 | This library supports *locked updates* to prevent the [thundering herd problem](https://en.wikipedia.org/wiki/Thundering_herd_problem) on cache misses. The `get_for_update` API acquires a per-key lock and returns a `KeyUpdater`; subsequent `get_for_update` calls on the same key will block until the previous `KeyUpdater` is dropped. 23 | 24 | ```rust 25 | let (updater, current_value) = topic.get_for_update("hello").await.unwrap(); 26 | let new_value = expensive_computation(current_value).await; 27 | updater.write(new_value, Duration::from_secs(60)).unwrap(); 28 | ``` 29 | 30 | ## Benchmark 31 | 32 | These results are from running `benches/cache_benchmark.rs` on an Apple M1 processor. 33 | 34 | - `mt(4)`: Per-thread operation latency when running the same task on 4 threads. 35 | 36 | ``` 37 | lookup - cache size 10000 38 | time: [1.5978 us 1.6051 us 1.6130 us] 39 | lookup mt(4) - cache size 10000 40 | time: [9.7801 us 9.8464 us 9.9329 us] 41 | insert - cache size 10000 42 | time: [4.6316 us 4.6785 us 4.7169 us] 43 | insert mt(4) - cache size 10000 44 | time: [21.195 us 21.420 us 21.614 us] 45 | ``` 46 | -------------------------------------------------------------------------------- /benches/cache_benchmark.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | sync::Arc, 3 | time::{Duration, Instant}, 4 | }; 5 | 6 | use criterion::{criterion_group, criterion_main, Bencher, Criterion}; 7 | use data_encoding::BASE64; 8 | use rand::{ 9 | prelude::{SliceRandom, ThreadRng}, 10 | thread_rng, RngCore, 11 | }; 12 | use rusqlite::Connection; 13 | use sqlite_cache::{Cache, CacheConfig}; 14 | 15 | fn gen_pairs(rng: &mut ThreadRng, size: usize) -> Vec<(String, Vec)> { 16 | (0..size) 17 | .map(|_| { 18 | let mut k = vec![0u8; (rng.next_u64() % 100 + 32) as usize]; 19 | let mut v = vec![0u8; (rng.next_u64() % 1000) as usize]; 20 | rng.fill_bytes(&mut k); 21 | rng.fill_bytes(&mut v); 22 | (BASE64.encode(&k), v) 23 | }) 24 | .collect::>() 25 | } 26 | 27 | fn bench_get_mt(b: &mut Bencher, size: usize, num_threads: usize) { 28 | let mut rng = thread_rng(); 29 | let cache = Cache::new( 30 | CacheConfig::default(), 31 | Connection::open_in_memory().unwrap(), 32 | ) 33 | .unwrap(); 34 | let topic = cache.topic("test").unwrap(); 35 | let pairs = gen_pairs(&mut rng, size); 36 | for (k, v) in &pairs { 37 | topic.set(k, v, Duration::from_secs(3600)).unwrap(); 38 | } 39 | 40 | let pairs = Arc::new(pairs); 41 | 42 | b.iter_custom(|n| { 43 | let start = Instant::now(); 44 | let handles = (0..num_threads) 45 | .map(|_| { 46 | let pairs = pairs.clone(); 47 | let topic = topic.clone(); 48 | std::thread::spawn(move || { 49 | let mut rng = thread_rng(); 50 | for _ in 0..n { 51 | let (k, v) = pairs.choose(&mut rng).unwrap(); 52 | let got_value = topic.get(k).unwrap().unwrap(); 53 | assert_eq!(&got_value.data, v); 54 | } 55 | }) 56 | }) 57 | .collect::>(); 58 | for h in handles { 59 | h.join().unwrap(); 60 | } 61 | start.elapsed() 62 | }); 63 | } 64 | 65 | fn bench_get(b: &mut Bencher, size: usize) { 66 | bench_get_mt(b, size, 1) 67 | } 68 | 69 | fn bench_set_mt(b: &mut Bencher, size: usize, num_threads: usize) { 70 | let mut rng = thread_rng(); 71 | let cache = Cache::new( 72 | CacheConfig::default(), 73 | Connection::open_in_memory().unwrap(), 74 | ) 75 | .unwrap(); 76 | let topic = cache.topic("test").unwrap(); 77 | let pairs = gen_pairs(&mut rng, size); 78 | 79 | let pairs = Arc::new(pairs); 80 | 81 | b.iter_custom(|n| { 82 | let start = Instant::now(); 83 | let handles = (0..num_threads) 84 | .map(|_| { 85 | let pairs = pairs.clone(); 86 | let topic = topic.clone(); 87 | std::thread::spawn(move || { 88 | let mut rng = thread_rng(); 89 | for _ in 0..n { 90 | let (k, v) = pairs.choose(&mut rng).unwrap(); 91 | topic.set(k, v, Duration::from_secs(3600)).unwrap(); 92 | } 93 | }) 94 | }) 95 | .collect::>(); 96 | for h in handles { 97 | h.join().unwrap(); 98 | } 99 | start.elapsed() 100 | }); 101 | } 102 | 103 | fn bench_set(b: &mut Bencher, size: usize) { 104 | bench_set_mt(b, size, 1) 105 | } 106 | 107 | fn criterion_benchmark(c: &mut Criterion) { 108 | //c.bench_function("lookup - cache size 1000", |b| bench_get(b, 1000)); 109 | c.bench_function("lookup - cache size 10000", |b| bench_get(b, 10000)); 110 | c.bench_function("lookup mt(4) - cache size 10000", |b| { 111 | bench_get_mt(b, 10000, 4) 112 | }); 113 | //c.bench_function("lookup - cache size 50000", |b| bench_get(b, 50000)); 114 | //c.bench_function("insert - cache size 1000", |b| bench_set(b, 1000)); 115 | c.bench_function("insert - cache size 10000", |b| bench_set(b, 10000)); 116 | c.bench_function("insert mt(4) - cache size 10000", |b| { 117 | bench_set_mt(b, 10000, 4) 118 | }); 119 | //c.bench_function("insert - cache size 50000", |b| bench_set(b, 50000)); 120 | } 121 | 122 | criterion_group!(benches, criterion_benchmark); 123 | criterion_main!(benches); 124 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | #[cfg(test)] 2 | mod lib_test; 3 | 4 | use data_encoding::BASE32_NOPAD; 5 | use futures::channel::oneshot::{channel, Receiver, Sender}; 6 | pub use rusqlite; 7 | 8 | use std::{ 9 | collections::HashMap, 10 | sync::{mpsc, Arc, Mutex, Weak}, 11 | time::{Duration, SystemTime, UNIX_EPOCH}, 12 | }; 13 | 14 | use rusqlite::{Connection, OptionalExtension}; 15 | 16 | #[derive(Clone)] 17 | pub struct Cache { 18 | inner: Arc, 19 | } 20 | 21 | #[derive(Clone, Debug)] 22 | pub struct CacheConfig { 23 | pub flush_interval: Duration, 24 | pub flush_gc_ratio: u64, 25 | pub max_ttl: Option, 26 | } 27 | 28 | impl Default for CacheConfig { 29 | fn default() -> Self { 30 | CacheConfig { 31 | flush_interval: Duration::from_secs(10), 32 | flush_gc_ratio: 30, 33 | max_ttl: None, 34 | } 35 | } 36 | } 37 | 38 | #[derive(Clone)] 39 | pub struct Topic { 40 | inner: Arc, 41 | } 42 | 43 | struct CacheImpl { 44 | config: CacheConfig, 45 | conn: Mutex, 46 | lazy_expiry_update: Mutex, String), u64>>, 47 | stop_tx: Mutex>, 48 | completion_rx: Mutex>, 49 | } 50 | 51 | struct TopicImpl { 52 | cache: Cache, 53 | table_name: Arc, 54 | listeners: Mutex>>>, 55 | } 56 | 57 | impl Drop for CacheImpl { 58 | fn drop(&mut self) { 59 | self.stop_tx.lock().unwrap().send(()).unwrap(); 60 | self.completion_rx.lock().unwrap().recv().unwrap(); 61 | } 62 | } 63 | 64 | impl Cache { 65 | pub fn new(config: CacheConfig, conn: Connection) -> Result { 66 | assert!(config.flush_gc_ratio > 0); 67 | let (stop_tx, stop_rx) = mpsc::channel::<()>(); 68 | let (completion_tx, completion_rx) = mpsc::channel::<()>(); 69 | conn.execute_batch("pragma journal_mode = wal;")?; 70 | let inner = Arc::new(CacheImpl { 71 | conn: Mutex::new(conn), 72 | config: config.clone(), 73 | lazy_expiry_update: Mutex::new(HashMap::new()), 74 | stop_tx: Mutex::new(stop_tx), 75 | completion_rx: Mutex::new(completion_rx), 76 | }); 77 | let w = Arc::downgrade(&inner); 78 | std::thread::spawn(move || periodic_task(config, stop_rx, completion_tx, w)); 79 | Ok(Self { inner }) 80 | } 81 | 82 | fn flush(&self) { 83 | let lazy_expiry_update = std::mem::take(&mut *self.inner.lazy_expiry_update.lock().unwrap()); 84 | for ((table_name, key), expiry) in lazy_expiry_update { 85 | let res = self.inner.conn.lock().unwrap().execute( 86 | &format!("update {} set expiry = ? where k = ?", table_name), 87 | rusqlite::params![expiry, key], 88 | ); 89 | if let Err(e) = res { 90 | tracing::error!(table = &*table_name, key = key.as_str(), error = %e, "error updating expiry"); 91 | } 92 | } 93 | } 94 | 95 | fn gc(&self) -> Result<(), rusqlite::Error> { 96 | let now = SystemTime::now() 97 | .duration_since(UNIX_EPOCH) 98 | .unwrap() 99 | .as_secs(); 100 | let tables = self 101 | .inner 102 | .conn 103 | .lock() 104 | .unwrap() 105 | .unchecked_transaction()? 106 | .prepare("select name from sqlite_master where type = 'table' and name like 'topic_%'")? 107 | .query_map(rusqlite::params![], |x| x.get::<_, String>(0))? 108 | .collect::, rusqlite::Error>>()?; 109 | let mut total = 0usize; 110 | for table in tables { 111 | let count = self.inner.conn.lock().unwrap().execute( 112 | &format!("delete from {} where expiry < ?", table), 113 | rusqlite::params![now], 114 | )?; 115 | total += count; 116 | } 117 | if total != 0 { 118 | tracing::info!(total = total, "gc deleted rows"); 119 | } 120 | Ok(()) 121 | } 122 | 123 | pub fn topic(&self, key: &str) -> Result { 124 | let table_name = format!("topic_{}", BASE32_NOPAD.encode(key.as_bytes())); 125 | self.inner.conn.lock().unwrap().execute_batch(&format!( 126 | r#" 127 | begin transaction; 128 | create table if not exists {} ( 129 | k text primary key not null, 130 | v blob not null, 131 | created_at integer not null default (cast(strftime('%s', 'now') as integer)), 132 | expiry integer not null, 133 | ttl integer not null 134 | ); 135 | create index if not exists {}_by_expiry on {} (expiry); 136 | commit; 137 | "#, 138 | table_name, table_name, table_name, 139 | ))?; 140 | Ok(Topic { 141 | inner: Arc::new(TopicImpl { 142 | cache: self.clone(), 143 | table_name: Arc::from(table_name), 144 | listeners: Mutex::new(HashMap::new()), 145 | }), 146 | }) 147 | } 148 | } 149 | 150 | pub struct Value { 151 | pub data: Vec, 152 | pub created_at: u64, 153 | } 154 | 155 | impl Topic { 156 | pub fn get(&self, key: &str) -> Result, rusqlite::Error> { 157 | let conn = self.inner.cache.inner.conn.lock().unwrap(); 158 | let mut stmt = conn.prepare_cached(&format!( 159 | "select v, created_at, ttl from {} where k = ?", 160 | self.inner.table_name, 161 | ))?; 162 | let rsp: Option<(Vec, u64, u64)> = stmt 163 | .query_row(rusqlite::params![key], |x| { 164 | Ok((x.get(0)?, x.get(1)?, x.get(2)?)) 165 | }) 166 | .optional()?; 167 | if let Some((data, created_at, ttl)) = rsp { 168 | self.inner 169 | .cache 170 | .inner 171 | .lazy_expiry_update 172 | .lock() 173 | .unwrap() 174 | .insert( 175 | (self.inner.table_name.clone(), key.to_string()), 176 | SystemTime::now() 177 | .duration_since(UNIX_EPOCH) 178 | .unwrap() 179 | .as_secs() 180 | .saturating_add(ttl) 181 | .min(i64::MAX as u64), 182 | ); 183 | Ok(Some(Value { data, created_at })) 184 | } else { 185 | Ok(None) 186 | } 187 | } 188 | 189 | pub async fn get_for_update( 190 | &self, 191 | key: &str, 192 | ) -> Result<(KeyUpdater, Option), rusqlite::Error> { 193 | loop { 194 | let receiver: Option>; 195 | { 196 | let mut listeners = self.inner.listeners.lock().unwrap(); 197 | if let Some(arr) = listeners.get_mut(key) { 198 | let (tx, rx) = channel(); 199 | arr.push(tx); 200 | receiver = Some(rx); 201 | } else { 202 | receiver = None; 203 | listeners.insert(key.to_string(), vec![]); 204 | } 205 | } 206 | 207 | if let Some(receiver) = receiver { 208 | let _ = receiver.await; 209 | } else { 210 | break; 211 | } 212 | } 213 | 214 | let data = self.get(key)?; 215 | Ok(( 216 | KeyUpdater { 217 | topic: self.clone(), 218 | key: key.to_string(), 219 | }, 220 | data, 221 | )) 222 | } 223 | 224 | pub fn set(&self, key: &str, value: &[u8], ttl: Duration) -> Result<(), rusqlite::Error> { 225 | let conn = self.inner.cache.inner.conn.lock().unwrap(); 226 | let mut stmt = conn.prepare_cached(&format!( 227 | "replace into {} (k, v, expiry, ttl) values(?, ?, ?, ?)", 228 | self.inner.table_name 229 | ))?; 230 | let mut ttl = ttl.as_secs(); 231 | if let Some(max_ttl) = self.inner.cache.inner.config.max_ttl { 232 | let max_ttl = max_ttl.as_secs(); 233 | ttl = ttl.min(max_ttl); 234 | } 235 | ttl = ttl.min(i64::MAX as u64); 236 | let expiry = SystemTime::now() 237 | .duration_since(UNIX_EPOCH) 238 | .unwrap() 239 | .as_secs() 240 | .saturating_add(ttl) 241 | .min(i64::MAX as u64); 242 | stmt.execute(rusqlite::params![key, value, expiry, ttl])?; 243 | self.inner 244 | .cache 245 | .inner 246 | .lazy_expiry_update 247 | .lock() 248 | .unwrap() 249 | .remove(&(self.inner.table_name.clone(), key.to_string())); 250 | Ok(()) 251 | } 252 | 253 | pub fn delete(&self, key: &str) -> Result<(), rusqlite::Error> { 254 | let conn = self.inner.cache.inner.conn.lock().unwrap(); 255 | let mut stmt = conn.prepare_cached(&format!( 256 | "delete from {} where k = ?", 257 | self.inner.table_name 258 | ))?; 259 | stmt.execute(rusqlite::params![key])?; 260 | Ok(()) 261 | } 262 | } 263 | 264 | pub struct KeyUpdater { 265 | topic: Topic, 266 | key: String, 267 | } 268 | 269 | impl Drop for KeyUpdater { 270 | fn drop(&mut self) { 271 | let mut listeners = self.topic.inner.listeners.lock().unwrap(); 272 | listeners.remove(self.key.as_str()).unwrap(); 273 | } 274 | } 275 | 276 | impl KeyUpdater { 277 | pub fn write(self, value: &[u8], ttl: Duration) -> Result<(), rusqlite::Error> { 278 | self.topic.set(&self.key, value, ttl)?; 279 | Ok(()) 280 | } 281 | } 282 | 283 | fn periodic_task( 284 | config: CacheConfig, 285 | stop_rx: mpsc::Receiver<()>, 286 | completion_tx: mpsc::Sender<()>, 287 | w: Weak, 288 | ) { 289 | let mut gc_ratio_counter = 0u64; 290 | loop { 291 | let tx = stop_rx.recv_timeout(config.flush_interval); 292 | if tx.is_ok() { 293 | break; 294 | } 295 | 296 | let inner = if let Some(x) = w.upgrade() { 297 | x 298 | } else { 299 | break; 300 | }; 301 | let cache = Cache { inner }; 302 | cache.flush(); 303 | gc_ratio_counter += 1; 304 | if gc_ratio_counter == config.flush_gc_ratio { 305 | gc_ratio_counter = 0; 306 | if let Err(e) = cache.gc() { 307 | tracing::error!(error = %e, "gc failed"); 308 | } 309 | } 310 | } 311 | tracing::info!("exiting periodic task"); 312 | completion_tx.send(()).unwrap(); 313 | } 314 | -------------------------------------------------------------------------------- /src/lib_test.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use rand::{thread_rng, Rng}; 4 | use rusqlite::Connection; 5 | 6 | use crate::{Cache, CacheConfig}; 7 | use tracing_test::traced_test; 8 | 9 | #[test] 10 | #[traced_test] 11 | fn test_new_drop() { 12 | let cache = Cache::new( 13 | CacheConfig::default(), 14 | Connection::open_in_memory().unwrap(), 15 | ) 16 | .unwrap(); 17 | drop(cache); 18 | } 19 | 20 | #[traced_test] 21 | #[tokio::test] 22 | async fn test_get_update() { 23 | let cache = Cache::new( 24 | CacheConfig::default(), 25 | Connection::open_in_memory().unwrap(), 26 | ) 27 | .unwrap(); 28 | let topic = cache.topic("test-topic").unwrap(); 29 | cache.topic("test-topic").unwrap(); // not used 30 | cache.topic("test-topic-2").unwrap(); // not used 31 | assert!(topic.get("hello").unwrap().is_none()); 32 | 33 | let (updater, value) = topic.get_for_update("hello").await.unwrap(); 34 | assert!(value.is_none()); 35 | updater.write(b"world", Duration::from_secs(1)).unwrap(); 36 | assert!(&topic.get("hello").unwrap().unwrap().data[..] == b"world"); 37 | topic.delete("hello").unwrap(); 38 | assert!(topic.get("hello").unwrap().is_none()); 39 | } 40 | 41 | #[traced_test] 42 | #[tokio::test] 43 | async fn test_update_lock() { 44 | let cache = Cache::new( 45 | CacheConfig::default(), 46 | Connection::open_in_memory().unwrap(), 47 | ) 48 | .unwrap(); 49 | let topic = cache.topic("test-topic").unwrap(); 50 | let (updater, _) = topic.get_for_update("hello").await.unwrap(); 51 | tokio::select! { 52 | _ = topic.get_for_update("hello") => { 53 | panic!("lock failed"); 54 | } 55 | _ = tokio::time::sleep(Duration::from_millis(100)) => { 56 | 57 | } 58 | } 59 | 60 | let attempt = { 61 | let topic = topic.clone(); 62 | tokio::task::spawn(async move { topic.get_for_update("hello").await.unwrap().1 }) 63 | }; 64 | tokio::time::sleep(Duration::from_millis(100)).await; 65 | updater.write(b"world", Duration::from_secs(1)).unwrap(); 66 | let value = attempt.await.unwrap().unwrap(); 67 | assert_eq!(&value.data[..], b"world"); 68 | } 69 | 70 | #[traced_test] 71 | #[tokio::test] 72 | async fn test_gc() { 73 | let cache = Cache::new( 74 | CacheConfig { 75 | flush_interval: Duration::from_millis(100), 76 | flush_gc_ratio: 5, 77 | ..CacheConfig::default() 78 | }, 79 | Connection::open_in_memory().unwrap(), 80 | ) 81 | .unwrap(); 82 | let topic = cache.topic("test-topic").unwrap(); 83 | let (updater, value) = topic.get_for_update("hello").await.unwrap(); 84 | assert!(value.is_none()); 85 | updater 86 | .write(b"world", Duration::from_millis(1000)) 87 | .unwrap(); 88 | assert!(&topic.get("hello").unwrap().unwrap().data[..] == b"world"); 89 | 90 | for _ in 0..20 { 91 | if thread_rng().gen_bool(0.5) { 92 | assert!(&topic.get("hello").unwrap().unwrap().data[..] == b"world"); 93 | } else { 94 | topic 95 | .set("hello", b"world", Duration::from_millis(1000)) 96 | .unwrap(); 97 | } 98 | tokio::time::sleep(Duration::from_millis(700)).await; 99 | } 100 | tokio::time::sleep(Duration::from_secs(3)).await; 101 | assert!(topic.get("hello").unwrap().is_none()); 102 | } 103 | 104 | // https://github.com/losfair/sqlite-cache/issues/1 105 | #[test] 106 | #[traced_test] 107 | fn test_large_durations() { 108 | let cache = Cache::new( 109 | CacheConfig::default(), 110 | Connection::open_in_memory().unwrap(), 111 | ) 112 | .unwrap(); 113 | 114 | let topic = cache.topic("test-topic").unwrap(); 115 | topic.set("key", b"value", Duration::MAX).unwrap(); 116 | } 117 | --------------------------------------------------------------------------------