Skip to content

Commit

Permalink
Clean up factor a bit
Browse files Browse the repository at this point in the history
Signed-off-by: Ryan Levick <[email protected]>
  • Loading branch information
rylev committed Sep 18, 2024
1 parent 2bc860d commit 5db8565
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 40 deletions.
26 changes: 9 additions & 17 deletions crates/factor-key-value/src/host.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
use crate::util::EmptyStoreManager;
use anyhow::{Context, Result};
use spin_core::{async_trait, wasmtime::component::Resource};
use spin_world::v2::key_value;
Expand Down Expand Up @@ -40,23 +39,22 @@ pub struct KeyValueDispatch {
}

impl KeyValueDispatch {
pub fn new() -> Self {
Self::new_with_capacity(DEFAULT_STORE_TABLE_CAPACITY)
pub fn new(allowed_stores: HashSet<String>, manager: Arc<dyn StoreManager>) -> Self {
Self::new_with_capacity(allowed_stores, manager, DEFAULT_STORE_TABLE_CAPACITY)
}

pub fn new_with_capacity(capacity: u32) -> Self {
pub fn new_with_capacity(
allowed_stores: HashSet<String>,
manager: Arc<dyn StoreManager>,
capacity: u32,
) -> Self {
Self {
allowed_stores: HashSet::new(),
manager: Arc::new(EmptyStoreManager),
allowed_stores,
manager,
stores: Table::new(capacity),
}
}

pub fn init(&mut self, allowed_stores: HashSet<String>, manager: Arc<dyn StoreManager>) {
self.allowed_stores = allowed_stores;
self.manager = manager;
}

pub fn get_store(&self, store: Resource<key_value::Store>) -> anyhow::Result<&Arc<dyn Store>> {
self.stores.get(store.rep()).context("invalid store")
}
Expand All @@ -66,12 +64,6 @@ impl KeyValueDispatch {
}
}

impl Default for KeyValueDispatch {
fn default() -> Self {
Self::new()
}
}

#[async_trait]
impl key_value::Host for KeyValueDispatch {}

Expand Down
14 changes: 8 additions & 6 deletions crates/factor-key-value/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ impl Factor for KeyValueFactor {

let delegating_manager = DelegatingStoreManager::new(store_managers);
let caching_manager = CachingStoreManager::new(delegating_manager);
let store_manager_manager = Arc::new(caching_manager);
let store_manager = Arc::new(caching_manager);

// Build component -> allowed stores map
let mut component_allowed_stores = HashMap::new();
Expand All @@ -65,7 +65,7 @@ impl Factor for KeyValueFactor {
for label in &key_value_stores {
// TODO: port nicer errors from KeyValueComponent (via error type?)
ensure!(
store_manager_manager.is_defined(label),
store_manager.is_defined(label),
"unknown key_value_stores label {label:?} for component {component_id:?}"
);
}
Expand All @@ -74,7 +74,7 @@ impl Factor for KeyValueFactor {
}

Ok(AppState {
store_manager: store_manager_manager,
store_manager,
component_allowed_stores,
})
}
Expand Down Expand Up @@ -150,8 +150,10 @@ impl FactorInstanceBuilder for InstanceBuilder {
store_manager,
allowed_stores,
} = self;
let mut dispatch = KeyValueDispatch::new_with_capacity(u32::MAX);
dispatch.init(allowed_stores, store_manager);
Ok(dispatch)
Ok(KeyValueDispatch::new_with_capacity(
allowed_stores,
store_manager,
u32::MAX,
))
}
}
18 changes: 3 additions & 15 deletions crates/factor-key-value/src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,21 +13,7 @@ use tokio::{
};
use tracing::Instrument;

const DEFAULT_CACHE_SIZE: usize = 256;

pub struct EmptyStoreManager;

#[async_trait]
impl StoreManager for EmptyStoreManager {
async fn get(&self, _name: &str) -> Result<Arc<dyn Store>, Error> {
Err(Error::NoSuchStore)
}

fn is_defined(&self, _store_name: &str) -> bool {
false
}
}

/// A [`StoreManager`] which delegates to other `StoreManager`s based on the store label.
pub struct DelegatingStoreManager {
delegates: HashMap<String, Arc<dyn StoreManager>>,
}
Expand Down Expand Up @@ -89,6 +75,8 @@ pub struct CachingStoreManager<T> {
inner: T,
}

const DEFAULT_CACHE_SIZE: usize = 256;

impl<T> CachingStoreManager<T> {
pub fn new(inner: T) -> Self {
Self::new_with_capacity(NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(), inner)
Expand Down
3 changes: 1 addition & 2 deletions crates/key-value-spin/src/store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -159,8 +159,7 @@ mod test {

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn all() -> Result<()> {
let mut kv = KeyValueDispatch::new();
kv.init(
let mut kv = KeyValueDispatch::new(
["default", "foo"]
.into_iter()
.map(ToOwned::to_owned)
Expand Down

0 comments on commit 5db8565

Please sign in to comment.