fix(core): stop concurrent manifest writes from losing a served file
A snapshot_download fans out across a repo's files, so several transfers finish at once and each rewrites the same manifest document. The read-merge-write had no conditional put behind it: two writers that read the same document both wrote their own version, and the second silently dropped the first's entry. The verify-after-write retry caught most of that, but the write happens in the detached transfer task -- after the client already has its bytes -- so a lost entry surfaced only later, as a sealed instance claiming never to have seen a file it had served. That breaks the fetch-once retention guarantee. Adds a compare-and-set pair to the Store port (get_versioned / put_if_version, conditional on the etag last read) and uses it for manifest updates, so a racing write is refused and retried rather than merged away. The local filesystem store cannot do conditional updates, so it reports supports_compare_and_set() = false and the caller keeps the verify-after-write path for it. Also serialises manifest writes per key within the process. CAS alone is correct but turns eight parallel workers into eight writers contending for one document, burning the retry budget on self-inflicted collisions; the lock makes the common case uncontended and leaves the retry loop for genuinely concurrent writers -- another instance, or the CLI against the same bucket. Regression test fetches twelve files of one repo concurrently and asserts every one survives into the manifest and replays when sealed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XZG2i4AmfSqE97EJGBVb64
This commit is contained in:
@@ -384,6 +384,69 @@ async fn a_ranged_request_resumes_from_the_offset_it_asks_for() {
|
||||
assert_eq!(status, StatusCode::RANGE_NOT_SATISFIABLE);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_fetches_of_different_files_all_reach_the_manifest() {
|
||||
let hub = FakeHub::start().await;
|
||||
// A snapshot_download fans out across a repo's files, so every one of
|
||||
// these lands in the same manifest document at the same moment.
|
||||
let files: Vec<(String, HubFile)> = (0..12)
|
||||
.map(|i| {
|
||||
(
|
||||
format!("shard-{i:02}.safetensors"),
|
||||
HubFile::lfs(format!("shard {i} bytes").as_bytes()),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let published: Vec<(&str, HubFile)> = files
|
||||
.iter()
|
||||
.map(|(path, file)| (path.as_str(), file.clone()))
|
||||
.collect();
|
||||
hub.publish("Qwen/Qwen3-32B", "main", COMMIT_A, &published)
|
||||
.await;
|
||||
|
||||
let bucket = bucket_dir("manifest-race");
|
||||
let rf = instance(&bucket, Some(&hub), |_| {});
|
||||
|
||||
let mut clients = Vec::new();
|
||||
for (path, _) in &files {
|
||||
let app = rf.app.clone();
|
||||
let uri = format!("/Qwen/Qwen3-32B/resolve/main/{path}");
|
||||
clients.push(tokio::spawn(
|
||||
async move { call(&app, "GET", &uri, &[]).await },
|
||||
));
|
||||
}
|
||||
for client in clients {
|
||||
assert_eq!(client.await.unwrap().0, StatusCode::OK);
|
||||
}
|
||||
|
||||
// A file that was served but never recorded is invisible the moment the
|
||||
// instance is sealed, which is the failure this guards against.
|
||||
let manifest = rf
|
||||
.registry
|
||||
.load_manifest(RepoType::Models, &qwen(), COMMIT_A)
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("a manifest");
|
||||
assert_eq!(
|
||||
manifest.entries.len(),
|
||||
files.len(),
|
||||
"every served file must survive the concurrent manifest writes"
|
||||
);
|
||||
|
||||
let sealed = instance(&bucket, None, |_| {});
|
||||
for (path, file) in &files {
|
||||
let (status, _, body) = call(
|
||||
&sealed.app,
|
||||
"GET",
|
||||
&format!("/Qwen/Qwen3-32B/resolve/main/{path}"),
|
||||
&[],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK, "{path}");
|
||||
assert_eq!(body, file.body, "{path}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_gated_repo_answers_403_with_the_code_the_client_branches_on() {
|
||||
let hub = FakeHub::start().await;
|
||||
|
||||
@@ -23,6 +23,16 @@ pub struct StoredObject {
|
||||
pub size: u64,
|
||||
}
|
||||
|
||||
/// Opaque token identifying the version of an object that was read.
|
||||
///
|
||||
/// Backends key conditional writes off differing combinations of etag and
|
||||
/// version id, so both are carried and neither is interpreted here.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct Version {
|
||||
pub e_tag: Option<String>,
|
||||
pub version: Option<String>,
|
||||
}
|
||||
|
||||
/// A readable blob: the byte stream plus what the store said about it.
|
||||
pub struct BlobRead {
|
||||
/// Bytes of the requested range, in order.
|
||||
@@ -54,6 +64,30 @@ pub trait Store: Send + Sync + 'static {
|
||||
/// Fetch a small object whole. `Ok(None)` when it does not exist.
|
||||
async fn get(&self, key: &str) -> Result<Option<Bytes>>;
|
||||
|
||||
/// As [`Store::get`], but also returning the version token needed to
|
||||
/// compare-and-set the object back.
|
||||
async fn get_versioned(&self, key: &str) -> Result<Option<(Bytes, Version)>>;
|
||||
|
||||
/// Write only if the object still has the version last read.
|
||||
///
|
||||
/// `version` of `None` means "only if absent". Returns `false` when
|
||||
/// somebody else wrote first, which is a retry, not an error.
|
||||
///
|
||||
/// Backends that cannot do this ([`Store::supports_compare_and_set`] is
|
||||
/// `false`) overwrite unconditionally and always return `true`; callers
|
||||
/// must then verify their write survived rather than trusting it.
|
||||
async fn put_if_version(
|
||||
&self,
|
||||
key: &str,
|
||||
body: Bytes,
|
||||
version: Option<&Version>,
|
||||
) -> Result<bool>;
|
||||
|
||||
/// Whether [`Store::put_if_version`] is genuinely conditional.
|
||||
fn supports_compare_and_set(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Overwrite an object.
|
||||
async fn put(&self, key: &str, body: Bytes) -> Result<()>;
|
||||
|
||||
|
||||
@@ -89,6 +89,15 @@ struct Inner {
|
||||
policy: Policy,
|
||||
config: Config,
|
||||
flights: Flights,
|
||||
/// One lock per manifest key.
|
||||
///
|
||||
/// Compare-and-set alone is correct but turns a `snapshot_download`'s
|
||||
/// eight parallel workers into eight writers contending for one document,
|
||||
/// which burns the retry budget on self-inflicted collisions. Serialising
|
||||
/// within the process makes the common case uncontended and leaves the
|
||||
/// retry loop for the genuinely concurrent case: another rustingface, or
|
||||
/// an operator running the CLI against the same bucket.
|
||||
manifest_locks: dashmap::DashMap<String, Arc<tokio::sync::Mutex<()>>>,
|
||||
}
|
||||
|
||||
impl Registry {
|
||||
@@ -117,6 +126,7 @@ impl Registry {
|
||||
policy,
|
||||
config,
|
||||
flights: Flights::new(),
|
||||
manifest_locks: dashmap::DashMap::new(),
|
||||
}),
|
||||
})
|
||||
}
|
||||
@@ -345,41 +355,76 @@ impl Registry {
|
||||
F: FnMut(&mut Manifest),
|
||||
{
|
||||
let manifest_key = key::manifest(repo_type, repo_id, commit);
|
||||
let lock = Arc::clone(
|
||||
&self
|
||||
.inner
|
||||
.manifest_locks
|
||||
.entry(manifest_key.clone())
|
||||
.or_default(),
|
||||
);
|
||||
let _guard = lock.lock().await;
|
||||
|
||||
let mut attempt = 0;
|
||||
loop {
|
||||
let mut manifest = self
|
||||
.load_manifest(repo_type, repo_id, commit)
|
||||
.await?
|
||||
.unwrap_or_else(|| {
|
||||
Manifest::new(
|
||||
repo_type,
|
||||
repo_id.clone(),
|
||||
commit.to_owned(),
|
||||
Some(self.inner.upstream.endpoint().to_owned()),
|
||||
)
|
||||
});
|
||||
let (mut manifest, version) =
|
||||
match self.inner.store.get_versioned(&manifest_key).await? {
|
||||
Some((bytes, version)) => {
|
||||
let manifest: Manifest = serde_json::from_slice(&bytes).map_err(|e| {
|
||||
Error::Storage(format!("{manifest_key} is not a valid manifest: {e}"))
|
||||
})?;
|
||||
(manifest, Some(version))
|
||||
}
|
||||
None => (
|
||||
Manifest::new(
|
||||
repo_type,
|
||||
repo_id.clone(),
|
||||
commit.to_owned(),
|
||||
Some(self.inner.upstream.endpoint().to_owned()),
|
||||
),
|
||||
None,
|
||||
),
|
||||
};
|
||||
mutate(&mut manifest);
|
||||
self.inner
|
||||
|
||||
let body = Bytes::from(encode_json(&manifest)?);
|
||||
let won = self
|
||||
.inner
|
||||
.store
|
||||
.put(&manifest_key, Bytes::from(encode_json(&manifest)?))
|
||||
.put_if_version(&manifest_key, body, version.as_ref())
|
||||
.await?;
|
||||
|
||||
// Confirm our contribution survived a concurrent write.
|
||||
let landed = self.load_manifest(repo_type, repo_id, commit).await?;
|
||||
let lost = match &landed {
|
||||
Some(current) => manifest
|
||||
.entries
|
||||
.iter()
|
||||
.any(|e| current.entry(&e.path).is_none()),
|
||||
None => true,
|
||||
};
|
||||
if !lost {
|
||||
return Ok(landed.unwrap_or(manifest));
|
||||
if won && self.inner.store.supports_compare_and_set() {
|
||||
// The write was conditional on what we read, so nothing can
|
||||
// have been merged away.
|
||||
self.inner.manifest_locks.remove(&manifest_key);
|
||||
return Ok(manifest);
|
||||
}
|
||||
|
||||
if won {
|
||||
// No conditional write available: confirm our contribution
|
||||
// actually survived rather than assuming it did. This is the
|
||||
// difference between a served file being recorded and a
|
||||
// sealed instance later claiming never to have seen it.
|
||||
let landed = self.load_manifest(repo_type, repo_id, commit).await?;
|
||||
let intact = landed.as_ref().is_some_and(|current| {
|
||||
manifest
|
||||
.entries
|
||||
.iter()
|
||||
.all(|e| current.entry(&e.path).is_some())
|
||||
});
|
||||
if intact {
|
||||
self.inner.manifest_locks.remove(&manifest_key);
|
||||
return Ok(landed.unwrap_or(manifest));
|
||||
}
|
||||
}
|
||||
|
||||
attempt += 1;
|
||||
if attempt > self.inner.config.storage.manifest_write_retries {
|
||||
self.inner.manifest_locks.remove(&manifest_key);
|
||||
return Err(Error::Storage(format!(
|
||||
"{manifest_key} kept losing entries to concurrent writers after {attempt} attempts"
|
||||
"{manifest_key} could not be updated after {attempt} attempts: another writer \
|
||||
kept winning the race. Nothing was lost -- the blob is stored -- but this \
|
||||
entry is unrecorded and the file will read as a miss."
|
||||
)));
|
||||
}
|
||||
debug!(key = %manifest_key, attempt, "manifest write raced; merging and retrying");
|
||||
|
||||
@@ -18,9 +18,9 @@ use object_store::path::Path;
|
||||
use object_store::signer::Signer;
|
||||
use object_store::{
|
||||
GetOptions, GetRange, ObjectStore, ObjectStoreExt, PutMode, PutOptions, PutPayload,
|
||||
WriteMultipart,
|
||||
UpdateVersion, WriteMultipart,
|
||||
};
|
||||
use rustingface_core::ports::{BlobRead, BlobWriter, Store, StoredObject};
|
||||
use rustingface_core::ports::{BlobRead, BlobWriter, Store, StoredObject, Version};
|
||||
use rustingface_entities::config::Storage as StorageConfig;
|
||||
use rustingface_entities::error::{Error, Result};
|
||||
use tracing::debug;
|
||||
@@ -29,6 +29,9 @@ use tracing::debug;
|
||||
pub struct ObjectStoreAdapter {
|
||||
inner: Arc<dyn ObjectStore>,
|
||||
signer: Option<Arc<AmazonS3>>,
|
||||
/// The local filesystem store implements no conditional update, so a
|
||||
/// caller has to fall back to verifying its write landed.
|
||||
compare_and_set: bool,
|
||||
}
|
||||
|
||||
impl ObjectStoreAdapter {
|
||||
@@ -53,6 +56,7 @@ impl ObjectStoreAdapter {
|
||||
return Ok(Self {
|
||||
inner: Arc::new(store),
|
||||
signer: None,
|
||||
compare_and_set: false,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -85,6 +89,7 @@ impl ObjectStoreAdapter {
|
||||
Ok(Self {
|
||||
inner: Arc::clone(&s3) as Arc<dyn ObjectStore>,
|
||||
signer: Some(s3),
|
||||
compare_and_set: true,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -97,6 +102,7 @@ impl ObjectStoreAdapter {
|
||||
Ok(Self {
|
||||
inner: Arc::new(store),
|
||||
signer: None,
|
||||
compare_and_set: false,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -124,6 +130,65 @@ impl Store for ObjectStoreAdapter {
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_versioned(&self, key: &str) -> Result<Option<(Bytes, Version)>> {
|
||||
match self.inner.get(&Path::from(key)).await {
|
||||
Ok(result) => {
|
||||
let version = Version {
|
||||
e_tag: result.meta.e_tag.clone(),
|
||||
version: result.meta.version.clone(),
|
||||
};
|
||||
let bytes = result
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| storage(&format!("reading {key}"), e))?;
|
||||
Ok(Some((bytes, version)))
|
||||
}
|
||||
Err(err) if absent(&err) => Ok(None),
|
||||
Err(err) => Err(storage(&format!("reading {key}"), err)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn put_if_version(
|
||||
&self,
|
||||
key: &str,
|
||||
body: Bytes,
|
||||
version: Option<&Version>,
|
||||
) -> Result<bool> {
|
||||
if !self.compare_and_set {
|
||||
self.put(key, body).await?;
|
||||
return Ok(true);
|
||||
}
|
||||
let mode = match version {
|
||||
None => PutMode::Create,
|
||||
Some(version) => PutMode::Update(UpdateVersion {
|
||||
e_tag: version.e_tag.clone(),
|
||||
version: version.version.clone(),
|
||||
}),
|
||||
};
|
||||
match self
|
||||
.inner
|
||||
.put_opts(
|
||||
&Path::from(key),
|
||||
PutPayload::from_bytes(body),
|
||||
PutOptions {
|
||||
mode,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => Ok(true),
|
||||
Err(object_store::Error::AlreadyExists { .. }) => Ok(false),
|
||||
Err(object_store::Error::Precondition { .. }) => Ok(false),
|
||||
Err(object_store::Error::NotModified { .. }) => Ok(false),
|
||||
Err(err) => Err(storage(&format!("conditionally writing {key}"), err)),
|
||||
}
|
||||
}
|
||||
|
||||
fn supports_compare_and_set(&self) -> bool {
|
||||
self.compare_and_set
|
||||
}
|
||||
|
||||
async fn put(&self, key: &str, body: Bytes) -> Result<()> {
|
||||
self.inner
|
||||
.put(&Path::from(key), PutPayload::from_bytes(body))
|
||||
|
||||
Reference in New Issue
Block a user