jcoffey-dev
7dae9b29fd
Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f Enterprise-only files removed or emptied: 63 Enterprise-only snippets removed: 117 in 50 files Dangling module declarations removed: 5 Cargo edits turning enterprise off: 14 Verification: clean Enterprise feature gates left for rebuilt features: 19 in 18 files Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.
301 lines
8.6 KiB
Rust
301 lines
8.6 KiB
Rust
/*
|
|
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
|
*
|
|
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
|
*/
|
|
|
|
use crate::utils::{
|
|
cleanup::{store_assert_is_empty, store_destroy},
|
|
server::TestServerBuilder,
|
|
};
|
|
use registry::schema::structs::Rate;
|
|
use registry::types::duration::Duration;
|
|
use store::{InMemoryStore, dispatch::lookup::KeyValue};
|
|
|
|
#[tokio::test]
|
|
pub async fn lookup_tests() {
|
|
let test = TestServerBuilder::new("lookup_tests").await.build().await;
|
|
let store = test.server.in_memory_store().clone();
|
|
let rate = Rate {
|
|
count: 1,
|
|
period: Duration::from_millis(1000),
|
|
};
|
|
|
|
println!(
|
|
"Testing in-memory store {}...",
|
|
std::env::var("MEMORY_STORE").unwrap_or_else(|_| "default".to_string())
|
|
);
|
|
if let InMemoryStore::Store(store) = &store {
|
|
store_destroy(store).await;
|
|
} else {
|
|
// Reset redis counter
|
|
store
|
|
.key_set(KeyValue::new("abc", "0".as_bytes().to_vec()))
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
// Test key
|
|
let key = "xyz".as_bytes().to_vec();
|
|
store
|
|
.key_set(KeyValue::new(key.clone(), "world".to_string().into_bytes()))
|
|
.await
|
|
.unwrap();
|
|
store.purge_in_memory_store().await.unwrap();
|
|
assert_eq!(
|
|
store.key_get::<String>(key.clone()).await.unwrap(),
|
|
Some("world".to_string())
|
|
);
|
|
|
|
// Test value expiry
|
|
store
|
|
.key_set(KeyValue::new(key.clone(), "hello".to_string().into_bytes()).expires(1))
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(
|
|
store.key_get::<String>(key.clone()).await.unwrap(),
|
|
Some("hello".to_string())
|
|
);
|
|
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
|
|
assert_eq!(None, store.key_get::<String>(key.clone()).await.unwrap());
|
|
|
|
store.purge_in_memory_store().await.unwrap();
|
|
if let InMemoryStore::Store(store) = &store {
|
|
store_assert_is_empty(store, store.clone().into(), false).await;
|
|
}
|
|
|
|
// Test counter
|
|
let key = "abc".as_bytes().to_vec();
|
|
store
|
|
.counter_incr(KeyValue::new(key.clone(), 1), true)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
|
|
store
|
|
.counter_incr(KeyValue::new(key.clone(), 2), true)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(3, store.counter_get(key.clone()).await.unwrap());
|
|
store
|
|
.counter_incr(KeyValue::new(key.clone(), -3), false)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
|
|
|
|
// Test counter expiry
|
|
let key = "fgh".as_bytes().to_vec();
|
|
store
|
|
.counter_incr(KeyValue::new(key.clone(), 1).expires(1), false)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
|
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
|
store.purge_in_memory_store().await.unwrap();
|
|
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
|
|
|
|
// Test rate limiter
|
|
assert!(
|
|
store
|
|
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
|
|
.await
|
|
.unwrap()
|
|
.is_none()
|
|
);
|
|
assert!(
|
|
store
|
|
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
|
|
.await
|
|
.unwrap()
|
|
.is_some()
|
|
);
|
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
|
assert!(
|
|
store
|
|
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
|
|
.await
|
|
.unwrap()
|
|
.is_none()
|
|
);
|
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
|
store.purge_in_memory_store().await.unwrap();
|
|
if let InMemoryStore::Store(store) = &store {
|
|
store_assert_is_empty(store, store.clone().into(), false).await;
|
|
}
|
|
|
|
// Test locking
|
|
for iteration in [1, 2] {
|
|
let mut tasks = Vec::new();
|
|
for _ in 0..100 {
|
|
let store = store.clone();
|
|
tasks.push(tokio::spawn(async move {
|
|
store.try_lock(0, "lock".as_bytes(), 1).await.unwrap()
|
|
}));
|
|
}
|
|
// Only one should return true
|
|
let mut count = 0;
|
|
for task in tasks {
|
|
if task.await.unwrap() {
|
|
count += 1;
|
|
}
|
|
}
|
|
assert_eq!(1, count, "Iteration {}", iteration);
|
|
|
|
// Wait 2 seconds for the lock to expire
|
|
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
|
|
}
|
|
store.purge_in_memory_store().await.unwrap();
|
|
|
|
// Failed lock attempts must not extend the lock expiry
|
|
assert!(store.try_lock(0, "lock".as_bytes(), 2).await.unwrap());
|
|
let mut acquired = false;
|
|
for _ in 0..12 {
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
|
if store.try_lock(0, "lock".as_bytes(), 2).await.unwrap() {
|
|
acquired = true;
|
|
break;
|
|
}
|
|
}
|
|
assert!(acquired, "Abandoned lock was never released");
|
|
store.remove_lock(0, "lock".as_bytes()).await.unwrap();
|
|
store.purge_in_memory_store().await.unwrap();
|
|
if let InMemoryStore::Store(store) = &store {
|
|
store_assert_is_empty(store, store.clone().into(), false).await;
|
|
}
|
|
|
|
// Test prefix delete
|
|
store
|
|
.key_set(KeyValue::with_prefix(
|
|
1,
|
|
[0],
|
|
"hello".to_string().into_bytes(),
|
|
))
|
|
.await
|
|
.unwrap();
|
|
for v in 0u32..2020u32 {
|
|
store
|
|
.key_set(KeyValue::with_prefix(
|
|
0,
|
|
pack_u32(0, v),
|
|
"world".to_string().into_bytes(),
|
|
))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.counter_incr(
|
|
KeyValue::with_prefix(0, pack_u32(1, v), 123).expires(3600),
|
|
false,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
// Make sure the keys are there
|
|
assert_eq!(
|
|
Some("hello"),
|
|
store
|
|
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
|
|
.await
|
|
.unwrap()
|
|
.as_deref()
|
|
);
|
|
for v in [0, 1000, 1001, 2000, 2001] {
|
|
assert_eq!(
|
|
Some("world"),
|
|
store
|
|
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
|
|
.await
|
|
.unwrap()
|
|
.as_deref()
|
|
);
|
|
}
|
|
for v in [0, 1000, 1001, 2000, 2001] {
|
|
assert_ne!(
|
|
0,
|
|
store
|
|
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
|
|
.await
|
|
.unwrap()
|
|
);
|
|
}
|
|
|
|
// Delete [0, 0, 0, 0, 1] prefix and make sure only the keys with that prefix are gone
|
|
store
|
|
.key_delete_prefix(&KeyValue::<()>::build_key(0, 1u32.to_be_bytes()))
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
Some("hello"),
|
|
store
|
|
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
|
|
.await
|
|
.unwrap()
|
|
.as_deref()
|
|
);
|
|
for v in [0, 1000, 1001, 2000, 2001] {
|
|
assert_eq!(
|
|
Some("world"),
|
|
store
|
|
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
|
|
.await
|
|
.unwrap()
|
|
.as_deref()
|
|
);
|
|
}
|
|
|
|
for v in [0, 1000, 1001, 2000, 2001] {
|
|
assert_eq!(
|
|
0,
|
|
store
|
|
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
|
|
.await
|
|
.unwrap()
|
|
);
|
|
}
|
|
|
|
// Delete [0, 0, 0, 0, 0] prefix and make sure only the keys with that prefix are gone
|
|
store
|
|
.key_delete_prefix(&KeyValue::<()>::build_key(0, 0u32.to_be_bytes()))
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
Some("hello"),
|
|
store
|
|
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
|
|
.await
|
|
.unwrap()
|
|
.as_deref()
|
|
);
|
|
for v in [0, 1000, 1001, 2000, 2001] {
|
|
assert_eq!(
|
|
None,
|
|
store
|
|
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
|
|
.await
|
|
.unwrap()
|
|
.as_deref()
|
|
);
|
|
}
|
|
|
|
// Delete [1, ...] prefix and make sure it's all gone
|
|
store.key_delete_prefix(&[1u8]).await.unwrap();
|
|
|
|
assert_eq!(
|
|
None,
|
|
store
|
|
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
|
|
.await
|
|
.unwrap()
|
|
.as_deref()
|
|
);
|
|
|
|
if let InMemoryStore::Store(store) = &store {
|
|
store_assert_is_empty(store, store.clone().into(), false).await;
|
|
}
|
|
}
|
|
|
|
fn pack_u32(a: u32, b: u32) -> Vec<u8> {
|
|
(((a as u64) << 32) | b as u64).to_be_bytes().to_vec()
|
|
}
|