diff --git a/seaweed-volume/src/server/grpc_server.rs b/seaweed-volume/src/server/grpc_server.rs index 32bb3e794..7574f1ef7 100644 --- a/seaweed-volume/src/server/grpc_server.rs +++ b/seaweed-volume/src/server/grpc_server.rs @@ -17,7 +17,7 @@ use crate::pb::master_pb; use crate::pb::master_pb::seaweed_client::SeaweedClient; use crate::pb::volume_server_pb; use crate::pb::volume_server_pb::volume_server_server::VolumeServer; -use crate::storage::erasure_coding::ec_shard::DATA_SHARDS_COUNT; +use crate::storage::erasure_coding::ec_shard::{DATA_SHARDS_COUNT, ShardId, shard_id_try_from}; use crate::storage::needle::needle::{self, Needle}; use crate::storage::types::*; use crate::storage::volume::VolumeSpec; @@ -3322,6 +3322,14 @@ impl VolumeServer for VolumeGrpcService { let req = request.into_inner(); let vid = VolumeId(req.volume_id); + // Validate wire shard ids at the boundary: ShardId is u8 but only + // 0..MAX_SHARD_COUNT are valid. Rejects 256 (would truncate to 0) + // and 270 (would alias 14). + let mut shard_ids: Vec = Vec::with_capacity(req.shard_ids.len()); + for &sid in &req.shard_ids { + shard_ids.push(shard_id_try_from(sid).map_err(Status::invalid_argument)?); + } + // Select target location: // When disk_id > 0: use that specific location. // When disk_id == 0 (unset): auto-select via @@ -3355,7 +3363,7 @@ impl VolumeServer for VolumeGrpcService { // writing them all to one disk would duplicate the other // disks' claims. Refuse so the caller splits the batch per // shard (or chooses explicitly via disk_id). - let owners = store.ec_shard_owner_disks(vid, &req.shard_ids); + let owners = store.ec_shard_owner_disks(vid, &shard_ids); if owners.len() > 1 { let dirs: Vec<&str> = owners .iter() @@ -3363,14 +3371,14 @@ impl VolumeServer for VolumeGrpcService { .collect(); return Err(Status::failed_precondition(format!( "volume {} shards {:?} are already owned by multiple local disks {:?}: no single destination; copy per shard or pass disk_id", - req.volume_id, req.shard_ids, dirs + req.volume_id, shard_ids, dirs ))); } match store.find_ec_shard_target_location( &req.collection, vid, DATA_SHARDS_COUNT as u32, - &req.shard_ids, + &shard_ids, ) { Some(i) => { let loc = &store.locations[i]; @@ -3417,7 +3425,7 @@ impl VolumeServer for VolumeGrpcService { .max_encoding_message_size(GRPC_MAX_MESSAGE_SIZE); // Copy each shard - for &shard_id in &req.shard_ids { + for &shard_id in &shard_ids { let ext = format!(".ec{:02}", shard_id); let copy_req = volume_server_pb::CopyFileRequest { volume_id: req.volume_id, @@ -3674,7 +3682,11 @@ impl VolumeServer for VolumeGrpcService { } let mut store = self.state.store.write().unwrap(); - store.delete_ec_shards(vid, &req.collection, &req.shard_ids); + let mut shard_ids: Vec = Vec::with_capacity(req.shard_ids.len()); + for &sid in &req.shard_ids { + shard_ids.push(shard_id_try_from(sid).map_err(Status::invalid_argument)?); + } + store.delete_ec_shards(vid, &req.collection, &shard_ids); drop(store); self.state.volume_state_notify.notify_one(); Ok(Response::new( @@ -3691,6 +3703,17 @@ impl VolumeServer for VolumeGrpcService { let req = request.into_inner(); let vid = VolumeId(req.volume_id); + // Pre-validate the ENTIRE batch before any mutation: validating inside + // the mount loop would mount a prefix (e.g. shard 0 of [0, 32]) and + // then fail, leaving a partial mutation while skipping the sidecar + // reload + notify below. Reject up front so invalid batches change no + // state. + // Validate wire shard ids at the boundary (rejects truncation aliases like 256→0). + let mut validated: Vec = Vec::with_capacity(req.shard_ids.len()); + for &sid in &req.shard_ids { + validated.push(shard_id_try_from(sid).map_err(Status::invalid_argument)?); + } + // Fetch a missing .ecx from a peer first so on-disk shards that never had // a local index can be mounted (issue #10104). Driven on demand by // ec.rebuild. volume_id 0 recovers every orphan on this server, including @@ -3702,7 +3725,7 @@ impl VolumeServer for VolumeGrpcService { // Mount one shard at a time, returning error on first failure. // Matches Go: for _, shardId := range req.ShardIds { err = vs.store.MountEcShards(...) } let mut store = self.state.store.write().unwrap(); - for &shard_id in &req.shard_ids { + for &shard_id in &validated { store .mount_ec_shard(vid, &req.collection, shard_id, &req.source_disk_type) .map_err(|e| { @@ -3746,10 +3769,20 @@ impl VolumeServer for VolumeGrpcService { let req = request.into_inner(); let vid = VolumeId(req.volume_id); + // Pre-validate the ENTIRE batch before acquiring the write lock or + // mutating: validating inside the unmount loop would unmount a prefix + // (e.g. shard 0 of [0, 32]) and then fail, leaving a partial mutation. + // Reject up front so invalid batches change no state. Auth stays first. + // Validate wire shard ids at the boundary (rejects truncation aliases like 256→0). + let mut validated: Vec = Vec::with_capacity(req.shard_ids.len()); + for &sid in &req.shard_ids { + validated.push(shard_id_try_from(sid).map_err(Status::invalid_argument)?); + } + // Unmount one shard at a time, returning error on first failure. // Matches Go: for _, shardId := range req.ShardIds { err = vs.store.UnmountEcShards(...) } let mut store = self.state.store.write().unwrap(); - for &shard_id in &req.shard_ids { + for &shard_id in &validated { store .unmount_ec_shard(vid, shard_id, req.encode_ts_ns) .map_err(|e| { @@ -3770,6 +3803,7 @@ impl VolumeServer for VolumeGrpcService { ) -> Result, Status> { let req = request.into_inner(); let vid = VolumeId(req.volume_id); + let shard_id = shard_id_try_from(req.shard_id).map_err(Status::invalid_argument)?; let store = self.state.store.read().unwrap(); // Reconciled EC volumes can have their shards split across @@ -3778,7 +3812,7 @@ impl VolumeServer for VolumeGrpcService { // rather than first-match `find_ec_volume(vid)` which would // miss shards that live on a sibling. Mirrors Go's findEcShard. let ec_vol = store - .find_ec_volume_with_shard(vid, req.shard_id) + .find_ec_volume_with_shard(vid, shard_id) .ok_or_else(|| { Status::not_found(format!( "ec volume {} shard {} not found", @@ -3822,7 +3856,7 @@ impl VolumeServer for VolumeGrpcService { // find_ec_volume_with_shard already verified it. let shard = ec_vol .shards - .get(req.shard_id as usize) + .get(shard_id as usize) .and_then(|s| s.as_ref()) .ok_or_else(|| { Status::not_found(format!( @@ -7712,6 +7746,91 @@ mod tests { assert_eq!(resp.total_files, 1); } + /// Batch atomicity: mount pre-validates the ENTIRE shard_ids before + /// acquiring the write lock or mounting anything. A batch like [0, 32] + /// must fail with InvalidArgument and mount NOTHING — not the valid + /// prefix (shard 0). Regression test for the in-loop validation that + /// mounted 0 then returned InvalidArgument, skipping the sidecar reload + /// + notify. + #[tokio::test] + async fn test_mount_rejects_invalid_batch_without_partial_mutation() { + let (service, _tmp) = make_local_service_with_volume("", None); + service + .volume_ec_shards_generate(Request::new( + volume_server_pb::VolumeEcShardsGenerateRequest { + volume_id: 1, + collection: String::new(), + }, + )) + .await + .unwrap(); + + let err = service + .volume_ec_shards_mount(Request::new(volume_server_pb::VolumeEcShardsMountRequest { + volume_id: 1, + collection: String::new(), + shard_ids: vec![0, 32], + source_disk_type: String::new(), + recover_missing_index: false, + })) + .await + .expect_err("batch with shard 32 must be rejected"); + assert_eq!(err.code(), tonic::Code::InvalidArgument, "{}", err); + + // No partial mutation: shard 0 must NOT be mounted. + let store = service.state.store.read().unwrap(); + assert!( + store.find_ec_volume_with_shard(VolumeId(1), 0).is_none(), + "invalid batch must mount nothing, but shard 0 is mounted" + ); + } + + /// Batch atomicity (unmount side): unmount pre-validates the ENTIRE + /// shard_ids after admin auth but before the write lock / mutation. A + /// batch like [0, 32] must fail with InvalidArgument and unmount NOTHING. + #[tokio::test] + async fn test_unmount_rejects_invalid_batch_without_partial_mutation() { + let (service, _tmp) = make_local_service_with_volume("", None); + service + .volume_ec_shards_generate(Request::new( + volume_server_pb::VolumeEcShardsGenerateRequest { + volume_id: 1, + collection: String::new(), + }, + )) + .await + .unwrap(); + service + .volume_ec_shards_mount(Request::new(volume_server_pb::VolumeEcShardsMountRequest { + volume_id: 1, + collection: String::new(), + shard_ids: (0..14).collect(), + source_disk_type: String::new(), + recover_missing_index: false, + })) + .await + .unwrap(); + + let err = service + .volume_ec_shards_unmount(Request::new( + volume_server_pb::VolumeEcShardsUnmountRequest { + volume_id: 1, + shard_ids: vec![0, 32], + encode_ts_ns: 0, + }, + )) + .await + .expect_err("batch with shard 32 must be rejected"); + assert_eq!(err.code(), tonic::Code::InvalidArgument, "{}", err); + + // No partial mutation: shard 0 must STILL be mounted. + let store = service.state.store.read().unwrap(); + assert!( + store.find_ec_volume_with_shard(VolumeId(1), 0).is_some(), + "invalid unmount batch must unmount nothing, but shard 0 is gone" + ); + } + /// Two locations, one vid: the split-disk / cross-disk-reconcile layout /// `build_split_disk_store` exercises in store_ec_reconcile.rs. Neither /// disk holds a `.dat`, so each disk's orphan shard survives diff --git a/seaweed-volume/src/server/store_ec.rs b/seaweed-volume/src/server/store_ec.rs index 20bc6557f..845849c2b 100644 --- a/seaweed-volume/src/server/store_ec.rs +++ b/seaweed-volume/src/server/store_ec.rs @@ -45,7 +45,7 @@ use crate::pb::volume_server_pb::{ use crate::server::grpc_client::{GRPC_MAX_MESSAGE_SIZE, build_grpc_endpoint, parse_grpc_address}; use crate::server::request_id::outgoing_request_id_interceptor; use crate::server::volume_server::{VolumeServerState, to_http_address}; -use crate::storage::erasure_coding::ec_shard::ShardId; +use crate::storage::erasure_coding::ec_shard::{ShardId, shard_id_try_from}; use crate::storage::needle::needle::{Needle, NeedleError, get_actual_size}; use crate::storage::store_ec_reconcile::EcVolumeMissingIndex; use crate::storage::types::*; @@ -898,7 +898,12 @@ async fn cached_lookup_ec_shard_locations( .iter() .map(format_location_as_server_address) .collect(); - out.insert(entry.shard_id as ShardId, addrs); + // Defensive: skip out-of-range shard ids from the master instead of + // truncating (256 would alias 0). Valid replies are unaffected. + let Ok(sid) = shard_id_try_from(entry.shard_id) else { + continue; + }; + out.insert(sid, addrs); } Ok(out) } @@ -1207,7 +1212,10 @@ async fn recover_one_remote_ec_shard_interval( // shard from a different encode run must not be fed to Reed-Solomon; // lenient only when the caller carries no identity (pre-upgrade). // Mirrors Go's `readLocalEcShardInterval`. - let owner = match store.find_ec_volume_with_shard(vid, sid as u32) { + let Ok(sid_shard) = ShardId::try_from(sid) else { + continue; + }; + let owner = match store.find_ec_volume_with_shard(vid, sid_shard) { Some(ecv) if expected_encode_ts_ns == 0 || ecv.encode_ts_ns == expected_encode_ts_ns => { diff --git a/seaweed-volume/src/storage/disk_location.rs b/seaweed-volume/src/storage/disk_location.rs index 142453ecd..9bbbd1859 100644 --- a/seaweed-volume/src/storage/disk_location.rs +++ b/seaweed-volume/src/storage/disk_location.rs @@ -16,7 +16,7 @@ use crate::config::MinFreeSpace; use crate::storage::erasure_coding::ec_bitrot::remove_bitrot_sidecars; use crate::storage::erasure_coding::ec_shard::{ DATA_SHARDS_COUNT, ERASURE_CODING_LARGE_BLOCK_SIZE, ERASURE_CODING_SMALL_BLOCK_SIZE, - EcVolumeShard, + EcVolumeShard, ShardId, }; use crate::storage::erasure_coding::ec_volume::EcVolume; use crate::storage::needle_map::NeedleMapKind; @@ -817,7 +817,7 @@ impl DiskLocation { &mut self, vid: VolumeId, collection: &str, - shard_ids: &[u32], + shard_ids: &[ShardId], source_disk_type: &str, ) -> Result<(), VolumeError> { let idx_dir = self.idx_directory.clone(); @@ -839,7 +839,7 @@ impl DiskLocation { &mut self, vid: VolumeId, collection: &str, - shard_ids: &[u32], + shard_ids: &[ShardId], idx_dir: &str, source_disk_type: &str, ) -> Result<(), VolumeError> { @@ -873,10 +873,10 @@ impl DiskLocation { // keep the existing registration (mirrors Go's AddEcVolumeShard // added=false) — re-adding would replace a serving fd and bump // the ec_shards gauge without growing the mounted count. - if ec_vol.has_shard(shard_id as u8) { + if ec_vol.has_shard(shard_id) { continue; } - let mut shard = EcVolumeShard::new(&dir, collection, vid, shard_id as u8); + let mut shard = EcVolumeShard::new(&dir, collection, vid, shard_id); shard.disk_type = ec_vol.disk_type.clone(); if let Err(e) = ec_vol.add_shard(shard) { // The shard was dropped (its descriptors closed) inside the @@ -904,14 +904,14 @@ impl DiskLocation { /// caller passes a shard that lives on a sibling disk /// (cross-disk reconcile makes that the common case for the same /// `vid` after reconciliation). - pub fn unmount_ec_shards(&mut self, vid: VolumeId, shard_ids: &[u32]) { + pub fn unmount_ec_shards(&mut self, vid: VolumeId, shard_ids: &[ShardId]) { if let Some(ec_vol) = self.ec_volumes.get_mut(&vid) { let collection = ec_vol.collection.clone(); for &shard_id in shard_ids { - if !ec_vol.has_shard(shard_id as u8) { + if !ec_vol.has_shard(shard_id) { continue; } - ec_vol.remove_shard(shard_id as u8); + let _ = ec_vol.remove_shard(shard_id); crate::metrics::VOLUME_GAUGE .with_label_values(&[&collection, "ec_shards"]) .dec(); @@ -971,7 +971,7 @@ impl DiskLocation { } entries.sort(); - let mut same_volume_shards: Vec<(String, u32)> = Vec::new(); // (filename, shard_id) + let mut same_volume_shards: Vec<(String, ShardId)> = Vec::new(); // (filename, shard_id) let mut prev_vid: Option = None; let mut prev_collection: String = String::new(); @@ -1036,7 +1036,12 @@ impl DiskLocation { /// Validate + mount a (collection, vid) group when its `.ecx` is /// found. Mirrors `handleFoundEcxFile` in /// `weed/storage/disk_location_ec.go`. - fn handle_found_ecx_file(&mut self, shards: &[(String, u32)], collection: &str, vid: VolumeId) { + fn handle_found_ecx_file( + &mut self, + shards: &[(String, ShardId)], + collection: &str, + vid: VolumeId, + ) { let base = volume_file_name(&self.directory, collection, vid); let dat_path = format!("{}.dat", base); let dat_exists = check_dat_file_exists(&dat_path); @@ -1050,7 +1055,7 @@ impl DiskLocation { return; } - let shard_ids: Vec = shards.iter().map(|(_, sid)| *sid).collect(); + let shard_ids: Vec = shards.iter().map(|(_, sid)| *sid).collect(); if let Err(e) = self.mount_ec_shards(vid, collection, &shard_ids, "") { // A mount failure (corrupt/locked .ecx, EMFILE, transient I/O) is // not proof the shards are disposable -- validate_ec_volume already @@ -1072,7 +1077,7 @@ impl DiskLocation { /// distributed-EC shards waiting for cross-disk reconciliation. fn check_orphaned_shards( &self, - shards: &[(String, u32)], + shards: &[(String, ShardId)], collection: &str, vid: VolumeId, ) -> bool { @@ -1263,7 +1268,7 @@ fn parse_collection_volume_id(base: &str) -> Option<(String, VolumeId)> { /// `pub(crate)` re-export of [`parse_ec_shard_extension`] for the /// cross-disk reconcile in `store_ec_reconcile.rs`. -pub(crate) fn is_ec_shard_extension(ext: &str) -> Option { +pub(crate) fn is_ec_shard_extension(ext: &str) -> Option { parse_ec_shard_extension(ext) } @@ -1277,7 +1282,7 @@ pub(crate) fn is_ec_shard_extension(ext: &str) -> Option { /// shardId > 255` guard. The 3-digit form (`.ec100`–`.ec255`) is /// retained so the parser can still recognise shards from custom /// 32+ ratios that fit in a u8 even though OSS only ships 10+4. -fn parse_ec_shard_extension(ext: &str) -> Option { +fn parse_ec_shard_extension(ext: &str) -> Option { let rest = ext.strip_prefix(".ec")?; if rest.len() < 2 || rest.len() > 3 { return None; @@ -1286,7 +1291,7 @@ fn parse_ec_shard_extension(ext: &str) -> Option { if id > 255 { return None; } - Some(id) + ShardId::try_from(id).ok() } /// Robust check that a `.dat` with actual data exists. An empty `.dat` diff --git a/seaweed-volume/src/storage/erasure_coding/ec_shard.rs b/seaweed-volume/src/storage/erasure_coding/ec_shard.rs index 030fd0795..538fba8e8 100644 --- a/seaweed-volume/src/storage/erasure_coding/ec_shard.rs +++ b/seaweed-volume/src/storage/erasure_coding/ec_shard.rs @@ -16,6 +16,20 @@ pub const ERASURE_CODING_SMALL_BLOCK_SIZE: usize = 1024 * 1024; // 1MB pub type ShardId = u8; +/// Validate a wire shard id. `ShardId` is `u8` but only 0..MAX_SHARD_COUNT are valid. +/// Rejects 256 (would truncate to 0 and delete .ec00) and 270 (would alias 14). +pub fn shard_id_try_from(v: u32) -> Result { + if v < MAX_SHARD_COUNT as u32 { + Ok(v as ShardId) + } else { + Err(format!( + "invalid shard id {} (max {})", + v, + MAX_SHARD_COUNT - 1 + )) + } +} + /// A single erasure-coded shard file. pub struct EcVolumeShard { pub volume_id: VolumeId, @@ -251,4 +265,42 @@ mod tests { let shard = EcVolumeShard::new("/data", "", VolumeId(7), 13); assert_eq!(shard.file_name(), "/data/7.ec13"); } + + #[test] + fn test_shard_id_try_from_u32_rejects_overflow() { + use super::{MAX_SHARD_COUNT, shard_id_try_from}; + assert_eq!(shard_id_try_from(0).unwrap(), 0u8); + assert_eq!(shard_id_try_from(14).unwrap(), 14u8); + assert_eq!(shard_id_try_from(31).unwrap(), 31u8); + assert!(shard_id_try_from(32).is_err()); + assert!(shard_id_try_from(256).is_err()); + assert!(shard_id_try_from(270).is_err()); + assert!(shard_id_try_from(u32::MAX).is_err()); + assert_eq!(MAX_SHARD_COUNT, 32); + } + + #[test] + fn test_shard_batch_validation_is_atomic_rejects_without_partial_prefix() { + use super::shard_id_try_from; + // The mount/unmount handlers pre-validate the ENTIRE req.shard_ids into + // a Vec BEFORE acquiring the write lock or mutating any EC + // state. This test pins the validation half of that contract at the + // unit level: a batch like [0, 32] must fail as a whole, so by + // construction no validated prefix (e.g. shard 0) is ever applied. + // The handler-level tests below assert the no-state-change half. + let batch = vec![0u32, 32u32]; + let validated: Result, _> = + batch.iter().map(|&sid| shard_id_try_from(sid)).collect(); + assert!( + validated.is_err(), + "batch {:?} must be rejected as a whole", + batch + ); + // A fully-valid batch still validates cleanly. + let ok: Result, _> = [0u32, 1u32, 13u32] + .iter() + .map(|&sid| shard_id_try_from(sid)) + .collect(); + assert_eq!(ok.unwrap(), vec![0u8, 1u8, 13u8]); + } } diff --git a/seaweed-volume/src/storage/erasure_coding/ec_volume.rs b/seaweed-volume/src/storage/erasure_coding/ec_volume.rs index e94837bb0..26cd19ea6 100644 --- a/seaweed-volume/src/storage/erasure_coding/ec_volume.rs +++ b/seaweed-volume/src/storage/erasure_coding/ec_volume.rs @@ -841,11 +841,23 @@ impl EcVolume { } /// Remove and close a shard. - pub fn remove_shard(&mut self, shard_id: ShardId) { - if let Some(ref mut shard) = self.shards[shard_id as usize] { + pub fn remove_shard(&mut self, shard_id: ShardId) -> io::Result<()> { + let idx = shard_id as usize; + if idx >= self.shards.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "invalid shard id {} (max {})", + shard_id, + self.shards.len().saturating_sub(1) + ), + )); + } + if let Some(ref mut shard) = self.shards[idx] { shard.close(); } - self.shards[shard_id as usize] = None; + self.shards[idx] = None; + Ok(()) } /// Get a ShardBits bitmap of locally available shards. @@ -867,7 +879,7 @@ impl EcVolume { /// Reports whether `shard_id` is currently registered to this /// EcVolume (used by the cross-disk reconcile to skip already- /// loaded shards). - pub fn has_shard(&self, shard_id: u8) -> bool { + pub fn has_shard(&self, shard_id: ShardId) -> bool { self.shards .get(shard_id as usize) .map(|s| s.is_some()) diff --git a/seaweed-volume/src/storage/store.rs b/seaweed-volume/src/storage/store.rs index f3d6a1b63..c415ac214 100644 --- a/seaweed-volume/src/storage/store.rs +++ b/seaweed-volume/src/storage/store.rs @@ -12,7 +12,7 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use crate::config::MinFreeSpace; use crate::pb::master_pb; use crate::storage::disk_location::DiskLocation; -use crate::storage::erasure_coding::ec_shard::{EcVolumeShard, MAX_SHARD_COUNT}; +use crate::storage::erasure_coding::ec_shard::{EcVolumeShard, MAX_SHARD_COUNT, ShardId}; use crate::storage::erasure_coding::ec_volume::EcVolume; use crate::storage::needle::needle::Needle; use crate::storage::needle_map::NeedleMapKind; @@ -300,7 +300,7 @@ impl Store { collection: &str, vid: VolumeId, data_shard_count: u32, - shard_ids: &[u32], + shard_ids: &[ShardId], ) -> Option { const TIER_ANY_DISK: u8 = 1; const TIER_HDD: u8 = 2; @@ -359,7 +359,7 @@ impl Store { /// owner (`volume_ec_shards_copy` refuses such a batch instead of /// guessing). Mirrors `Store.EcShardOwnerDisks` in /// `weed/storage/store_ec.go`. - pub fn ec_shard_owner_disks(&self, vid: VolumeId, shard_ids: &[u32]) -> Vec { + pub fn ec_shard_owner_disks(&self, vid: VolumeId, shard_ids: &[ShardId]) -> Vec { self.locations .iter() .enumerate() @@ -882,7 +882,7 @@ impl Store { &mut self, vid: VolumeId, collection: &str, - shard_ids: &[u32], + shard_ids: &[ShardId], ) -> Result<(), VolumeError> { // Find the location where the EC files live let loc_idx = self.find_ec_location(vid, collection).ok_or_else(|| { @@ -903,12 +903,12 @@ impl Store { &mut self, vid: VolumeId, collection: &str, - shard_id: u32, + shard_id: ShardId, source_disk_type: &str, ) -> Result<(), VolumeError> { for loc in &mut self.locations { // Check if the shard file exists on this location - let shard = EcVolumeShard::new(&loc.directory, collection, vid, shard_id as u8); + let shard = EcVolumeShard::new(&loc.directory, collection, vid, shard_id); if std::path::Path::new(&shard.file_name()).exists() { loc.mount_ec_shards(vid, collection, &[shard_id], source_disk_type)?; return Ok(()); @@ -929,7 +929,7 @@ impl Store { /// the same store (#9252). DiskLocation::unmount_ec_shards /// already skips shards that aren't mounted, so this is safe to /// fan out blindly. - pub fn unmount_ec_shards(&mut self, vid: VolumeId, shard_ids: &[u32]) { + pub fn unmount_ec_shards(&mut self, vid: VolumeId, shard_ids: &[ShardId]) { for loc in &mut self.locations { if loc.has_ec_volume(vid) { loc.unmount_ec_shards(vid, shard_ids); @@ -942,7 +942,7 @@ impl Store { pub fn unmount_ec_shard( &mut self, vid: VolumeId, - shard_id: u32, + shard_id: ShardId, req_encode_ts_ns: i64, ) -> Result<(), VolumeError> { // Walk all locations rather than stopping at the first with the @@ -950,7 +950,7 @@ impl Store { // multiple disks, with the target shard on any of them. for disk_id in 0..self.locations.len() { let ec_vol = self.locations[disk_id].find_ec_volume(vid); - let has_shard = ec_vol.is_some_and(|ec_vol| ec_vol.has_shard(shard_id as u8)); + let has_shard = ec_vol.is_some_and(|ec_vol| ec_vol.has_shard(shard_id)); if !has_shard { continue; } @@ -1069,10 +1069,10 @@ impl Store { /// disks (each holding a disjoint subset of the shards). Without /// this, callers using `find_ec_volume(vid)` would only see the /// first disk and miss shards that live on a sibling. - pub fn find_ec_shard_location(&self, vid: VolumeId, shard_id: u32) -> Option { + pub fn find_ec_shard_location(&self, vid: VolumeId, shard_id: ShardId) -> Option { for (i, loc) in self.locations.iter().enumerate() { if let Some(ecv) = loc.find_ec_volume(vid) - && ecv.has_shard(shard_id as u8) + && ecv.has_shard(shard_id) { return Some(i); } @@ -1083,10 +1083,10 @@ impl Store { /// Like [`Self::find_ec_shard_location`] but returns the EcVolume /// reference directly. Borrows the store immutably for the /// EcVolume's lifetime. - pub fn find_ec_volume_with_shard(&self, vid: VolumeId, shard_id: u32) -> Option<&EcVolume> { + pub fn find_ec_volume_with_shard(&self, vid: VolumeId, shard_id: ShardId) -> Option<&EcVolume> { for loc in &self.locations { if let Some(ecv) = loc.find_ec_volume(vid) - && ecv.has_shard(shard_id as u8) + && ecv.has_shard(shard_id) { return Some(ecv); } @@ -1118,7 +1118,10 @@ impl Store { found_vol = Some(ecv); } for (shard_id, dir) in dirs.iter_mut().enumerate() { - if dir.is_none() && ecv.has_shard(shard_id as u8) { + let Ok(sid) = ShardId::try_from(shard_id) else { + continue; + }; + if dir.is_none() && ecv.has_shard(sid) { *dir = Some(loc.directory.clone()); } } @@ -1214,12 +1217,12 @@ impl Store { } /// Delete EC shard files from disk. - pub fn delete_ec_shards(&mut self, vid: VolumeId, collection: &str, shard_ids: &[u32]) { + pub fn delete_ec_shards(&mut self, vid: VolumeId, collection: &str, shard_ids: &[ShardId]) { // Delete shard files from disk, tracking which locations actually held one. let mut deleted_at = vec![false; self.locations.len()]; for (i, loc) in self.locations.iter().enumerate() { for &shard_id in shard_ids { - let shard = EcVolumeShard::new(&loc.directory, collection, vid, shard_id as u8); + let shard = EcVolumeShard::new(&loc.directory, collection, vid, shard_id); if std::fs::remove_file(shard.file_name()).is_ok() { deleted_at[i] = true; } @@ -1555,13 +1558,13 @@ fn ec_free_shard_count(loc: &DiskLocation, data_shard_count: u32) -> i64 { /// the in-memory registration the read path and heartbeats use. /// /// Mirrors `ownedEcShardCount` in `weed/storage/store_ec.go`. -fn owned_ec_shard_count(loc: &DiskLocation, vid: VolumeId, shard_ids: &[u32]) -> usize { +fn owned_ec_shard_count(loc: &DiskLocation, vid: VolumeId, shard_ids: &[ShardId]) -> usize { let Some(ecv) = loc.find_ec_volume(vid) else { return 0; }; shard_ids .iter() - .filter(|&&shard_id| ecv.has_shard(shard_id as u8)) + .filter(|&&shard_id| ecv.has_shard(shard_id)) .count() } @@ -2661,7 +2664,7 @@ mod tests { // Fill disk 0 past its shard-slot budget so ec_free_shard_count is 0. let filler = VolumeId(10000); let filler_base = volume_file_name(&store.locations[0].directory, collection, filler); - let filler_shards: Vec = (0..10).collect(); + let filler_shards: Vec = (0..10).collect(); for shard_id in &filler_shards { std::fs::write(format!("{}.ec{:02}", filler_base, shard_id), b"x").unwrap(); } diff --git a/seaweed-volume/src/storage/store_ec_reconcile.rs b/seaweed-volume/src/storage/store_ec_reconcile.rs index 7057761ab..f5253832d 100644 --- a/seaweed-volume/src/storage/store_ec_reconcile.rs +++ b/seaweed-volume/src/storage/store_ec_reconcile.rs @@ -21,7 +21,7 @@ use std::fs; use tracing::{error, info, warn}; use crate::storage::disk_location::{is_ec_shard_extension, parse_collection_volume_id_pub}; -use crate::storage::erasure_coding::ec_shard::DATA_SHARDS_COUNT; +use crate::storage::erasure_coding::ec_shard::{DATA_SHARDS_COUNT, ShardId}; use crate::storage::store::Store; use crate::storage::types::VolumeId; @@ -83,7 +83,7 @@ struct EcxOwnerInfo { /// One unit of reconcile work: the disk holding orphan shards, the volume /// they belong to, the shard files, the `.ecx` owner, and whether the /// mirror already installed sidecars locally (`use_local_idx`). -type OrphanShardLoad = (usize, EcKey, Vec<(String, u32)>, EcxOwnerInfo, bool); +type OrphanShardLoad = (usize, EcKey, Vec<(String, ShardId)>, EcxOwnerInfo, bool); impl Store { /// Run cross-disk orphan-shard reconciliation. Should be called @@ -135,7 +135,7 @@ impl Store { for (loc_idx, key, shards, owner, use_local_idx) in to_load { let shard_names: Vec<&str> = shards.iter().map(|(n, _)| n.as_str()).collect(); let loc_dir = self.locations[loc_idx].directory.clone(); - let shard_ids: Vec = shards.iter().map(|(_, sid)| *sid).collect(); + let shard_ids: Vec = shards.iter().map(|(_, sid)| *sid).collect(); if use_local_idx { info!( @@ -477,13 +477,13 @@ impl Store { /// Unlike `reconcile_ec_shards_across_disks` it needs no sibling disk, so a /// single-disk store recovers once its index has been fetched from a peer. fn load_orphan_ec_shards_with_local_index(&mut self) { - let mut work: Vec<(usize, EcKey, Vec)> = Vec::new(); + let mut work: Vec<(usize, EcKey, Vec)> = Vec::new(); for (loc_idx, loc) in self.locations.iter().enumerate() { for (key, shards) in collect_orphan_ec_shards(loc, loc_idx) { if !loc.has_ecx_file_on_disk(&key.collection, key.vid) { continue; } - let ids: Vec = shards.iter().map(|(_, sid)| *sid).collect(); + let ids: Vec = shards.iter().map(|(_, sid)| *sid).collect(); work.push((loc_idx, key, ids)); } } @@ -510,8 +510,8 @@ impl Store { fn collect_orphan_ec_shards( loc: &crate::storage::disk_location::DiskLocation, _loc_idx: usize, -) -> HashMap> { - let mut orphans: HashMap> = HashMap::new(); +) -> HashMap> { + let mut orphans: HashMap> = HashMap::new(); let Ok(read) = fs::read_dir(&loc.directory) else { return orphans; }; @@ -539,7 +539,7 @@ fn collect_orphan_ec_shards( }; // Skip shards that are already registered to an EcVolume. if let Some(ecv) = loc.find_ec_volume(vid) - && ecv.has_shard(shard_id as u8) + && ecv.has_shard(shard_id) { continue; } diff --git a/seaweed-volume/tests/http_integration.rs b/seaweed-volume/tests/http_integration.rs index b0cbd91d4..a622d045a 100644 --- a/seaweed-volume/tests/http_integration.rs +++ b/seaweed-volume/tests/http_integration.rs @@ -957,6 +957,7 @@ async fn replicate_write_does_not_re_replicate() { #[tokio::test] async fn chunk_manifest_expands_chunk_stored_on_ec_volume() { use seaweed_volume::storage::erasure_coding::ec_encoder::write_ec_files; + use seaweed_volume::storage::erasure_coding::ec_shard::ShardId; use seaweed_volume::storage::needle::needle::{FileId, Needle}; use seaweed_volume::storage::types::{Cookie, NeedleId}; use seaweed_volume::storage::volume::{Volume, VolumeSpec}; @@ -998,7 +999,7 @@ async fn chunk_manifest_expands_chunk_stored_on_ec_volume() { // after ec.encode retired the regular volume. { let mut store = state.store.write().unwrap(); - let shard_ids: Vec = (0..14).collect(); + let shard_ids: Vec = (0..14).collect(); store.mount_ec_shards(VolumeId(2), "", &shard_ids).unwrap(); }