aws_smithy_runtime/static_partition_map.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/
use once_cell::sync::OnceCell;
use std::collections::HashMap;
use std::hash::Hash;
use std::sync::{Mutex, MutexGuard};
/// A data structure for persisting and sharing state between multiple clients.
///
/// Some state should be shared between multiple clients. For example, when creating multiple clients
/// for the same service, it's desirable to share a client rate limiter. This way, when one client
/// receives a throttling response, the other clients will be aware of it as well.
///
/// Whether clients share state is dependent on their partition key `K`. Going back to the client
/// rate limiter example, `K` would be a struct containing the name of the service as well as the
/// client's configured region, since receiving throttling responses in `us-east-1` shouldn't
/// throttle requests to the same service made in other regions.
///
/// Values stored in a `StaticPartitionMap` will be cloned whenever they are requested. Values must
/// be initialized before they can be retrieved, and the `StaticPartitionMap::get_or_init` method is
/// how you can ensure this.
///
/// # Example
///
/// ```
///use std::sync::{Arc, Mutex};
/// use aws_smithy_runtime::static_partition_map::StaticPartitionMap;
///
/// // The shared state must be `Clone` and will be internally mutable. Deriving `Default` isn't
/// // necessary, but allows us to use the `StaticPartitionMap::get_or_init_default` method.
/// #[derive(Clone, Default)]
/// pub struct SomeSharedState {
/// inner: Arc<Mutex<Inner>>
/// }
///
/// #[derive(Default)]
/// struct Inner {
/// // Some shared state...
/// }
///
/// // `Clone`, `Hash`, and `Eq` are all required trait impls for partition keys
/// #[derive(Clone, Hash, PartialEq, Eq)]
/// pub struct SharedStatePartition {
/// region: String,
/// service_name: String,
/// }
///
/// impl SharedStatePartition {
/// pub fn new(region: impl Into<String>, service_name: impl Into<String>) -> Self {
/// Self { region: region.into(), service_name: service_name.into() }
/// }
/// }
///
/// static SOME_SHARED_STATE: StaticPartitionMap<SharedStatePartition, SomeSharedState> = StaticPartitionMap::new();
///
/// struct Client {
/// shared_state: SomeSharedState,
/// }
///
/// impl Client {
/// pub fn new() -> Self {
/// let key = SharedStatePartition::new("us-east-1", "example_service_20230628");
/// Self {
/// // If the stored value implements `Default`, you can call the
/// // `StaticPartitionMap::get_or_init_default` convenience method.
/// shared_state: SOME_SHARED_STATE.get_or_init_default(key),
/// }
/// }
/// }
/// ```
#[derive(Debug, Default)]
pub struct StaticPartitionMap<K, V> {
inner: OnceCell<Mutex<HashMap<K, V>>>,
}
impl<K, V> StaticPartitionMap<K, V> {
/// Creates a new `StaticPartitionMap`.
pub const fn new() -> Self {
Self {
inner: OnceCell::new(),
}
}
}
impl<K, V> StaticPartitionMap<K, V>
where
K: Eq + Hash,
{
fn get_or_init_inner(&self) -> MutexGuard<'_, HashMap<K, V>> {
self.inner
// At the very least, we'll always be storing the default state.
.get_or_init(|| Mutex::new(HashMap::with_capacity(1)))
.lock()
.unwrap()
}
}
impl<K, V> StaticPartitionMap<K, V>
where
K: Eq + Hash,
V: Clone,
{
/// Gets the value for the given partition key.
#[must_use]
pub fn get(&self, partition_key: K) -> Option<V> {
self.get_or_init_inner().get(&partition_key).cloned()
}
/// Gets the value for the given partition key, initializing it with `init` if it doesn't exist.
#[must_use]
pub fn get_or_init<F>(&self, partition_key: K, init: F) -> V
where
F: FnOnce() -> V,
{
let mut inner = self.get_or_init_inner();
let v = inner.entry(partition_key).or_insert_with(init);
v.clone()
}
}
impl<K, V> StaticPartitionMap<K, V>
where
K: Eq + Hash,
V: Clone + Default,
{
/// Gets the value for the given partition key, initializing it if it doesn't exist.
#[must_use]
pub fn get_or_init_default(&self, partition_key: K) -> V {
self.get_or_init(partition_key, V::default)
}
}
#[cfg(test)]
mod tests {
use super::StaticPartitionMap;
#[test]
fn test_keyed_partition_returns_same_value_for_same_key() {
let kp = StaticPartitionMap::new();
let _ = kp.get_or_init("A", || "A".to_owned());
let actual = kp.get_or_init("A", || "B".to_owned());
let expected = "A".to_owned();
assert_eq!(expected, actual);
}
#[test]
fn test_keyed_partition_returns_different_value_for_different_key() {
let kp = StaticPartitionMap::new();
let _ = kp.get_or_init("A", || "A".to_owned());
let actual = kp.get_or_init("B", || "B".to_owned());
let expected = "B".to_owned();
assert_eq!(expected, actual);
let actual = kp.get("A").unwrap();
let expected = "A".to_owned();
assert_eq!(expected, actual);
}
}