Skip to content

refactor(spider-storage): Replace DashMap with tokio::sync::RwLock<HashMap> in JobCache to avoid holding shard locks across .await. #321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sitaowang1998 merged 1 commit into from
May 12, 2026
Merged

refactor(spider-storage): Replace DashMap with tokio::sync::RwLock<HashMap> in JobCache to avoid holding shard locks across .await. #321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion components/spider-storage/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ path = "tests/test_spider_storage.rs"
async-channel = "2.3.1"
async-trait = "0.1.89"
const_format = "0.2.35"
dashmap = "6.1.0"
rmp-serde = "1.3.1"
secrecy = { version = "0.10.3", features = ["serde"] }
serde = { version = "1.0.228", features = ["derive"] }
Expand All @@ -35,6 +34,7 @@ uuid = { version = "1.19.0", features = ["serde"] }

[dev-dependencies]
anyhow = "1.0.98"
dashmap = "6.1.0"
rand = "0.9.1"
serial_test = { version = "3.2.0", features = ["file_locks"] }
tabled = "0.20.0"
Expand Down
65 changes: 36 additions & 29 deletions components/spider-storage/src/state/job_cache.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
use dashmap::{DashMap, mapref::entry::Entry};
use std::collections::{HashMap, hash_map::Entry};

use spider_core::types::id::JobId;
use tokio::sync::RwLock;

use crate::{
cache::job::SharedJobControlBlock,
Expand All @@ -11,8 +13,8 @@ use crate::{

/// An in-memory cache for job control blocks.
///
/// This type provides concurrent access to job control blocks via a `DashMap`. It is generic over
/// the same type parameters as [`SharedJobControlBlock`].
/// This type provides concurrent access to job control blocks via a [`tokio::sync::RwLock`] over a
/// [`HashMap`]. It is generic over the same type parameters as [`SharedJobControlBlock`].
///
/// # Type Parameters
///
Expand All @@ -24,9 +26,15 @@ pub struct JobCache<
DbConnectorType: InternalJobOrchestration,
TaskInstancePoolConnectorType: TaskInstancePoolConnector,
> {
jobs: DashMap<
JobId,
SharedJobControlBlock<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType>,
jobs: RwLock<
HashMap<
JobId,
SharedJobControlBlock<
ReadyQueueSenderType,
DbConnectorType,
TaskInstancePoolConnectorType,
>,
>,
>,
}

Expand All @@ -40,7 +48,7 @@ impl<
#[must_use]
pub fn new() -> Self {
Self {
jobs: DashMap::new(),
jobs: RwLock::new(HashMap::new()),
}
}

Expand All @@ -52,7 +60,7 @@ impl<
///
/// * [`StorageServerError::JobAlreadyExists`] if a job control block with the same ID already
/// exists.
pub fn insert(
pub async fn insert(
&self,
jcb: SharedJobControlBlock<
ReadyQueueSenderType,
Expand All @@ -61,7 +69,7 @@ impl<
>,
) -> Result<(), StorageServerError> {
let job_id = jcb.id();
match self.jobs.entry(job_id) {
match self.jobs.write().await.entry(job_id) {
Entry::Vacant(e) => {
e.insert(jcb);
Ok(())
Expand All @@ -75,29 +83,27 @@ impl<
/// # Returns
///
/// The job control block of the given ID if it exists, [`None`] otherwise.
#[must_use]
pub fn get(
pub async fn get(
&self,
job_id: JobId,
) -> Option<
SharedJobControlBlock<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType>,
> {
self.jobs.get(&job_id).map(|entry| entry.clone())
self.jobs.read().await.get(&job_id).cloned()
}

/// Removes a job control block from the cache.
///
/// # Returns
///
/// The removed job control block if it existed, [`None`] otherwise.
#[must_use]
pub fn remove(
pub async fn remove(
&self,
job_id: JobId,
) -> Option<
SharedJobControlBlock<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType>,
> {
self.jobs.remove(&job_id).map(|(_, v)| v)
self.jobs.write().await.remove(&job_id)
}

/// Resends all ready tasks for every job in the cache to the ready queue.
Expand All @@ -108,8 +114,8 @@ impl<
///
/// * Forwards [`SharedJobControlBlock::resend_ready_tasks`]'s return values on failure.
pub async fn resend_ready_tasks(&self) -> Result<(), StorageServerError> {
for entry in &self.jobs {
entry.value().resend_ready_tasks().await?;
for jcb in self.jobs.read().await.values() {
jcb.resend_ready_tasks().await?;
}
Ok(())
}
Expand Down Expand Up @@ -299,9 +305,9 @@ mod tests {
let job_id = JobId::new();

let jcb = create_test_jcb(job_id).await;
cache.insert(jcb)?;
cache.insert(jcb).await?;

let result = cache.get(job_id);
let result = cache.get(job_id).await;
assert!(result.is_some(), "inserted JCB should be retrievable");
Ok(())
}
Expand All @@ -313,12 +319,12 @@ mod tests {
let job_id = JobId::new();

let jcb = create_test_jcb(job_id).await;
cache.insert(jcb)?;
cache.insert(jcb).await?;

let removed = cache.remove(job_id);
let removed = cache.remove(job_id).await;
assert!(removed.is_some(), "remove should return the JCB");

let result = cache.get(job_id);
let result = cache.get(job_id).await;
assert!(result.is_none(), "JCB should no longer exist after removal");
Ok(())
}
Expand All @@ -329,7 +335,7 @@ mod tests {
JobCache::new();
let job_id = JobId::new();

let result = cache.get(job_id);
let result = cache.get(job_id).await;
assert!(
result.is_none(),
"get should return None for nonexistent job"
Expand All @@ -344,10 +350,10 @@ mod tests {
let job_id = JobId::new();

let jcb1 = create_test_jcb(job_id).await;
cache.insert(jcb1)?;
cache.insert(jcb1).await?;

let jcb2 = create_test_jcb(job_id).await;
let result = cache.insert(jcb2);
let result = cache.insert(jcb2).await;
assert!(
matches!(result, Err(StorageServerError::JobAlreadyExists(_))),
"insert should return JobAlreadyExists error for duplicate key"
Expand Down Expand Up @@ -376,15 +382,16 @@ mod tests {
let jcb = create_test_jcb(job_id).await;
cache
.insert(jcb)
.await
.expect("insert should succeed for new job");

let result = cache.get(job_id);
let result = cache.get(job_id).await;
assert!(result.is_some(), "task {i} should find inserted JCB");

let removed = cache.remove(job_id);
let removed = cache.remove(job_id).await;
assert!(removed.is_some(), "task {i} should remove inserted JCB");

let result = cache.get(job_id);
let result = cache.get(job_id).await;
assert!(result.is_none(), "task {i} should not find removed JCB");
});
}
Expand Down Expand Up @@ -478,7 +485,7 @@ mod tests {
MockDbConnector,
MockTaskInstancePoolConnector,
> = JobCache::new();
cache.insert(jcb)?;
cache.insert(jcb).await?;

cache.resend_ready_tasks().await?;

Expand Down
Loading