snapshot of https://github.com/seaweedfs/seaweedfs.git (master, 2026-09-16; blobs > 131072B omitted)
Some checks failed
go: build dev binaries / cleanup (push) Successful in 4s
Codespell / Check for spelling errors (push) Failing after 3s
docker: build dev containers / build-rust-binaries (amd64, x86_64-unknown-linux-musl) (push) Failing after 3s
docker: build dev containers / build-dev-containers (push) Has been skipped
docker: build dev containers / build-rust-binaries (arm64, true, aarch64-unknown-linux-musl) (push) Failing after 3s
End to End / FUSE Mount (push) Failing after 2s
EC Integration Tests / EC Integration Tests (push) Failing after 30s
FUSE DLM Integration Tests / FUSE DLM Integration Tests (push) Failing after 3s
FUSE Volume Server Failover Tests / FUSE Volume Server Failover (push) Failing after 3s
FUSE Integration Tests / FUSE Integration Testing (push) Failing after 2s
FUSE P2P Peer Chunk Sharing Integration Tests / FUSE P2P Peer Chunk Sharing (push) Failing after 3s
go: build binary / Go Vet (push) Failing after 2s
go: build binary / Go Vet 32-bit (push) Failing after 3s
go: build binary / Build (push) Failing after 4s
go: build binary / Build cross-platform (freebsd/amd64) (push) Failing after 3s
go: build binary / Build cross-platform (windows/amd64) (push) Failing after 3s
go: build binary / Build cross-platform (darwin/arm64) (push) Failing after 3s
go: build binary / Build cross-platform (windows/arm64) (push) Failing after 3s
go: build binary / Test (push) Failing after 2s
go: build binary / Test 32-bit (push) Failing after 3s
helm: lint and test charts / lint-test (push) Failing after 3s
Java Client Integration Tests / Java Integration Tests (11) (push) Failing after 3s
Java Client Integration Tests / Java Integration Tests (17) (push) Failing after 3s
go: build dev binaries / build_dev_linux_windows (amd64, linux) (push) Failing after 2s
go: build dev binaries / build_dev_linux_windows (amd64, windows) (push) Failing after 3s
go: build dev binaries / build_dev_darwin (amd64, darwin) (push) Has been skipped
go: build dev binaries / build_dev_darwin (arm64, darwin) (push) Has been skipped
Java Client Unit Tests / Java Unit Tests (11) (push) Failing after 3s
Java Client Unit Tests / Java Unit Tests (17) (push) Failing after 2s
Java Client Unit Tests / Java Unit Tests (21) (push) Failing after 3s
Java Client Unit Tests / Java Unit Tests (8) (push) Failing after 3s
Kafka Quick Test (Load Test with Schema Registry) / Kafka Client Load Test (Quick) (push) Failing after 3s
Kafka Gateway Tests / Kafka Unit Tests (unit-tests-1) (push) Failing after 11s
Kafka Gateway Tests / Kafka Integration Tests (Critical) (integration-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka End-to-End Tests (with SMQ) (e2e-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka Consumer Group Tests (Highly Isolated) (consumer-group-1) (push) Failing after 2s
Kafka Gateway Tests / Kafka Client Compatibility (with SMQ) (client-compat-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka SMQ Integration Tests (Full Stack) (smq-integration-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka Protocol Tests (Isolated) (protocol-1) (push) Failing after 3s
KMS Tests / KMS Provider Integration Tests (push) Failing after 2s
KMS Tests / S3 KMS End-to-End Tests (push) Failing after 3s
Master Cold Start Tests / Master Cold Start Tests (push) Failing after 4s
Metadata Subscribe Integration Tests / Metadata Subscribe Integration Tests (push) Failing after 3s
Multi-Master Tests / Multi-Master Failover Tests (push) Failing after 3s
Performance / CPU and Heap Profile (push) Failing after 2s
Performance / S3 Read/Write Benchmark (rust volume) (push) Failing after 3s
pjdfstest POSIX Compliance / pjdfstest (push) Failing after 2s
Plugin Worker Integration Tests / Plugin Worker: erasure_coding (push) Failing after 3s
Plugin Worker Integration Tests / Plugin Worker: vacuum (push) Failing after 3s
Rust Volume Server Tests / Rust Integration Tests (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 1) (push) Failing after 2s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 1) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 2) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 2) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 3) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Distributed Tests (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Performance Tests (push) Failing after 3s
Performance / Throughput Benchmark (go volume) (push) Failing after 3s
Performance / Throughput Benchmark (rust volume) (push) Failing after 2s
Performance / S3 Read/Write Benchmark (go volume) (push) Failing after 3s
Plugin Worker Integration Tests / Plugin Worker: volume_balance (push) Failing after 3s
PostgreSQL Gateway Tests / PostgreSQL Basic Tests (push) Failing after 3s
Rust Volume Server Tests / Rust Unit Tests (push) Failing after 2s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 3) (push) Failing after 3s
Rust Plugin Worker Tests / Rust Plugin Worker Build and Unit Tests (push) Failing after 3s
rust: build dev volume server binaries / cleanup (push) Successful in 3s
S3 ETag and ACL Tests / S3 ETag + ACL Integration Tests (push) Failing after 3s
S3 IAM Integration Tests / IAM Unit Tests (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (advanced) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (basic) (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Integration Tests (group) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (policy-enforcement) (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Integration Tests (sts) (push) Failing after 4s
S3 Keycloak Integration Tests / S3 Keycloak Integration Tests (push) Failing after 3s
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.11) (push) Failing after 3s
S3 PyArrow Parquet Tests / Go Unit Tests (Implicit Directory) (push) Has been cancelled
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.12) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-darwin (darwin-amd64, x86_64-apple-darwin) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-darwin (darwin-arm64, aarch64-apple-darwin) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-linux (linux-amd64, x86_64-unknown-linux-gnu) (push) Has been cancelled
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.9) (push) Has been cancelled
S3 Proxy Signature Tests / S3 Proxy Signature Verification Tests (push) Has been cancelled
S3 SDK V2 Route Disambiguation Tests / S3 SDK V2 Routing Tests (push) Has been cancelled
S3 SSE Tests / S3 SSE Integration Tests (quick) (push) Has been cancelled
S3 SSE Tests / S3 SSE Compatibility Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Metadata Persistence Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Copy Operations Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Multipart Upload Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Performance Test (push) Has been cancelled
S3 SSE Tests / S3 Volume Encryption Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Integration Tests (comprehensive) (push) Has been cancelled
Spark Integration Tests / Spark Integration Tests (push) Failing after 2s
Telemetry Integration Tests / telemetry-integration-test (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 1) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 1) (push) Failing after 3s
Ceph S3 tests / Basic S3 tests (KV store) (push) Has been cancelled
Ceph S3 tests / S3 Versioning & Object Lock tests (push) Has been cancelled
S3 Policy Integration Tests / S3 Policy Unit Tests (push) Failing after 3s
Ceph S3 tests / S3 CORS tests (push) Has been cancelled
Ceph S3 tests / SeaweedFS Custom S3 Copy tests (push) Has been cancelled
Ceph S3 tests / Basic S3 tests (SQL store) (push) Has been cancelled
test s3 over https using aws-cli / awscli-tests (push) Has been cancelled
TLS Rotation Integration Tests / TLS Rotation Integration Tests (push) Has been cancelled
Vacuum Integration Tests / Vacuum Integration Tests (push) Has been cancelled
S3 Policy Integration Tests / S3 Policy Variables Integration Tests (push) Failing after 3s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (advanced-policy) (push) Failing after 3s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (basic-policy) (push) Failing after 2s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (contextual-policy) (push) Failing after 4s
S3 Policy Integration Tests / Trusted Proxy Detection Tests (push) Failing after 2s
S3 Policy Integration Tests / S3 Policy Shell Integration Tests (push) Failing after 3s
Samba on FUSE Integration / samba-integration (push) Failing after 2s
SFTP Integration Tests / SFTP Integration Testing (push) Failing after 3s
terraform: validate and test modules / fmt, validate, plan-level tests (push) Failing after 3s
terraform: validate and test modules / local cluster smoke test (real weed) (push) Failing after 2s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 3) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 2) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 2) (push) Failing after 2s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 3) (push) Failing after 5s
mount: windows conformance / WinFsp conformance (push) Has been cancelled
mount: windows / Mount on Windows (push) Has been cancelled

This commit is contained in:
upstream-mirror
2026-09-16 03:23:42 +05:00
commit 89f5be005d
3997 changed files with 961970 additions and 0 deletions

1804
seaweed-volume/src/config.rs Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,275 @@
//! JPEG EXIF orientation auto-fix, matching Go's `FixJpgOrientation`.
//!
//! Reads the EXIF orientation tag from JPEG data and rotates/flips the image
//! to normalize it to orientation 1 (top-left). If EXIF parsing fails or
//! orientation is already normal, returns the original data unchanged.
use std::io::Cursor;
use image::{DynamicImage, GenericImageView, ImageFormat, RgbaImage};
/// EXIF orientation tag values.
/// See: <http://sylvana.net/jpegcrop/exif_orientation.html>
const TOP_LEFT_SIDE: u32 = 1;
const TOP_RIGHT_SIDE: u32 = 2;
const BOTTOM_RIGHT_SIDE: u32 = 3;
const BOTTOM_LEFT_SIDE: u32 = 4;
const LEFT_SIDE_TOP: u32 = 5;
const RIGHT_SIDE_TOP: u32 = 6;
const RIGHT_SIDE_BOTTOM: u32 = 7;
const LEFT_SIDE_BOTTOM: u32 = 8;
/// Fix JPEG orientation based on EXIF data.
///
/// Reads the EXIF orientation tag and applies the appropriate rotation/flip
/// to normalize the image to orientation 1 (top-left). Re-encodes as JPEG.
///
/// Returns the original data unchanged if:
/// - EXIF data cannot be parsed
/// - No orientation tag is present
/// - Orientation is already 1 (normal)
/// - Image decoding or re-encoding fails
pub fn fix_jpg_orientation(data: &[u8]) -> Vec<u8> {
// Parse EXIF data
let orientation = match read_exif_orientation(data) {
Some(o) => o,
None => return data.to_vec(),
};
// Orientation 1 means normal — no transformation needed
if orientation == TOP_LEFT_SIDE {
return data.to_vec();
}
// Determine rotation angle and flip mode
let (angle, flip_horizontal) = match orientation {
TOP_RIGHT_SIDE => (0, true),
BOTTOM_RIGHT_SIDE => (180, false),
BOTTOM_LEFT_SIDE => (180, true),
LEFT_SIDE_TOP => (-90, true),
RIGHT_SIDE_TOP => (-90, false),
RIGHT_SIDE_BOTTOM => (90, true),
LEFT_SIDE_BOTTOM => (90, false),
_ => return data.to_vec(),
};
// Decode the image
let src_image = match image::load_from_memory_with_format(data, ImageFormat::Jpeg) {
Ok(img) => img,
Err(_) => return data.to_vec(),
};
// Apply rotation then flip (matching Go's flip(rotate(img, angle), flipMode))
let transformed = flip_horizontal_if(rotate(src_image, angle), flip_horizontal);
// Re-encode as JPEG
let mut buf = Cursor::new(Vec::new());
match transformed.write_to(&mut buf, ImageFormat::Jpeg) {
Ok(_) => buf.into_inner(),
Err(_) => data.to_vec(),
}
}
/// Read the EXIF orientation tag from JPEG data.
/// Returns None if EXIF cannot be parsed or orientation tag is not present.
fn read_exif_orientation(data: &[u8]) -> Option<u32> {
let exif_reader = exif::Reader::new();
let mut cursor = Cursor::new(data);
let exif_data = exif_reader.read_from_container(&mut cursor).ok()?;
let orientation_field = exif_data.get_field(exif::Tag::Orientation, exif::In::PRIMARY)?;
match orientation_field.value {
exif::Value::Short(ref v) if !v.is_empty() => Some(v[0] as u32),
_ => orientation_field.value.get_uint(0),
}
}
/// Rotate an image by the given angle (counter-clockwise, in degrees).
/// Matches Go's rotate function.
fn rotate(img: DynamicImage, angle: i32) -> DynamicImage {
let (width, height) = img.dimensions();
match angle {
90 => {
// 90 degrees counter-clockwise
let new_w = height;
let new_h = width;
let mut out = RgbaImage::new(new_w, new_h);
for y in 0..new_h {
for x in 0..new_w {
out.put_pixel(x, y, img.get_pixel(new_h - 1 - y, x));
}
}
DynamicImage::ImageRgba8(out)
}
-90 => {
// 90 degrees clockwise (or 270 counter-clockwise)
let new_w = height;
let new_h = width;
let mut out = RgbaImage::new(new_w, new_h);
for y in 0..new_h {
for x in 0..new_w {
out.put_pixel(x, y, img.get_pixel(y, new_w - 1 - x));
}
}
DynamicImage::ImageRgba8(out)
}
180 | -180 => {
let mut out = RgbaImage::new(width, height);
for y in 0..height {
for x in 0..width {
out.put_pixel(x, y, img.get_pixel(width - 1 - x, height - 1 - y));
}
}
DynamicImage::ImageRgba8(out)
}
_ => img,
}
}
/// Flip the image horizontally if requested.
/// In Go, flipMode 2 == FlipHorizontal. We simplify since only horizontal flip is used.
fn flip_horizontal_if(img: DynamicImage, do_flip: bool) -> DynamicImage {
if !do_flip {
return img;
}
let (width, height) = img.dimensions();
let mut out = RgbaImage::new(width, height);
for y in 0..height {
for x in 0..width {
out.put_pixel(x, y, img.get_pixel(width - 1 - x, y));
}
}
DynamicImage::ImageRgba8(out)
}
/// Returns true if the given MIME type or file path extension indicates a JPEG file.
pub fn is_jpeg(mime_type: &str, path: &str) -> bool {
if mime_type == "image/jpeg" {
return true;
}
let lower = path.to_lowercase();
lower.ends_with(".jpg") || lower.ends_with(".jpeg")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_non_jpeg_data_returned_unchanged() {
let data = b"not a jpeg file at all";
let result = fix_jpg_orientation(data);
assert_eq!(result, data);
}
#[test]
fn test_jpeg_without_exif_returned_unchanged() {
// Create a minimal JPEG without EXIF data
let img = DynamicImage::ImageRgba8(RgbaImage::new(2, 2));
let mut buf = Cursor::new(Vec::new());
img.write_to(&mut buf, ImageFormat::Jpeg).unwrap();
let jpeg_data = buf.into_inner();
let result = fix_jpg_orientation(&jpeg_data);
// Should return data unchanged (no EXIF orientation tag)
// Just verify it's still valid JPEG
assert!(!result.is_empty());
assert_eq!(&result[0..2], &[0xFF, 0xD8]); // JPEG magic bytes
}
#[test]
fn test_is_jpeg() {
assert!(is_jpeg("image/jpeg", ""));
assert!(is_jpeg("", "/3,abc.jpg"));
assert!(is_jpeg("", "/3,abc.JPEG"));
assert!(is_jpeg("application/octet-stream", "/3,abc.JPG"));
assert!(!is_jpeg("image/png", "/3,abc.png"));
assert!(!is_jpeg("", "/3,abc.png"));
}
#[test]
fn test_rotate_180() {
// Create a 2x2 image with distinct pixel colors
let mut img = RgbaImage::new(2, 2);
img.put_pixel(0, 0, image::Rgba([255, 0, 0, 255])); // red top-left
img.put_pixel(1, 0, image::Rgba([0, 255, 0, 255])); // green top-right
img.put_pixel(0, 1, image::Rgba([0, 0, 255, 255])); // blue bottom-left
img.put_pixel(1, 1, image::Rgba([255, 255, 0, 255])); // yellow bottom-right
let dynamic = DynamicImage::ImageRgba8(img);
let rotated = rotate(dynamic, 180);
let (w, h) = rotated.dimensions();
assert_eq!((w, h), (2, 2));
// After 180 rotation: top-left should be yellow, top-right should be blue
assert_eq!(rotated.get_pixel(0, 0), image::Rgba([255, 255, 0, 255]));
assert_eq!(rotated.get_pixel(1, 0), image::Rgba([0, 0, 255, 255]));
assert_eq!(rotated.get_pixel(0, 1), image::Rgba([0, 255, 0, 255]));
assert_eq!(rotated.get_pixel(1, 1), image::Rgba([255, 0, 0, 255]));
}
#[test]
fn test_rotate_90_ccw() {
// Create 3x2 image (width=3, height=2)
let mut img = RgbaImage::new(3, 2);
img.put_pixel(0, 0, image::Rgba([1, 0, 0, 255]));
img.put_pixel(1, 0, image::Rgba([2, 0, 0, 255]));
img.put_pixel(2, 0, image::Rgba([3, 0, 0, 255]));
img.put_pixel(0, 1, image::Rgba([4, 0, 0, 255]));
img.put_pixel(1, 1, image::Rgba([5, 0, 0, 255]));
img.put_pixel(2, 1, image::Rgba([6, 0, 0, 255]));
let dynamic = DynamicImage::ImageRgba8(img);
let rotated = rotate(dynamic, 90);
let (w, h) = rotated.dimensions();
// 90 CCW: width=3,height=2 -> new_w=2, new_h=3
assert_eq!((w, h), (2, 3));
// Top-right (2,0) should move to top-left (0,0) in CCW 90
assert_eq!(rotated.get_pixel(0, 0)[0], 3);
assert_eq!(rotated.get_pixel(1, 0)[0], 6);
}
#[test]
fn test_rotate_neg90_cw() {
// Create 3x2 image
let mut img = RgbaImage::new(3, 2);
img.put_pixel(0, 0, image::Rgba([1, 0, 0, 255]));
img.put_pixel(1, 0, image::Rgba([2, 0, 0, 255]));
img.put_pixel(2, 0, image::Rgba([3, 0, 0, 255]));
img.put_pixel(0, 1, image::Rgba([4, 0, 0, 255]));
img.put_pixel(1, 1, image::Rgba([5, 0, 0, 255]));
img.put_pixel(2, 1, image::Rgba([6, 0, 0, 255]));
let dynamic = DynamicImage::ImageRgba8(img);
let rotated = rotate(dynamic, -90);
let (w, h) = rotated.dimensions();
assert_eq!((w, h), (2, 3));
// -90 (CW 90): top-left (0,0) should go to top-right
assert_eq!(rotated.get_pixel(0, 0)[0], 4);
assert_eq!(rotated.get_pixel(1, 0)[0], 1);
}
#[test]
fn test_flip_horizontal() {
let mut img = RgbaImage::new(2, 1);
img.put_pixel(0, 0, image::Rgba([10, 0, 0, 255]));
img.put_pixel(1, 0, image::Rgba([20, 0, 0, 255]));
let dynamic = DynamicImage::ImageRgba8(img);
let flipped = flip_horizontal_if(dynamic, true);
assert_eq!(flipped.get_pixel(0, 0)[0], 20);
assert_eq!(flipped.get_pixel(1, 0)[0], 10);
}
#[test]
fn test_flip_horizontal_noop() {
let mut img = RgbaImage::new(2, 1);
img.put_pixel(0, 0, image::Rgba([10, 0, 0, 255]));
img.put_pixel(1, 0, image::Rgba([20, 0, 0, 255]));
let dynamic = DynamicImage::ImageRgba8(img);
let not_flipped = flip_horizontal_if(dynamic, false);
assert_eq!(not_flipped.get_pixel(0, 0)[0], 10);
assert_eq!(not_flipped.get_pixel(1, 0)[0], 20);
}
}

28
seaweed-volume/src/lib.rs Normal file
View File

@@ -0,0 +1,28 @@
pub mod config;
pub mod images;
pub mod malloc_tuning;
pub mod metrics;
pub mod remote_storage;
pub mod security;
pub mod server;
pub mod storage;
pub mod version;
/// Generated protobuf modules.
pub mod pb {
pub const FILE_DESCRIPTOR_SET: &[u8] =
tonic::include_file_descriptor_set!("seaweed_descriptor");
pub mod remote_pb {
tonic::include_proto!("remote_pb");
}
pub mod volume_server_pb {
tonic::include_proto!("volume_server_pb");
}
pub mod master_pb {
tonic::include_proto!("master_pb");
}
pub mod filer_pb {
tonic::include_proto!("filer_pb");
}
}

1135
seaweed-volume/src/main.rs Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,467 @@
//! Keep glibc from silently converting large short-lived buffers into heap the
//! process never gives back.
//!
//! glibc serves an allocation with `mmap` when it is at least
//! `M_MMAP_THRESHOLD` (128 KiB by default), and `munmap`s it on free, so the
//! pages go straight back to the OS. That threshold is **adaptive**: whenever a
//! block that came from `mmap` is freed, glibc raises the threshold to that
//! block's size — up to 32 MiB — on the theory that a workload repeatedly
//! allocating buffers of that size is better served from the heap.
//!
//! For a volume server that theory is wrong in a specific, expensive way. EC
//! reconstruction and needle reassembly allocate large, short-lived buffers.
//! The first few are mmap'd and freed, which trains the threshold upward; every
//! later buffer of that size is then carved out of the heap instead. Heap
//! memory is only returned to the OS from the top of the arena, so those pages
//! stay resident as anonymous memory for the life of the process. They are
//! still *reusable* — this is not a leak, and a repeat workload does not grow
//! the footprint further — but under a hard cgroup `MemoryMax` they are
//! indistinguishable from a leak, because anonymous pages cannot be reclaimed
//! under pressure the way page cache can. The retained footprint eats exactly
//! the headroom that a burst of maintenance work needs, and the process is
//! OOM-killed while most of its resident memory is free-but-unreturned.
//!
//! Measured on a 17-node cluster (EC 10+4, `--index=redb`), one node, two
//! identical `ec.scrub -mode full` rounds over 10912 EC files each, comparing
//! the same unit restarted with and without a pinned threshold:
//!
//! | | baseline | round 1 | round 2 | 60s idle |
//! |---|---|---|---|---|
//! | default (adaptive) | 10 MB | 84 MB | 88 MB | **88 MB** |
//! | pinned threshold | 10 MB | 13 MB | 14 MB | **14 MB** |
//!
//! 78 MB retained versus 4 MB for identical work. On that cluster's heavier
//! mixed scrub workloads the same effect reached ~600 MB of retained anonymous
//! memory per volume server, against a 3 GiB cap.
//!
//! Calling `mallopt(M_MMAP_THRESHOLD, ...)` sets the threshold *and* disables
//! the dynamic adjustment, which is the documented behaviour of setting it
//! explicitly. We pin it to glibc's own default rather than inventing a value:
//! the goal is to stop the adaptation, not to second-guess the default.
/// glibc's own default `M_MMAP_THRESHOLD`. Pinning to this value changes
/// nothing about which allocations use `mmap` on a freshly started process; it
/// only prevents the threshold from drifting upward later.
#[cfg(all(target_os = "linux", target_env = "gnu"))]
const DEFAULT_MMAP_THRESHOLD: libc::c_int = 128 * 1024;
/// Legacy environment variable glibc reads for the same setting. If an operator
/// has set it, honour their value and do not override it.
pub const MMAP_THRESHOLD_ENV: &str = "MALLOC_MMAP_THRESHOLD_";
/// Modern glibc tunables environment variable. Operators may set the threshold
/// via `GLIBC_TUNABLES=glibc.malloc.mmap_threshold=...` instead of the legacy
/// variable; that override is honoured too.
pub const GLIBC_TUNABLES_ENV: &str = "GLIBC_TUNABLES";
/// The tunable name within `GLIBC_TUNABLES` that maps to `M_MMAP_THRESHOLD`.
#[cfg(all(target_os = "linux", target_env = "gnu"))]
const MMAP_THRESHOLD_TUNABLE: &str = "glibc.malloc.mmap_threshold";
/// Outcome of the tuning attempt, so the caller can log it and tests can assert
/// on it without inspecting global allocator state.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MallocTuning {
/// Threshold pinned to `DEFAULT_MMAP_THRESHOLD`; dynamic adjustment is off.
Pinned(i32),
/// An allocator override (`MALLOC_MMAP_THRESHOLD_` or
/// `GLIBC_TUNABLES=glibc.malloc.mmap_threshold=...`) was set, so the
/// operator's value wins.
DeferredToEnv,
/// `mallopt` reported failure. Not fatal — the server runs, it just keeps
/// glibc's adaptive behaviour.
Failed,
/// Not glibc, so there is no adaptive threshold to pin.
NotApplicable,
}
/// Pin glibc's mmap threshold unless the operator has set an allocator override.
/// Safe to call more than once; call it before serving traffic, since the point
/// is to prevent the threshold from being trained upward by early allocations.
pub fn pin_mmap_threshold() -> MallocTuning {
pin_mmap_threshold_inner()
}
#[cfg(all(target_os = "linux", target_env = "gnu"))]
fn pin_mmap_threshold_inner() -> MallocTuning {
if operator_mmap_threshold_override_active() {
return MallocTuning::DeferredToEnv;
}
// SAFETY: `mallopt` is a libc entry point that takes two ints and mutates
// only allocator-internal tunables. It has no preconditions and no effect
// on memory this process already owns.
let rc = unsafe { libc::mallopt(libc::M_MMAP_THRESHOLD, DEFAULT_MMAP_THRESHOLD) };
if rc == 1 {
MallocTuning::Pinned(DEFAULT_MMAP_THRESHOLD)
} else {
MallocTuning::Failed
}
}
#[cfg(not(all(target_os = "linux", target_env = "gnu")))]
fn pin_mmap_threshold_inner() -> MallocTuning {
MallocTuning::NotApplicable
}
#[cfg(all(target_os = "linux", target_env = "gnu"))]
fn operator_mmap_threshold_override_active() -> bool {
// MALLOC_MMAP_THRESHOLD_: glibc calls atoi(value) then mallopt, which
// always sets the threshold and disables dynamic adjustment — even for
// empty, negative, or non-numeric values (atoi returns 0). So any presence
// of the variable means the operator's override is in effect.
std::env::var_os(MMAP_THRESHOLD_ENV).is_some()
|| usable_glibc_tunable_threshold(std::env::var_os(GLIBC_TUNABLES_ENV))
}
/// Look for `glibc.malloc.mmap_threshold=<value>` among the colon-separated
/// tunables in `GLIBC_TUNABLES`. glibc's `parse_tunables_string` (elf/dl-tunables.c)
/// rejects the **entire** string (returns -1) if it reaches `\0` before finding
/// `=` in a name (last entry has no `=`), or if any entry's value contains a
/// duplicate `=`. When `parse_tunables_string` returns -1, `parse_tunables`
/// prints a warning and returns immediately without applying ANY tunable —
/// including ones already parsed into the tunables array. We match that by
/// returning `false` for the entire string on any of those conditions.
///
/// glibc parses tunable values with `_dl_strtoul`, which accepts decimal,
/// `0x` hex, `0` octal, an optional sign (negatives wrap to `unsigned long`),
/// and requires the entire value to be consumed; we match that with
/// `dl_strtoul_consumes_all`.
#[cfg(all(target_os = "linux", target_env = "gnu"))]
fn usable_glibc_tunable_threshold(tunables: Option<std::ffi::OsString>) -> bool {
let s = match tunables.and_then(|v| v.into_string().ok()) {
Some(s) => s,
None => return false,
};
if s.is_empty() {
return false;
}
// Parse the string character-by-character, matching glibc's
// parse_tunables_string logic exactly. Using split(':') would lose the
// distinction between an entry terminated by ':' (skip) and one terminated
// by '\0' with no '=' (reject entire string).
let bytes = s.as_bytes();
let mut pos = 0;
let mut found_threshold = false;
loop {
// Find where the name ends ('=', ':', or end of string).
let name_start = pos;
while pos < bytes.len() && bytes[pos] != b'=' && bytes[pos] != b':' {
pos += 1;
}
// End of string before '=' → glibc returns -1 (reject entire string).
if pos >= bytes.len() {
return false;
}
// ':' before '=' → glibc skips this entry and continues.
if bytes[pos] == b':' {
pos += 1;
continue;
}
// Skip the '='.
let name_end = pos;
pos += 1;
// Find where the value ends ('=', ':', or end of string).
let val_start = pos;
while pos < bytes.len() && bytes[pos] != b'=' && bytes[pos] != b':' {
pos += 1;
}
// '=' in value → glibc returns -1 (reject entire string).
if pos < bytes.len() && bytes[pos] == b'=' {
return false;
}
let key = &s[name_start..name_end];
let val = &s[val_start..pos];
if key == MMAP_THRESHOLD_TUNABLE && dl_strtoul_consumes_all(val) {
found_threshold = true;
}
// End of string → done.
if pos >= bytes.len() {
break;
}
// Skip the ':'.
pos += 1;
}
found_threshold
}
/// Replicate glibc's `_dl_strtoul` (elf/dl-misc.c) just enough to determine
/// whether it would consume the entire string — which is what
/// `tunable_parse_num` checks (`endptr == strval + len`). Returns `true` if
/// glibc would accept the value and apply it.
///
/// `_dl_strtoul` skips leading spaces/tabs, accepts an optional `+`/`-` sign,
/// and parses `0x`-prefixed hex, `0`-prefixed octal, or plain decimal. A
/// negative result wraps to `unsigned long` (`-1` → `SIZE_MAX`). If no digit is
/// found after the sign, the end pointer stays at the current position — which
/// still counts as "consumed" when the string is empty or whitespace-only
/// (value 0). On overflow, `_dl_strtoul` stops at the overflowing digit (endptr
/// does not reach the end), so `tunable_parse_num` rejects the value.
#[cfg(all(target_os = "linux", target_env = "gnu"))]
fn dl_strtoul_consumes_all(s: &str) -> bool {
let bytes = s.as_bytes();
let mut pos = 0;
// Skip leading whitespace (spaces and tabs, matching _dl_strtoul).
while pos < bytes.len() && (bytes[pos] == b' ' || bytes[pos] == b'\t') {
pos += 1;
}
// Optional sign.
if pos < bytes.len() && (bytes[pos] == b'-' || bytes[pos] == b'+') {
pos += 1;
}
// Must have at least one digit (0-9) to start parsing, unless we're already
// at the end (empty / whitespace-only / sign-only → value 0, consumed).
if pos >= bytes.len() {
return true;
}
if bytes[pos] < b'0' || bytes[pos] > b'9' {
return false;
}
// Determine base: 0x → hex, 0 → octal, else decimal. _dl_strtoul unconditionally
// advances past "0x"/"0X" when the first char is '0' and the next is 'x'/'X',
// even if no hex digit follows — in that case the digit loop breaks immediately,
// endptr reaches the end, and the value is 0.
let base: u32 = if bytes[pos] == b'0'
&& pos + 1 < bytes.len()
&& (bytes[pos + 1] == b'x' || bytes[pos + 1] == b'X')
{
pos += 2; // skip "0x"
16
} else if bytes[pos] == b'0' {
8
} else {
10
};
// Parse digits with overflow detection, matching _dl_strtoul's cutoff/cutlim
// logic. On overflow, _dl_strtoul sets endptr to the overflowing digit and
// returns UINT64_MAX — so the value is NOT fully consumed and
// tunable_parse_num rejects it.
let mut result: u64 = 0;
let cutoff = u64::MAX / base as u64;
let cutlim = u64::MAX % base as u64;
while pos < bytes.len() {
let b = bytes[pos];
let digval: u32 = match digit_value(b, base) {
Some(v) => v,
None => break,
};
if result > cutoff || (result == cutoff && digval as u64 > cutlim) {
// Overflow: _dl_strtoul stops here, endptr points at this digit.
return false;
}
result *= base as u64;
result += digval as u64;
pos += 1;
}
// The entire string must be consumed (matching tunable_parse_num's check).
pos == bytes.len()
}
/// Returns the numeric value of a digit byte in the given base, or `None` if
/// the byte is not a valid digit in that base.
#[cfg(all(target_os = "linux", target_env = "gnu"))]
fn digit_value(b: u8, base: u32) -> Option<u32> {
if (b'0'..=b'0' + (base - 1).min(9) as u8).contains(&b) {
return Some((b - b'0') as u32);
}
if base == 16 {
if (b'a'..=b'f').contains(&b) {
return Some((b - b'a' + 10) as u32);
}
if (b'A'..=b'F').contains(&b) {
return Some((b - b'A' + 10) as u32);
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn env_override_constants_match_glibc_names() {
// Verified against the real accessor rather than a copy of the name, so
// renaming the constant cannot silently break the override contract.
assert_eq!(MMAP_THRESHOLD_ENV, "MALLOC_MMAP_THRESHOLD_");
assert_eq!(GLIBC_TUNABLES_ENV, "GLIBC_TUNABLES");
}
#[test]
fn calling_twice_is_stable() {
// Startup paths get re-entered in tests and in `weed mini`; the second
// call must not report a different outcome from the first.
let first = pin_mmap_threshold();
let second = pin_mmap_threshold();
assert_eq!(first, second);
}
#[cfg(all(target_os = "linux", target_env = "gnu"))]
#[test]
fn pins_threshold_on_glibc_when_no_override_is_set() {
// The env override is not set in the test process, so this exercises the
// mallopt path. If an override happens to be present, defer to it.
if operator_mmap_threshold_override_active() {
assert_eq!(pin_mmap_threshold(), MallocTuning::DeferredToEnv);
return;
}
assert_eq!(
pin_mmap_threshold(),
MallocTuning::Pinned(DEFAULT_MMAP_THRESHOLD),
"mallopt(M_MMAP_THRESHOLD) should succeed on glibc"
);
}
#[cfg(not(all(target_os = "linux", target_env = "gnu")))]
#[test]
fn is_a_noop_off_glibc() {
// No glibc adaptive threshold exists off glibc, so there is nothing to
// pin regardless of any environment variables that happen to be set.
assert_eq!(pin_mmap_threshold(), MallocTuning::NotApplicable);
}
#[cfg(all(target_os = "linux", target_env = "gnu"))]
#[test]
fn dl_strtoul_consumes_all_matches_glibc_parser() {
// Decimal — any non-empty decimal integer is accepted, including
// negative (wraps to unsigned) and zero.
assert!(dl_strtoul_consumes_all("131072"));
assert!(dl_strtoul_consumes_all("0"));
assert!(dl_strtoul_consumes_all("-1"));
assert!(dl_strtoul_consumes_all("-131072"));
// Values above i64::MAX are valid for glibc's unsigned parser.
assert!(dl_strtoul_consumes_all("9223372036854775808"));
// Hex with 0x prefix.
assert!(dl_strtoul_consumes_all("0x20000"));
assert!(dl_strtoul_consumes_all("0X20000"));
assert!(dl_strtoul_consumes_all("0x0"));
// Octal with leading 0.
assert!(dl_strtoul_consumes_all("010"));
// Leading whitespace (spaces and tabs) is skipped.
assert!(dl_strtoul_consumes_all(" 131072"));
assert!(dl_strtoul_consumes_all("\t0x20000"));
// Empty and whitespace-only strings are accepted (value 0).
assert!(dl_strtoul_consumes_all(""));
assert!(dl_strtoul_consumes_all(" "));
assert!(dl_strtoul_consumes_all("\t"));
// Sign-only strings are accepted: _dl_strtoul skips the sign, finds no
// digit, sets endptr to the position after the sign (== end of string),
// and returns 0. tunable_parse_num sees endptr == strval + len → true.
assert!(dl_strtoul_consumes_all("-"));
assert!(dl_strtoul_consumes_all("+"));
// Trailing garbage is rejected — _dl_strtoul stops at the first
// non-digit and tunable_parse_num requires the entire string consumed.
assert!(!dl_strtoul_consumes_all("131072abc"));
// In hex mode, a-f are digits, so "0x20000abc" is a valid hex number.
// Use a non-hex character like 'g' to test trailing garbage in hex.
assert!(!dl_strtoul_consumes_all("0x20000g"));
assert!(!dl_strtoul_consumes_all("128K"));
// Non-numeric strings are rejected.
assert!(!dl_strtoul_consumes_all("abc"));
// "0x" with no hex digits: _dl_strtoul advances past "0x", the digit loop
// breaks immediately (no hex digit), endptr reaches the end, value is 0.
// tunable_parse_num accepts it.
assert!(dl_strtoul_consumes_all("0x"));
assert!(dl_strtoul_consumes_all("0X"));
// Overflow: _dl_strtoul stops at the overflowing digit (endptr points
// there, not at the end), so tunable_parse_num rejects the value.
assert!(!dl_strtoul_consumes_all("18446744073709551616")); // u64::MAX + 1
assert!(!dl_strtoul_consumes_all("99999999999999999999")); // 20 nines
assert!(!dl_strtoul_consumes_all("0x10000000000000000")); // 2^64
// u64::MAX itself is accepted: the last digit (5) equals cutlim (=5),
// so the overflow check (digval > cutlim) is false.
assert!(dl_strtoul_consumes_all("18446744073709551615")); // u64::MAX
}
#[cfg(all(target_os = "linux", target_env = "gnu"))]
#[test]
fn usable_glibc_tunable_threshold_detects_mmap_threshold() {
// Decimal, hex, octal, negative, and zero values are all accepted by
// glibc's _dl_strtoul and cause the threshold to be pinned.
assert!(usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=131072".into()
)));
assert!(usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=0x20000".into()
)));
assert!(usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=0".into()
)));
assert!(usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=-1".into()
)));
assert!(usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=9223372036854775808".into()
)));
// u64::MAX is accepted by _dl_strtoul (last digit == cutlim, no overflow).
assert!(usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=18446744073709551615".into()
)));
// Appears alongside other tunables.
assert!(usable_glibc_tunable_threshold(Some(
"glibc.cpu.x=1:glibc.malloc.mmap_threshold=131072".into()
)));
// Leading ':' is accepted — glibc skips the empty entry and continues.
assert!(usable_glibc_tunable_threshold(Some(
":glibc.malloc.mmap_threshold=131072".into()
)));
// Empty value is accepted by _dl_strtoul (value 0).
assert!(usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=".into()
)));
// Non-numeric values are rejected by _dl_strtoul.
assert!(!usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=abc".into()
)));
assert!(!usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=128K".into()
)));
// A malformed sibling entry (duplicate '=') makes glibc reject the
// entire string, so we must not accept the threshold entry either.
// This applies regardless of whether the threshold is before or after
// the malformed entry — parse_tunables_string returns -1, and
// parse_tunables discards all tunables without applying any.
assert!(!usable_glibc_tunable_threshold(Some(
"glibc.malloc.check=2=2:glibc.malloc.mmap_threshold=131072".into()
)));
assert!(!usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=262144:glibc.malloc.check=2=2".into()
)));
// A trailing entry with no '=' makes glibc reject the entire string
// (parse_tunables_string hits '\0' before '=' and returns -1).
assert!(!usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=262144:glibc.cpu.x".into()
)));
// A trailing ':' makes glibc reject the entire string (the empty entry
// after ':' hits '\0' before '=' and returns -1).
assert!(!usable_glibc_tunable_threshold(Some(
"glibc.malloc.mmap_threshold=262144:".into()
)));
// Unrelated tunables do not count.
assert!(!usable_glibc_tunable_threshold(Some(
"glibc.cpu.x=1".into()
)));
assert!(!usable_glibc_tunable_threshold(None));
}
}

View File

@@ -0,0 +1,621 @@
//! Prometheus metrics for the volume server.
//!
//! Mirrors the Go SeaweedFS volume server metrics.
use prometheus::{
self, Encoder, GaugeVec, HistogramOpts, HistogramVec, IntCounter, IntCounterVec, IntGauge,
IntGaugeVec, Opts, Registry, TextEncoder,
};
use std::sync::{LazyLock, Once};
use crate::version;
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct PushGatewayConfig {
pub address: String,
pub interval_seconds: u32,
}
pub static REGISTRY: LazyLock<Registry> = LazyLock::new(Registry::new);
// ---- Request metrics (Go: VolumeServerRequestCounter, VolumeServerRequestHistogram) ----
/// Request counter with labels `type` (HTTP method) and `code` (HTTP status).
pub static REQUEST_COUNTER: LazyLock<IntCounterVec> = LazyLock::new(|| {
IntCounterVec::new(
Opts::new(
"SeaweedFS_volumeServer_request_total",
"Volume server requests",
),
&["type", "code"],
)
.expect("metric can be created")
});
/// Request duration histogram with label `type` (HTTP method).
pub static REQUEST_DURATION: LazyLock<HistogramVec> = LazyLock::new(|| {
HistogramVec::new(
HistogramOpts::new(
"SeaweedFS_volumeServer_request_seconds",
"Volume server request duration in seconds",
)
.buckets(exponential_buckets(0.0001, 2.0, 24)),
&["type"],
)
.expect("metric can be created")
});
// ---- Handler counters (Go: VolumeServerHandlerCounter) ----
/// Handler-level operation counter with label `type`.
pub static HANDLER_COUNTER: LazyLock<IntCounterVec> = LazyLock::new(|| {
IntCounterVec::new(
Opts::new(
"SeaweedFS_volumeServer_handler_total",
"Volume server handler counters",
),
&["type"],
)
.expect("metric can be created")
});
// ---- Vacuuming metrics (Go: VolumeServerVacuuming*) ----
/// Vacuuming compact counter with label `success` (true/false).
pub static VACUUMING_COMPACT_COUNTER: LazyLock<IntCounterVec> = LazyLock::new(|| {
IntCounterVec::new(
Opts::new(
"SeaweedFS_volumeServer_vacuuming_compact_count",
"Counter of volume vacuuming Compact counter",
),
&["success"],
)
.expect("metric can be created")
});
/// Vacuuming commit counter with label `success` (true/false).
pub static VACUUMING_COMMIT_COUNTER: LazyLock<IntCounterVec> = LazyLock::new(|| {
IntCounterVec::new(
Opts::new(
"SeaweedFS_volumeServer_vacuuming_commit_count",
"Counter of volume vacuuming commit counter",
),
&["success"],
)
.expect("metric can be created")
});
/// Vacuuming duration histogram with label `type` (compact/commit).
pub static VACUUMING_HISTOGRAM: LazyLock<HistogramVec> = LazyLock::new(|| {
HistogramVec::new(
HistogramOpts::new(
"SeaweedFS_volumeServer_vacuuming_seconds",
"Volume vacuuming duration in seconds",
)
.buckets(exponential_buckets(0.0001, 2.0, 24)),
&["type"],
)
.expect("metric can be created")
});
// ---- Volume gauges (Go: VolumeServerVolumeGauge, VolumeServerReadOnlyVolumeGauge) ----
/// Volumes per collection and type (volume/ec_shards).
pub static VOLUME_GAUGE: LazyLock<GaugeVec> = LazyLock::new(|| {
GaugeVec::new(
Opts::new("SeaweedFS_volumeServer_volumes", "Number of volumes"),
&["collection", "type"],
)
.expect("metric can be created")
});
/// Read-only volumes per collection and type.
pub static READ_ONLY_VOLUME_GAUGE: LazyLock<GaugeVec> = LazyLock::new(|| {
GaugeVec::new(
Opts::new(
"SeaweedFS_volumeServer_read_only_volumes",
"Number of read-only volumes.",
),
&["collection", "type"],
)
.expect("metric can be created")
});
/// Maximum number of volumes this server can hold.
pub static MAX_VOLUMES: LazyLock<IntGauge> = LazyLock::new(|| {
IntGauge::new(
"SeaweedFS_volumeServer_max_volumes",
"Maximum number of volumes",
)
.expect("metric can be created")
});
// ---- Disk size gauges (Go: VolumeServerDiskSizeGauge) ----
/// Actual disk size used by volumes per collection and type (normal/deleted_bytes/ec).
pub static DISK_SIZE_GAUGE: LazyLock<GaugeVec> = LazyLock::new(|| {
GaugeVec::new(
Opts::new(
"SeaweedFS_volumeServer_total_disk_size",
"Actual disk size used by volumes",
),
&["collection", "type"],
)
.expect("metric can be created")
});
// ---- Resource gauges (Go: VolumeServerResourceGauge) ----
/// Disk resource usage per directory and type (all/used/free/avail).
pub static RESOURCE_GAUGE: LazyLock<GaugeVec> = LazyLock::new(|| {
GaugeVec::new(
Opts::new("SeaweedFS_volumeServer_resource", "Server resource usage"),
&["name", "type"],
)
.expect("metric can be created")
});
// ---- In-flight gauges (Go: VolumeServerInFlightRequestsGauge, InFlightDownload/UploadSize) ----
/// In-flight requests per HTTP method.
pub static INFLIGHT_REQUESTS_GAUGE: LazyLock<IntGaugeVec> = LazyLock::new(|| {
IntGaugeVec::new(
Opts::new(
"SeaweedFS_volumeServer_in_flight_requests",
"Current number of in-flight requests being handled by volume server.",
),
&["type"],
)
.expect("metric can be created")
});
/// Concurrent download limit in bytes.
pub static CONCURRENT_DOWNLOAD_LIMIT: LazyLock<IntGauge> = LazyLock::new(|| {
IntGauge::new(
"SeaweedFS_volumeServer_concurrent_download_limit",
"Limit for total concurrent download size in bytes",
)
.expect("metric can be created")
});
/// Concurrent upload limit in bytes.
pub static CONCURRENT_UPLOAD_LIMIT: LazyLock<IntGauge> = LazyLock::new(|| {
IntGauge::new(
"SeaweedFS_volumeServer_concurrent_upload_limit",
"Limit for total concurrent upload size in bytes",
)
.expect("metric can be created")
});
/// Current in-flight download bytes.
pub static INFLIGHT_DOWNLOAD_SIZE: LazyLock<IntGauge> = LazyLock::new(|| {
IntGauge::new(
"SeaweedFS_volumeServer_in_flight_download_size",
"In flight total download size.",
)
.expect("metric can be created")
});
/// Current in-flight upload bytes.
pub static INFLIGHT_UPLOAD_SIZE: LazyLock<IntGauge> = LazyLock::new(|| {
IntGauge::new(
"SeaweedFS_volumeServer_in_flight_upload_size",
"In flight total upload size.",
)
.expect("metric can be created")
});
/// Upload error counter by HTTP status code. Code "0" = transport error (no response).
pub static UPLOAD_ERROR_COUNTER: LazyLock<IntCounterVec> = LazyLock::new(|| {
IntCounterVec::new(
Opts::new("SeaweedFS_upload_error_total",
"Counter of upload errors by HTTP status code. Code 0 means transport error (no response received)."),
&["code"],
).expect("metric can be created")
});
// ---- Scrubbing metrics (Go: VolumeServerScrub*) ----
/// Last scrub execution time, as seconds since UNIX epoch, with label `mode`.
pub static SCRUB_LAST_TIME_SECONDS: LazyLock<GaugeVec> = LazyLock::new(|| {
GaugeVec::new(
Opts::new(
"SeaweedFS_volumeServer_scrub_last_time_seconds",
"Last scrub execution time, as seconds since UNIX epoch.",
),
&["mode"],
)
.expect("metric can be created")
});
/// Counter of overall volumes with issues detected during scrubbing, with label `mode`.
pub static SCRUB_VOLUME_FAILURES: LazyLock<IntCounterVec> = LazyLock::new(|| {
IntCounterVec::new(
Opts::new(
"SeaweedFS_volumeServer_scrub_volume_failures",
"Counter of overall volumes with issues detected during scrubbing.",
),
&["mode"],
)
.expect("metric can be created")
});
/// Counter of overall EC shards with issues detected during scrubbing, with label `mode`.
pub static SCRUB_SHARD_FAILURES: LazyLock<IntCounterVec> = LazyLock::new(|| {
IntCounterVec::new(
Opts::new(
"SeaweedFS_volumeServer_scrub_shard_failures",
"Counter of overall EC shards with issues detected during scrubbing.",
),
&["mode"],
)
.expect("metric can be created")
});
/// Counter of storage read/write EIO errors on volumes and EC shards.
/// Mirrors Go's VolumeServerStorageIoErrorCounter.
pub static STORAGE_IO_ERROR_COUNTER: LazyLock<IntCounter> = LazyLock::new(|| {
IntCounter::new(
"SeaweedFS_volumeServer_storage_io_error_total",
"Counter of storage read/write EIO errors on volumes and EC shards.",
)
.expect("metric can be created")
});
/// Number of volumes quarantined due to storage IO errors.
/// Mirrors Go's VolumeServerIoQuarantineGauge.
pub static IO_QUARANTINE_GAUGE: LazyLock<IntGaugeVec> = LazyLock::new(|| {
IntGaugeVec::new(
Opts::new(
"SeaweedFS_volumeServer_io_quarantine",
"Number of volumes or EC shards quarantined due to storage IO errors.",
),
&["kind"],
)
.expect("metric can be created")
});
// ---- Legacy aliases for backward compat with existing code ----
/// Total number of volumes on this server (flat gauge).
pub static VOLUMES_TOTAL: LazyLock<IntGauge> = LazyLock::new(|| {
IntGauge::new("volume_server_volumes_total", "Total number of volumes")
.expect("metric can be created")
});
/// Disk size in bytes per directory.
pub static DISK_SIZE_BYTES: LazyLock<IntGaugeVec> = LazyLock::new(|| {
IntGaugeVec::new(
Opts::new("volume_server_disk_size_bytes", "Disk size in bytes"),
&["dir"],
)
.expect("metric can be created")
});
/// Disk free bytes per directory.
pub static DISK_FREE_BYTES: LazyLock<IntGaugeVec> = LazyLock::new(|| {
IntGaugeVec::new(
Opts::new("volume_server_disk_free_bytes", "Disk free space in bytes"),
&["dir"],
)
.expect("metric can be created")
});
/// Current number of in-flight requests (flat gauge).
pub static INFLIGHT_REQUESTS: LazyLock<IntGauge> = LazyLock::new(|| {
IntGauge::new(
"volume_server_inflight_requests",
"Current number of in-flight requests",
)
.expect("metric can be created")
});
/// Total number of files stored across all volumes.
pub static VOLUME_FILE_COUNT: LazyLock<IntGauge> = LazyLock::new(|| {
IntGauge::new(
"volume_server_volume_file_count",
"Total number of files stored across all volumes",
)
.expect("metric can be created")
});
// ---- Build info (Go: BuildInfo) ----
/// Build information gauge, always set to 1. Matches Go:
/// Namespace="SeaweedFS", Subsystem="build", Name="info",
/// labels: version, commit, sizelimit, goos, goarch.
pub static BUILD_INFO: LazyLock<GaugeVec> = LazyLock::new(|| {
GaugeVec::new(
Opts::new("SeaweedFS_build_info", "A metric with a constant '1' value labeled by version, commit, sizelimit, goos, and goarch from which SeaweedFS was built."),
&["version", "commit", "sizelimit", "goos", "goarch"],
).expect("metric can be created")
});
/// Generate exponential bucket boundaries for histograms.
fn exponential_buckets(start: f64, factor: f64, count: usize) -> Vec<f64> {
let mut buckets = Vec::with_capacity(count);
let mut val = start;
for _ in 0..count {
buckets.push(val);
val *= factor;
}
buckets
}
// Handler counter type constants (matches Go's metrics_names.go).
pub const WRITE_TO_LOCAL_DISK: &str = "writeToLocalDisk";
pub const WRITE_TO_REPLICAS: &str = "writeToReplicas";
pub const DOWNLOAD_LIMIT_COND: &str = "downloadLimitCondition";
pub const UPLOAD_LIMIT_COND: &str = "uploadLimitCondition";
pub const READ_PROXY_REQ: &str = "readProxyRequest";
pub const READ_REDIRECT_REQ: &str = "readRedirectRequest";
pub const EMPTY_READ_PROXY_LOC: &str = "emptyReadProxyLocaction";
pub const FAILED_READ_PROXY_REQ: &str = "failedReadProxyRequest";
// Error metric name constants.
pub const ERROR_SIZE_MISMATCH_OFFSET_SIZE: &str = "errorSizeMismatchOffsetSize";
pub const ERROR_SIZE_MISMATCH: &str = "errorSizeMismatch";
pub const ERROR_CRC: &str = "errorCRC";
pub const ERROR_INDEX_OUT_OF_RANGE: &str = "errorIndexOutOfRange";
pub const ERROR_GET_NOT_FOUND: &str = "errorGetNotFound";
pub const ERROR_GET_INTERNAL: &str = "errorGetInternal";
pub const ERROR_WRITE_TO_LOCAL_DISK: &str = "errorWriteToLocalDisk";
pub const ERROR_UNMARSHAL_PAIRS: &str = "errorUnmarshalPairs";
pub const ERROR_WRITE_TO_REPLICAS: &str = "errorWriteToReplicas";
// Go volume heartbeat metric label values.
pub const READ_ONLY_LABEL_IS_READ_ONLY: &str = "IsReadOnly";
pub const READ_ONLY_LABEL_NO_WRITE_OR_DELETE: &str = "noWriteOrDelete";
pub const READ_ONLY_LABEL_NO_WRITE_CAN_DELETE: &str = "noWriteCanDelete";
pub const READ_ONLY_LABEL_IS_DISK_SPACE_LOW: &str = "isDiskSpaceLow";
pub const DISK_SIZE_LABEL_NORMAL: &str = "normal";
pub const DISK_SIZE_LABEL_DELETED_BYTES: &str = "deleted_bytes";
pub const DISK_SIZE_LABEL_EC: &str = "ec";
static REGISTER_METRICS: Once = Once::new();
/// Register all metrics with the custom registry.
/// Call this once at startup.
pub fn register_metrics() {
REGISTER_METRICS.call_once(|| {
let metrics: Vec<Box<dyn prometheus::core::Collector>> = vec![
// New Go-compatible metrics
Box::new(REQUEST_COUNTER.clone()),
Box::new(REQUEST_DURATION.clone()),
Box::new(HANDLER_COUNTER.clone()),
Box::new(VACUUMING_COMPACT_COUNTER.clone()),
Box::new(VACUUMING_COMMIT_COUNTER.clone()),
Box::new(VACUUMING_HISTOGRAM.clone()),
Box::new(VOLUME_GAUGE.clone()),
Box::new(READ_ONLY_VOLUME_GAUGE.clone()),
Box::new(MAX_VOLUMES.clone()),
Box::new(DISK_SIZE_GAUGE.clone()),
Box::new(RESOURCE_GAUGE.clone()),
Box::new(INFLIGHT_REQUESTS_GAUGE.clone()),
Box::new(CONCURRENT_DOWNLOAD_LIMIT.clone()),
Box::new(CONCURRENT_UPLOAD_LIMIT.clone()),
Box::new(INFLIGHT_DOWNLOAD_SIZE.clone()),
Box::new(INFLIGHT_UPLOAD_SIZE.clone()),
Box::new(UPLOAD_ERROR_COUNTER.clone()),
Box::new(SCRUB_LAST_TIME_SECONDS.clone()),
Box::new(SCRUB_VOLUME_FAILURES.clone()),
Box::new(SCRUB_SHARD_FAILURES.clone()),
Box::new(STORAGE_IO_ERROR_COUNTER.clone()),
Box::new(IO_QUARANTINE_GAUGE.clone()),
// Legacy metrics
Box::new(VOLUMES_TOTAL.clone()),
Box::new(DISK_SIZE_BYTES.clone()),
Box::new(DISK_FREE_BYTES.clone()),
Box::new(INFLIGHT_REQUESTS.clone()),
Box::new(VOLUME_FILE_COUNT.clone()),
// Build info
Box::new(BUILD_INFO.clone()),
];
for m in metrics {
REGISTRY.register(m).expect("metric registered");
}
// Set build info gauge to 1 with version/commit/sizelimit/os/arch labels (matches Go).
BUILD_INFO
.with_label_values(&[
version::version(),
version::commit(),
version::size_limit(),
std::env::consts::OS,
std::env::consts::ARCH,
])
.set(1.0);
});
}
/// Gather all metrics and encode them in Prometheus text exposition format.
pub fn gather_metrics() -> String {
let encoder = TextEncoder::new();
let metric_families = REGISTRY.gather();
let mut buffer = Vec::new();
encoder
.encode(&metric_families, &mut buffer)
.expect("encoding metrics");
String::from_utf8(buffer).expect("metrics are valid UTF-8")
}
pub fn delete_collection_metrics(collection: &str) {
// Mirrors Go's DeletePartialMatch(prometheus.Labels{"collection": collection})
// which removes ALL metric entries matching the collection label, regardless
// of other label values (like "type"). We gather the metric families to discover
// all type values dynamically, matching Go's partial-match behavior.
delete_partial_match_collection(&VOLUME_GAUGE, collection);
delete_partial_match_collection(&READ_ONLY_VOLUME_GAUGE, collection);
delete_partial_match_collection(&DISK_SIZE_GAUGE, collection);
}
/// Drop a collection's volume server series once its last volume leaves this
/// server. These gauges are only ever set for collections still present, so the
/// values from the heartbeat that saw the last volume would otherwise stand
/// until the process restarts.
pub fn delete_volume_server_collection_metrics(collection: &str) {
let _ = DISK_SIZE_GAUGE.remove_label_values(&[collection, DISK_SIZE_LABEL_NORMAL]);
let _ = DISK_SIZE_GAUGE.remove_label_values(&[collection, DISK_SIZE_LABEL_DELETED_BYTES]);
delete_partial_match_collection(&READ_ONLY_VOLUME_GAUGE, collection);
}
/// Remove all metric entries from a GaugeVec where the "collection" label matches.
/// This emulates Go's `DeletePartialMatch(prometheus.Labels{"collection": collection})`.
fn delete_partial_match_collection(gauge: &GaugeVec, collection: &str) {
use prometheus::core::Collector;
let families = gauge.collect();
for family in &families {
for metric in family.get_metric() {
let labels = metric.get_label();
let mut matches_collection = false;
let mut type_value = None;
for label in labels {
if label.get_name() == "collection" && label.get_value() == collection {
matches_collection = true;
}
if label.get_name() == "type" {
type_value = Some(label.get_value().to_string());
}
}
if matches_collection && let Some(ref tv) = type_value {
let _ = gauge.remove_label_values(&[collection, tv]);
}
}
}
}
pub fn build_pushgateway_url(address: &str, job: &str, instance: &str) -> String {
let base = if address.starts_with("http://") || address.starts_with("https://") {
address.to_string()
} else {
format!("http://{}", address)
};
let base = base.trim_end_matches('/');
format!("{}/metrics/job/{}/instance/{}", base, job, instance)
}
pub async fn push_metrics_once(
client: &reqwest::Client,
address: &str,
job: &str,
instance: &str,
) -> Result<(), String> {
let url = build_pushgateway_url(address, job, instance);
let response = client
.put(&url)
.header(
reqwest::header::CONTENT_TYPE,
"text/plain; version=0.0.4; charset=utf-8",
)
.body(gather_metrics())
.send()
.await
.map_err(|e| format!("push metrics request failed: {}", e))?;
if response.status().is_success() {
Ok(())
} else {
Err(format!(
"push metrics failed with status {}",
response.status()
))
}
}
#[cfg(test)]
mod tests {
use super::*;
use axum::{Router, routing::put};
use std::sync::{Arc, Mutex};
#[test]
fn test_gather_metrics_returns_text() {
register_metrics();
REQUEST_COUNTER.with_label_values(&["GET", "200"]).inc();
let output = gather_metrics();
assert!(output.contains("SeaweedFS_volumeServer_request_total"));
}
#[test]
fn test_build_pushgateway_url() {
assert_eq!(
build_pushgateway_url("localhost:9091", "volumeServer", "test-instance"),
"http://localhost:9091/metrics/job/volumeServer/instance/test-instance"
);
assert_eq!(
build_pushgateway_url("https://push.example", "volumeServer", "node-a"),
"https://push.example/metrics/job/volumeServer/instance/node-a"
);
}
#[tokio::test]
async fn test_push_metrics_once() {
register_metrics();
// A CounterVec with no children emits nothing, so create a labelset
// instead of depending on another test having touched the counter
// first (test order is nondeterministic).
REQUEST_COUNTER.with_label_values(&["GET", "200"]).inc();
let captured = Arc::new(Mutex::new(None::<String>));
let captured_clone = captured.clone();
let app = Router::new().route(
"/metrics/job/volumeServer/instance/test-instance",
put(move |body: String| {
let captured = captured_clone.clone();
async move {
*captured.lock().unwrap() = Some(body);
"ok"
}
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
let client = reqwest::Client::new();
push_metrics_once(
&client,
&format!("127.0.0.1:{}", addr.port()),
"volumeServer",
"test-instance",
)
.await
.unwrap();
let body = captured.lock().unwrap().clone().unwrap();
assert!(body.contains("SeaweedFS_volumeServer_request_total"));
server.abort();
}
#[test]
fn test_delete_collection_metrics_removes_collection_labelsets() {
register_metrics();
VOLUME_GAUGE.with_label_values(&["pics", "volume"]).set(2.0);
VOLUME_GAUGE
.with_label_values(&["pics", "ec_shards"])
.set(3.0);
READ_ONLY_VOLUME_GAUGE
.with_label_values(&["pics", "volume"])
.set(1.0);
DISK_SIZE_GAUGE
.with_label_values(&["pics", "normal"])
.set(10.0);
DISK_SIZE_GAUGE
.with_label_values(&["pics", "deleted_bytes"])
.set(4.0);
delete_collection_metrics("pics");
let output = gather_metrics();
assert!(!output.contains("collection=\"pics\",type=\"volume\""));
assert!(!output.contains("collection=\"pics\",type=\"ec_shards\""));
assert!(!output.contains("collection=\"pics\",type=\"normal\""));
assert!(!output.contains("collection=\"pics\",type=\"deleted_bytes\""));
}
}

View File

@@ -0,0 +1,654 @@
//! SSRF guard for caller-supplied remote storage endpoints.
//!
//! `FetchAndWriteNeedle` accepts an S3 endpoint URL chosen by the caller and
//! dials it directly from the volume server, which typically has network
//! access to cluster-internal hosts. Without validation a caller could point
//! the server at loopback / link-local / RFC 1918 / cloud-metadata addresses
//! and read internal services (SSRF). This mirrors the Go volume server's
//! `validateRemoteEndpoint` (`weed/server/volume_grpc_remote.go`); operators
//! that legitimately fetch from private hosts opt out with
//! `-volume.allowUntrustedRemoteEndpoints`.
//!
//! NOTE: the Go server additionally pins the validated endpoint against DNS
//! rebinding with a custom dialer that re-checks the resolved IP at TCP connect
//! time (`guardedDialer`). The `aws-sdk-s3` client used here builds its own
//! connector, so that connect-time re-validation is not yet ported; the
//! up-front resolve-and-check below still blocks the common SSRF vectors. A
//! narrow TOCTOU window (a hostname that resolves to a public IP here and then
//! flips to a blocked one when the SDK dials) remains as a follow-up.
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
/// AWS/Azure/GCP IPv4 instance-metadata-service (IMDS) address. It is
/// link-local and thus already covered by [`is_link_local`], but is named
/// explicitly so the rejection reason is unambiguous in logs.
const IMDS_IPV4: IpAddr = IpAddr::V4(Ipv4Addr::new(169, 254, 169, 254));
/// Hostnames that target cloud instance metadata services, blocked regardless
/// of how they resolve because some environments alias the IMDS address under
/// a name.
fn is_blocked_imds_host(host: &str) -> bool {
matches!(host, "metadata.google.internal" | "metadata")
}
/// Whether `ip` is in a link-local range: IPv4 169.254.0.0/16, IPv6 fe80::/10
/// (unicast) and interface-local / link-local multicast (scope 1 and 2).
fn is_link_local(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => v4.is_link_local(),
IpAddr::V6(v6) => {
let seg0 = v6.segments()[0];
if (seg0 & 0xffc0) == 0xfe80 {
return true; // fe80::/10 link-local unicast
}
if (seg0 & 0xff00) == 0xff00 {
// ffXS:: multicast; reject interface-local (scope 1) and
// link-local (scope 2) the way Go's IsInterfaceLocalMulticast /
// IsLinkLocalMulticast do.
let scope = seg0 & 0x000f;
return scope == 1 || scope == 2;
}
false
}
}
}
/// Whether `ip` is in a private range: IPv4 RFC 1918, IPv6 fc00::/7 unique-local
/// (matching Go's `net.IP.IsPrivate`).
fn is_private(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => v4.is_private(),
IpAddr::V6(v6) => (v6.segments()[0] & 0xfe00) == 0xfc00,
}
}
/// Whether `ip` is in the RFC 6598 carrier-grade NAT range (100.64.0.0/10).
/// The IPv4 private check does not cover CGNAT, so check it explicitly.
fn is_cgnat(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => {
let o = v4.octets();
o[0] == 100 && (o[1] & 0xc0) == 0x40 // second octet 64..=127
}
IpAddr::V6(_) => false,
}
}
/// Returns the IPv4 address carried by an IPv6 transition address -- NAT64
/// 64:ff9b::/96 (RFC 6052), 6to4 2002::/16 (RFC 3056), Teredo 2001:0000::/32
/// (RFC 4380), and the deprecated IPv4-compatible ::/96 (RFC 4291) -- or None
/// when it is not one of those. IPv4-mapped ::ffff:0:0/96 is excluded; it is
/// already normalized via `to_ipv4_mapped`.
fn embedded_transition_ipv4(v6: Ipv6Addr) -> Option<Ipv4Addr> {
let o = v6.octets();
if o[0] == 0x00
&& o[1] == 0x64
&& o[2] == 0xff
&& o[3] == 0x9b
&& o[4..12].iter().all(|&b| b == 0)
{
return Some(Ipv4Addr::new(o[12], o[13], o[14], o[15]));
}
if o[0] == 0x20 && o[1] == 0x02 {
return Some(Ipv4Addr::new(o[2], o[3], o[4], o[5]));
}
if o[0] == 0x20 && o[1] == 0x01 && o[2] == 0x00 && o[3] == 0x00 {
// Teredo obfuscates the client IPv4 as its ones' complement.
return Some(Ipv4Addr::new(
o[12] ^ 0xff,
o[13] ^ 0xff,
o[14] ^ 0xff,
o[15] ^ 0xff,
));
}
if o[..12].iter().all(|&b| b == 0) {
// IPv4-compatible ::a.b.c.d; :: and ::1 are already handled by the
// unspecified / loopback checks before extraction runs.
return Some(Ipv4Addr::new(o[12], o[13], o[14], o[15]));
}
None
}
/// Returns an error if `ip` is not safe to dial from a server that can reach
/// cluster-internal hosts. Mirrors Go's `checkBlockedIP`.
pub fn check_blocked_ip(endpoint: &str, ip: IpAddr) -> Result<(), String> {
check_blocked_ip_policy(endpoint, ip, false)
}
/// Like [`check_blocked_ip`], but `allow_private` keeps RFC 1918 / CGNAT
/// reachable for callers whose target legitimately sits on an internal network
/// (peer volume servers), while still blocking loopback, link-local (IMDS) and
/// unspecified. Mirrors Go's `checkBlockedIPPolicy`.
pub fn check_blocked_ip_policy(
endpoint: &str,
ip: IpAddr,
allow_private: bool,
) -> Result<(), String> {
// Normalize IPv4-mapped IPv6 (`::ffff:a.b.c.d`) to its IPv4 form so the
// IPv4 deny rules apply. The OS routes these to the embedded IPv4 address,
// so without this `::ffff:127.0.0.1` / `::ffff:169.254.169.254` would slip
// past the IPv6-only checks. Mirrors Go's reliance on net.IP.To4().
let ip = match ip {
IpAddr::V6(v6) => match v6.to_ipv4_mapped() {
Some(v4) => IpAddr::V4(v4),
None => IpAddr::V6(v6),
},
other => other,
};
if ip == IMDS_IPV4 {
return Err(format!(
"remote endpoint {:?} targets instance metadata service {}",
endpoint, ip
));
}
if ip.is_loopback() {
return Err(format!(
"remote endpoint {:?} resolves to loopback address {}",
endpoint, ip
));
}
if ip.is_unspecified() {
return Err(format!(
"remote endpoint {:?} resolves to unspecified address {}",
endpoint, ip
));
}
if is_link_local(ip) {
return Err(format!(
"remote endpoint {:?} resolves to link-local address {}",
endpoint, ip
));
}
if !allow_private {
if is_private(ip) {
return Err(format!(
"remote endpoint {:?} resolves to private address {}",
endpoint, ip
));
}
if is_cgnat(ip) {
return Err(format!(
"remote endpoint {:?} resolves to CGNAT address {}",
endpoint, ip
));
}
}
// IPv6 transition addresses embed an IPv4 destination that routes to the
// same host wherever the matching relay exists (common in IPv6-only cloud).
// to_ipv4_mapped above only covers ::ffff: mapped addresses, so pull the
// embedded IPv4 out of the other forms and re-check it against the rules.
if let IpAddr::V6(v6) = ip
&& let Some(v4) = embedded_transition_ipv4(v6)
{
return check_blocked_ip_policy(endpoint, IpAddr::V4(v4), allow_private);
}
Ok(())
}
/// Outcome of the synchronous, DNS-free portion of endpoint validation.
#[derive(Debug)]
enum HostCheck {
/// Host was a literal IP and already passed [`check_blocked_ip`].
Validated,
/// Host is a name that must be resolved and each address checked.
NeedsResolution(String),
}
/// Validate the scheme and host without touching DNS. Pure and unit-testable.
fn precheck_endpoint(endpoint: &str) -> Result<HostCheck, String> {
let trimmed = endpoint.trim();
if trimmed.is_empty() {
return Err("remote endpoint is empty".to_string());
}
let Some(scheme_end) = trimmed.find("://") else {
return Err(format!(
"remote endpoint {:?} must use http or https",
endpoint
));
};
let scheme = trimmed[..scheme_end].to_ascii_lowercase();
if scheme != "http" && scheme != "https" {
return Err(format!(
"remote endpoint {:?} must use http or https, got {:?}",
endpoint,
&trimmed[..scheme_end]
));
}
// Authority is everything up to the first '/', '?', or '#'.
let after = &trimmed[scheme_end + 3..];
let authority_end = after.find(['/', '?', '#']).unwrap_or(after.len());
let authority = &after[..authority_end];
// Strip optional userinfo ("user:pass@").
let host_port = match authority.rfind('@') {
Some(at) => &authority[at + 1..],
None => authority,
};
// Extract the host, handling bracketed IPv6 literals and host:port.
let host = if let Some(rest) = host_port.strip_prefix('[') {
match rest.find(']') {
Some(end) => &rest[..end],
None => {
return Err(format!(
"remote endpoint {:?} has a malformed IPv6 host",
endpoint
));
}
}
} else {
match host_port.find(':') {
Some(colon) => &host_port[..colon],
None => host_port,
}
};
if host.is_empty() {
return Err(format!("remote endpoint {:?} has no host", endpoint));
}
if is_blocked_imds_host(&host.to_ascii_lowercase()) {
return Err(format!(
"remote endpoint {:?} targets instance metadata service",
endpoint
));
}
if let Ok(ip) = host.parse::<IpAddr>() {
check_blocked_ip(endpoint, ip)?;
return Ok(HostCheck::Validated);
}
Ok(HostCheck::NeedsResolution(host.to_string()))
}
/// Resolve `host` to its addresses with a short timeout, matching Go's 2s
/// resolver deadline.
async fn resolve_host(host: &str) -> Result<Vec<IpAddr>, String> {
let lookup = tokio::net::lookup_host((host, 0u16));
let addrs = tokio::time::timeout(std::time::Duration::from_secs(2), lookup)
.await
.map_err(|_| format!("resolve remote endpoint host {:?}: timed out", host))?
.map_err(|e| format!("resolve remote endpoint host {:?}: {}", host, e))?;
Ok(addrs.map(|sock| sock.ip()).collect())
}
/// Returns an error if `endpoint` is not safe to dial from a server with access
/// to cluster-internal hosts. Rejects empty / non-http(s) endpoints,
/// loopback / unspecified / link-local / RFC 1918 / CGNAT addresses, and
/// well-known IMDS hostnames. Hostnames are resolved and every returned address
/// is checked. Mirrors Go's `validateRemoteEndpoint`.
pub async fn validate_remote_endpoint(endpoint: &str) -> Result<(), String> {
match precheck_endpoint(endpoint)? {
HostCheck::Validated => Ok(()),
HostCheck::NeedsResolution(host) => {
let addrs = resolve_host(&host).await?;
if addrs.is_empty() {
return Err(format!(
"resolve remote endpoint host {:?}: no addresses",
host
));
}
for ip in addrs {
check_blocked_ip(endpoint, ip)?;
}
Ok(())
}
}
}
/// Returns an error if `target` could redirect a replica upload away from a peer
/// volume server. The target must be a bare `host:port` -- a scheme, userinfo,
/// path, query or fragment can smuggle a different destination into the
/// formatted upload URL -- whose host is not loopback, link-local (IMDS) or
/// unspecified. Cluster peers legitimately sit on private networks, so RFC 1918
/// / CGNAT are allowed. Mirrors Go's `validateReplicaTarget`.
pub async fn validate_replica_target(target: &str) -> Result<(), String> {
let trimmed = target.trim();
if trimmed.is_empty() {
return Err("replica target is empty".to_string());
}
if trimmed.contains("://") || trimmed.contains(['/', '?', '#', '@', '\\']) {
return Err(format!(
"replica target {:?} must be a bare host:port",
target
));
}
// Require an explicit host:port, handling `[IPv6]:port`. A bracketless IPv6
// literal (which carries its own colons) is rejected; peers are addressed as
// `[ipv6]:port`, matching Go's net.SplitHostPort.
let host = if let Some(rest) = trimmed.strip_prefix('[') {
match rest.split_once(']') {
Some((h, port)) if port.starts_with(':') && port.len() > 1 => h,
_ => {
return Err(format!(
"replica target {:?} must be a bare host:port",
target
));
}
}
} else {
match trimmed.rsplit_once(':') {
Some((h, port)) if !port.is_empty() && !h.contains(':') => h,
_ => {
return Err(format!(
"replica target {:?} must be a bare host:port",
target
));
}
}
};
if host.is_empty() {
return Err(format!("replica target {:?} has no host", target));
}
if is_blocked_imds_host(&host.to_ascii_lowercase()) {
return Err(format!(
"replica target {:?} targets instance metadata service",
target
));
}
if let Ok(ip) = host.parse::<IpAddr>() {
return check_blocked_ip_policy(target, ip, true);
}
let addrs = resolve_host(host).await?;
if addrs.is_empty() {
return Err(format!(
"resolve replica target host {:?}: no addresses",
host
));
}
for ip in addrs {
check_blocked_ip_policy(target, ip, true)?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn ip(s: &str) -> IpAddr {
s.parse().unwrap()
}
#[test]
fn rejects_blocked_literal_endpoints() {
let cases = [
("http://127.0.0.1:8080", "loopback"),
("http://[::1]:8080", "loopback"),
("http://169.254.169.254/", "metadata"),
("http://0.0.0.0/", "unspecified"),
("http://[fe80::1]/", "link-local"),
("http://10.0.0.1/", "private"),
("http://172.16.5.5/", "private"),
("http://192.168.0.1/", "private"),
("http://100.64.0.1/", "CGNAT"),
];
for (endpoint, want) in cases {
let err = precheck_endpoint(endpoint).expect_err(endpoint);
assert!(err.contains(want), "{endpoint}: {err:?} missing {want:?}");
}
}
#[test]
fn rejects_empty_and_bad_scheme() {
assert!(precheck_endpoint("").unwrap_err().contains("empty"));
assert!(
precheck_endpoint("ftp://example.com/")
.unwrap_err()
.contains("http or https")
);
assert!(
precheck_endpoint("example.com/")
.unwrap_err()
.contains("http or https")
);
}
#[test]
fn rejects_imds_hostnames() {
assert!(
precheck_endpoint("http://metadata.google.internal/")
.unwrap_err()
.contains("metadata service")
);
assert!(
precheck_endpoint("http://metadata/")
.unwrap_err()
.contains("metadata service")
);
}
#[test]
fn accepts_public_literal_and_defers_hostname() {
assert!(matches!(
precheck_endpoint("https://52.216.10.10/"),
Ok(HostCheck::Validated)
));
match precheck_endpoint("https://s3.us-east-1.amazonaws.com/") {
Ok(HostCheck::NeedsResolution(host)) => {
assert_eq!(host, "s3.us-east-1.amazonaws.com")
}
other => panic!("expected resolution, got {other:?}"),
}
}
#[test]
fn check_blocked_ip_matches_resolved_categories() {
// Mirror Go's "host resolves to X" cases at the address level.
assert!(
check_blocked_ip("e", ip("127.0.0.1"))
.unwrap_err()
.contains("loopback")
);
assert!(
check_blocked_ip("e", ip("169.254.10.20"))
.unwrap_err()
.contains("link-local")
);
assert!(
check_blocked_ip("e", ip("10.1.2.3"))
.unwrap_err()
.contains("private")
);
assert!(
check_blocked_ip("e", ip("172.20.0.5"))
.unwrap_err()
.contains("private")
);
assert!(
check_blocked_ip("e", ip("192.168.1.1"))
.unwrap_err()
.contains("private")
);
assert!(
check_blocked_ip("e", ip("100.64.0.42"))
.unwrap_err()
.contains("CGNAT")
);
assert!(
check_blocked_ip("e", ip("fc00::1"))
.unwrap_err()
.contains("private")
);
assert!(check_blocked_ip("e", ip("52.216.10.10")).is_ok());
assert!(check_blocked_ip("e", ip("2606:4700:4700::1111")).is_ok());
}
#[test]
fn rejects_ipv6_transition_addresses() {
// Every transition form encoding an internal IPv4 must be blocked.
assert!(
check_blocked_ip("e", ip("64:ff9b::a9fe:a9fe")) // NAT64 -> IMDS
.unwrap_err()
.contains("metadata")
);
assert!(
check_blocked_ip("e", ip("64:ff9b::7f00:1")) // NAT64 -> loopback
.unwrap_err()
.contains("loopback")
);
assert!(
check_blocked_ip("e", ip("2002:a00:1::")) // 6to4 -> 10.0.0.1
.unwrap_err()
.contains("private")
);
assert!(
check_blocked_ip("e", ip("2001:0:4136:e378:8000:63bf:80ff:fffe")) // Teredo -> loopback
.unwrap_err()
.contains("loopback")
);
assert!(
check_blocked_ip("e", ip("::7f00:1")) // IPv4-compatible -> loopback
.unwrap_err()
.contains("loopback")
);
// A NAT64 address outside the 64:ff9b::/96 well-known prefix is not
// decoded (its embedded IPv4 lives elsewhere), so it is left as-is.
assert!(check_blocked_ip("e", ip("64:ff9b:1::a9fe:a9fe")).is_ok());
// Every transition form embedding a public IPv4 (8.8.8.8) still passes:
// NAT64, 6to4, Teredo, IPv4-compatible.
assert!(check_blocked_ip("e", ip("64:ff9b::808:808")).is_ok());
assert!(check_blocked_ip("e", ip("2002:808:808::")).is_ok());
assert!(check_blocked_ip("e", ip("2001::f7f7:f7f7")).is_ok());
assert!(check_blocked_ip("e", ip("::808:808")).is_ok());
// Bracketed transition literal via the full endpoint path.
assert!(
precheck_endpoint("http://[64:ff9b::a9fe:a9fe]/")
.unwrap_err()
.contains("metadata")
);
}
#[test]
fn rejects_ipv4_mapped_ipv6() {
// IPv4-mapped IPv6 must be unmapped so the IPv4 rules catch it.
assert!(
check_blocked_ip("e", ip("::ffff:127.0.0.1"))
.unwrap_err()
.contains("loopback")
);
assert!(
check_blocked_ip("e", ip("::ffff:169.254.169.254"))
.unwrap_err()
.contains("metadata")
);
assert!(
check_blocked_ip("e", ip("::ffff:10.0.0.1"))
.unwrap_err()
.contains("private")
);
// A mapped public address still passes, and genuine IPv6 loopback is
// still caught by the V6 path.
assert!(check_blocked_ip("e", ip("::ffff:52.216.10.10")).is_ok());
assert!(
check_blocked_ip("e", ip("::1"))
.unwrap_err()
.contains("loopback")
);
// Bracketed mapped literal via the full endpoint path.
assert!(
precheck_endpoint("http://[::ffff:127.0.0.1]/")
.unwrap_err()
.contains("loopback")
);
}
#[test]
fn check_blocked_ip_policy_allows_private_peers() {
// The replica leg targets peer volume servers, which may be private.
assert!(check_blocked_ip_policy("e", ip("10.0.0.7"), true).is_ok());
assert!(check_blocked_ip_policy("e", ip("192.168.1.5"), true).is_ok());
assert!(check_blocked_ip_policy("e", ip("100.64.0.42"), true).is_ok());
// Loopback / IMDS / unspecified stay blocked even when private is allowed.
assert!(
check_blocked_ip_policy("e", ip("127.0.0.1"), true)
.unwrap_err()
.contains("loopback")
);
assert!(
check_blocked_ip_policy("e", ip("169.254.169.254"), true)
.unwrap_err()
.contains("metadata")
);
assert!(
check_blocked_ip_policy("e", ip("0.0.0.0"), true)
.unwrap_err()
.contains("unspecified")
);
}
#[tokio::test]
async fn validate_replica_target_rejects_and_allows() {
// A path plus a trailing ?a= would otherwise swallow ?type=replicate.
assert!(
validate_replica_target("127.0.0.1:7000/status/x/?a=")
.await
.unwrap_err()
.contains("bare host:port")
);
assert!(
validate_replica_target("http://10.0.0.7:8080")
.await
.unwrap_err()
.contains("bare host:port")
);
assert!(
validate_replica_target("user@10.0.0.7:8080")
.await
.unwrap_err()
.contains("bare host:port")
);
assert!(
validate_replica_target("10.0.0.7")
.await
.unwrap_err()
.contains("bare host:port")
);
assert!(
validate_replica_target("peer.example.com")
.await
.unwrap_err()
.contains("bare host:port")
);
assert!(
validate_replica_target("127.0.0.1:8080")
.await
.unwrap_err()
.contains("loopback")
);
assert!(
validate_replica_target("[::1]:8080")
.await
.unwrap_err()
.contains("loopback")
);
assert!(
validate_replica_target("169.254.169.254:80")
.await
.unwrap_err()
.contains("metadata")
);
assert!(
validate_replica_target("metadata:80")
.await
.unwrap_err()
.contains("metadata")
);
assert!(
validate_replica_target("")
.await
.unwrap_err()
.contains("empty")
);
// Legitimate peer volume servers on private networks pass.
assert!(validate_replica_target("10.0.0.7:8080").await.is_ok());
assert!(validate_replica_target("192.168.1.5:8080").await.is_ok());
assert!(validate_replica_target("[fd00::1]:8080").await.is_ok());
}
}

View File

@@ -0,0 +1,246 @@
//! Remote storage backends for tiered storage support.
//!
//! Provides a trait-based abstraction over cloud storage providers (S3, GCS, Azure, etc.)
//! and a registry to create clients from protobuf RemoteConf messages.
pub mod endpoint_guard;
pub mod s3;
pub mod s3_tier;
pub use endpoint_guard::{validate_remote_endpoint, validate_replica_target};
use crate::pb::remote_pb::{RemoteConf, RemoteStorageLocation};
/// Error type for remote storage operations.
#[derive(Debug, thiserror::Error)]
pub enum RemoteStorageError {
#[error("remote storage type {0} not found")]
TypeNotFound(String),
#[error("remote object not found: {0}")]
ObjectNotFound(String),
#[error("remote storage error: {0}")]
Other(String),
#[error("io error: {0}")]
Io(#[from] std::io::Error),
}
/// Metadata about a remote file entry.
#[derive(Debug, Clone)]
pub struct RemoteEntry {
pub size: i64,
pub last_modified_at: i64, // Unix seconds
pub e_tag: String,
pub storage_name: String,
}
/// Trait for remote storage clients. Matches Go's RemoteStorageClient interface.
#[async_trait::async_trait]
pub trait RemoteStorageClient: Send + Sync {
/// Read (part of) a file from remote storage.
async fn read_file(
&self,
loc: &RemoteStorageLocation,
offset: i64,
size: i64,
) -> Result<Vec<u8>, RemoteStorageError>;
/// Write a file to remote storage.
async fn write_file(
&self,
loc: &RemoteStorageLocation,
data: &[u8],
) -> Result<RemoteEntry, RemoteStorageError>;
/// Get metadata for a file in remote storage.
async fn stat_file(
&self,
loc: &RemoteStorageLocation,
) -> Result<RemoteEntry, RemoteStorageError>;
/// Delete a file from remote storage.
async fn delete_file(&self, loc: &RemoteStorageLocation) -> Result<(), RemoteStorageError>;
/// List all buckets.
async fn list_buckets(&self) -> Result<Vec<String>, RemoteStorageError>;
/// The RemoteConf used to create this client.
fn remote_conf(&self) -> &RemoteConf;
}
/// Create a new remote storage client from a RemoteConf.
pub fn make_remote_storage_client(
conf: &RemoteConf,
) -> Result<Box<dyn RemoteStorageClient>, RemoteStorageError> {
match conf.r#type.as_str() {
// All S3-compatible backends use the same client with different credentials
"s3" | "wasabi" | "backblaze" | "aliyun" | "tencent" | "baidu" | "filebase" | "storj"
| "contabo" => {
let (access_key, secret_key, endpoint, region) = extract_s3_credentials(conf);
Ok(Box::new(s3::S3RemoteStorageClient::new(
conf.clone(),
&access_key,
&secret_key,
&region,
&endpoint,
conf.s3_force_path_style,
)))
}
other => Err(RemoteStorageError::TypeNotFound(other.to_string())),
}
}
/// Endpoint URL that the volume server would dial directly for `conf`, or
/// `None` for non-S3-compatible backends. Every type handled here routes
/// through [`s3::S3RemoteStorageClient`] with a caller-supplied endpoint, so
/// the SSRF guard must validate each one. Keep this match in sync with
/// [`make_remote_storage_client`].
pub fn s3_compatible_endpoint(conf: &RemoteConf) -> Option<&str> {
match conf.r#type.as_str() {
"s3" => Some(&conf.s3_endpoint),
"wasabi" => Some(&conf.wasabi_endpoint),
"backblaze" => Some(&conf.backblaze_endpoint),
"aliyun" => Some(&conf.aliyun_endpoint),
"tencent" => Some(&conf.tencent_endpoint),
"baidu" => Some(&conf.baidu_endpoint),
"filebase" => Some(&conf.filebase_endpoint),
"storj" => Some(&conf.storj_endpoint),
"contabo" => Some(&conf.contabo_endpoint),
_ => None,
}
}
/// Extract S3-compatible credentials from a RemoteConf based on its type.
fn extract_s3_credentials(conf: &RemoteConf) -> (String, String, String, String) {
match conf.r#type.as_str() {
"s3" => (
conf.s3_access_key.clone(),
conf.s3_secret_key.clone(),
conf.s3_endpoint.clone(),
if conf.s3_region.is_empty() {
"us-east-1".to_string()
} else {
conf.s3_region.clone()
},
),
"wasabi" => (
conf.wasabi_access_key.clone(),
conf.wasabi_secret_key.clone(),
conf.wasabi_endpoint.clone(),
conf.wasabi_region.clone(),
),
"backblaze" => (
conf.backblaze_key_id.clone(),
conf.backblaze_application_key.clone(),
conf.backblaze_endpoint.clone(),
conf.backblaze_region.clone(),
),
"aliyun" => (
conf.aliyun_access_key.clone(),
conf.aliyun_secret_key.clone(),
conf.aliyun_endpoint.clone(),
conf.aliyun_region.clone(),
),
"tencent" => (
conf.tencent_secret_id.clone(),
conf.tencent_secret_key.clone(),
conf.tencent_endpoint.clone(),
String::new(),
),
"baidu" => (
conf.baidu_access_key.clone(),
conf.baidu_secret_key.clone(),
conf.baidu_endpoint.clone(),
conf.baidu_region.clone(),
),
"filebase" => (
conf.filebase_access_key.clone(),
conf.filebase_secret_key.clone(),
conf.filebase_endpoint.clone(),
String::new(),
),
"storj" => (
conf.storj_access_key.clone(),
conf.storj_secret_key.clone(),
conf.storj_endpoint.clone(),
String::new(),
),
"contabo" => (
conf.contabo_access_key.clone(),
conf.contabo_secret_key.clone(),
conf.contabo_endpoint.clone(),
conf.contabo_region.clone(),
),
_ => (
conf.s3_access_key.clone(),
conf.s3_secret_key.clone(),
conf.s3_endpoint.clone(),
conf.s3_region.clone(),
),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn s3_compatible_endpoint_covers_all_s3_backends() {
let s3 = RemoteConf {
r#type: "s3".to_string(),
s3_endpoint: "http://s3.internal".to_string(),
..Default::default()
};
assert_eq!(s3_compatible_endpoint(&s3), Some("http://s3.internal"));
// A non-"s3" S3-compatible type still surfaces its own endpoint, so the
// SSRF guard cannot be bypassed by picking a different alias.
let wasabi = RemoteConf {
r#type: "wasabi".to_string(),
wasabi_endpoint: "http://wasabi.internal".to_string(),
..Default::default()
};
assert_eq!(
s3_compatible_endpoint(&wasabi),
Some("http://wasabi.internal")
);
// Non-S3 backends do not dial a caller-supplied URL directly.
let gcs = RemoteConf {
r#type: "gcs".to_string(),
..Default::default()
};
assert_eq!(s3_compatible_endpoint(&gcs), None);
}
#[test]
fn azure_endpoint_has_no_ssrf_path() {
// The Go volume server guards the caller-supplied azure endpoint against
// SSRF. This server has no azure backend, so there is nothing to dial:
// azure is not S3-compatible (the endpoint guard does not apply) and
// make_remote_storage_client rejects the type before building a client.
let azure = RemoteConf {
r#type: "azure".to_string(),
azure_endpoint: "https://169.254.169.254/".to_string(),
..Default::default()
};
assert_eq!(s3_compatible_endpoint(&azure), None);
assert!(make_remote_storage_client(&azure).is_err());
}
#[test]
fn gcs_credentials_have_no_ssrf_path() {
// The Go volume server accepts only static-key gcs credentials and puts
// their token endpoint behind the SSRF guard, because the SDK dials
// whatever url, file or executable the credentials name. This server has
// no gcs backend, so make_remote_storage_client rejects the type before
// any credentials are parsed. Anyone adding one must carry both guards
// over with it.
let gcs = RemoteConf {
r#type: "gcs".to_string(),
gcs_google_application_credentials: r#"{"type":"external_account","credential_source":{"url":"http://169.254.169.254/"}}"#.to_string(),
..Default::default()
};
assert_eq!(s3_compatible_endpoint(&gcs), None);
assert!(make_remote_storage_client(&gcs).is_err());
}
}

View File

@@ -0,0 +1,186 @@
//! S3-compatible remote storage client.
//!
//! Works with AWS S3, MinIO, SeaweedFS S3, and all S3-compatible providers.
use aws_sdk_s3::Client;
use aws_sdk_s3::config::{BehaviorVersion, Credentials, Region};
use aws_sdk_s3::primitives::ByteStream;
use super::{RemoteEntry, RemoteStorageClient, RemoteStorageError};
use crate::pb::remote_pb::{RemoteConf, RemoteStorageLocation};
/// S3-compatible remote storage client.
pub struct S3RemoteStorageClient {
client: Client,
conf: RemoteConf,
}
impl S3RemoteStorageClient {
/// Create a new S3 client from credentials and endpoint configuration.
pub fn new(
conf: RemoteConf,
access_key: &str,
secret_key: &str,
region: &str,
endpoint: &str,
force_path_style: bool,
) -> Self {
let region = if region.is_empty() {
"us-east-1"
} else {
region
};
let credentials = Credentials::new(
access_key,
secret_key,
None, // session token
None, // expiry
"seaweedfs-volume",
);
let mut s3_config = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(region.to_string()))
.credentials_provider(credentials)
.force_path_style(force_path_style);
if !endpoint.is_empty() {
s3_config = s3_config.endpoint_url(endpoint);
}
let client = Client::from_conf(s3_config.build());
S3RemoteStorageClient { client, conf }
}
}
#[async_trait::async_trait]
impl RemoteStorageClient for S3RemoteStorageClient {
async fn read_file(
&self,
loc: &RemoteStorageLocation,
offset: i64,
size: i64,
) -> Result<Vec<u8>, RemoteStorageError> {
let key = loc.path.trim_start_matches('/');
let mut req = self.client.get_object().bucket(&loc.bucket).key(key);
// Set byte range if specified
if size > 0 {
let end = offset + size - 1;
req = req.range(format!("bytes={}-{}", offset, end));
} else if offset > 0 {
req = req.range(format!("bytes={}-", offset));
}
let resp = req.send().await.map_err(|e| {
let msg = format!("{}", e);
if msg.contains("NoSuchKey") || msg.contains("404") {
RemoteStorageError::ObjectNotFound(format!("{}/{}", loc.bucket, key))
} else {
RemoteStorageError::Other(format!("s3 get object: {}", e))
}
})?;
let data = resp
.body
.collect()
.await
.map_err(|e| RemoteStorageError::Other(format!("s3 read body: {}", e)))?;
Ok(data.into_bytes().to_vec())
}
async fn write_file(
&self,
loc: &RemoteStorageLocation,
data: &[u8],
) -> Result<RemoteEntry, RemoteStorageError> {
let key = loc.path.trim_start_matches('/');
let resp = self
.client
.put_object()
.bucket(&loc.bucket)
.key(key)
.body(ByteStream::from(data.to_vec()))
.send()
.await
.map_err(|e| RemoteStorageError::Other(format!("s3 put object: {}", e)))?;
Ok(RemoteEntry {
size: data.len() as i64,
last_modified_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64,
e_tag: resp.e_tag().unwrap_or_default().to_string(),
storage_name: loc.name.clone(),
})
}
async fn stat_file(
&self,
loc: &RemoteStorageLocation,
) -> Result<RemoteEntry, RemoteStorageError> {
let key = loc.path.trim_start_matches('/');
let resp = self
.client
.head_object()
.bucket(&loc.bucket)
.key(key)
.send()
.await
.map_err(|e| {
let msg = format!("{}", e);
if msg.contains("404") || msg.contains("NotFound") {
RemoteStorageError::ObjectNotFound(format!("{}/{}", loc.bucket, key))
} else {
RemoteStorageError::Other(format!("s3 head object: {}", e))
}
})?;
Ok(RemoteEntry {
size: resp.content_length().unwrap_or(0),
last_modified_at: resp.last_modified().map(|t| t.secs()).unwrap_or(0),
e_tag: resp.e_tag().unwrap_or_default().to_string(),
storage_name: loc.name.clone(),
})
}
async fn delete_file(&self, loc: &RemoteStorageLocation) -> Result<(), RemoteStorageError> {
let key = loc.path.trim_start_matches('/');
self.client
.delete_object()
.bucket(&loc.bucket)
.key(key)
.send()
.await
.map_err(|e| RemoteStorageError::Other(format!("s3 delete object: {}", e)))?;
Ok(())
}
async fn list_buckets(&self) -> Result<Vec<String>, RemoteStorageError> {
let resp = self
.client
.list_buckets()
.send()
.await
.map_err(|e| RemoteStorageError::Other(format!("s3 list buckets: {}", e)))?;
Ok(resp
.buckets()
.iter()
.filter_map(|b| b.name().map(String::from))
.collect())
}
fn remote_conf(&self) -> &RemoteConf {
&self.conf
}
}

View File

@@ -0,0 +1,572 @@
//! S3-compatible tiered storage backend for volume .dat file upload/download.
//!
//! Provides multipart upload and concurrent download with progress callbacks,
//! matching the Go SeaweedFS S3 backend behavior.
use std::collections::HashMap;
use std::future::Future;
use std::sync::{Arc, OnceLock, RwLock};
use aws_sdk_s3::Client;
use aws_sdk_s3::config::{BehaviorVersion, Credentials, Region};
use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart};
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
use tokio::sync::Semaphore;
/// Concurrency limit for multipart upload/download (matches Go's s3manager).
const CONCURRENCY: usize = 5;
/// Configuration for an S3 tier backend.
#[derive(Debug, Clone)]
pub struct S3TierConfig {
pub access_key: String,
pub secret_key: String,
pub region: String,
pub bucket: String,
pub endpoint: String,
pub storage_class: String,
pub force_path_style: bool,
}
/// S3 tier backend for uploading/downloading volume .dat files.
pub struct S3TierBackend {
client: Client,
pub bucket: String,
pub storage_class: String,
}
impl S3TierBackend {
/// Create a new S3 tier backend from configuration.
pub fn new(config: &S3TierConfig) -> Self {
let region = if config.region.is_empty() {
"us-east-1"
} else {
&config.region
};
let credentials = Credentials::new(
&config.access_key,
&config.secret_key,
None,
None,
"seaweedfs-volume-tier",
);
let mut s3_config = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(region.to_string()))
.credentials_provider(credentials)
.force_path_style(config.force_path_style);
if !config.endpoint.is_empty() {
s3_config = s3_config.endpoint_url(&config.endpoint);
}
let client = Client::from_conf(s3_config.build());
S3TierBackend {
client,
bucket: config.bucket.clone(),
storage_class: if config.storage_class.is_empty() {
"STANDARD_IA".to_string()
} else {
config.storage_class.clone()
},
}
}
/// Upload a local file to S3 using multipart upload with concurrent parts
/// and progress reporting.
///
/// Returns (s3_key, file_size) on success.
/// The progress callback receives (bytes_uploaded, percentage).
/// Uses 64MB part size and 5 concurrent uploads (matches Go s3manager).
/// `progress_fn` returning `Err` aborts the upload, mirroring Go's
/// `fn(progressed, percentage) error` in `s3_upload.go`, where the error
/// surfaces out of `ReadAt` and fails the transfer. It is what lets a
/// caller that has hung up stop the work it is no longer waiting for.
pub async fn upload_file<F>(
&self,
file_path: &str,
progress_fn: F,
) -> Result<(String, u64), String>
where
F: FnMut(i64, f32) -> Result<(), String> + Send + Sync + 'static,
{
let key = uuid::Uuid::new_v4().to_string();
let metadata = tokio::fs::metadata(file_path)
.await
.map_err(|e| format!("failed to stat file {}: {}", file_path, e))?;
let file_size = metadata.len();
// Calculate part size: start at 64MB, scale up for very large files (matches Go)
let mut part_size: u64 = 64 * 1024 * 1024;
while part_size * 1000 < file_size {
part_size *= 4;
}
// Initiate multipart upload
let create_resp = self
.client
.create_multipart_upload()
.bucket(&self.bucket)
.key(&key)
.storage_class(
self.storage_class
.parse()
.unwrap_or(aws_sdk_s3::types::StorageClass::StandardIa),
)
.send()
.await
.map_err(|e| format!("failed to create multipart upload: {}", e))?;
let upload_id = create_resp
.upload_id()
.ok_or_else(|| "no upload_id in multipart upload response".to_string())?
.to_string();
// Build list of (part_number, offset, size) for all parts
let mut parts_plan: Vec<(i32, u64, usize)> = Vec::new();
let mut offset: u64 = 0;
let mut part_number: i32 = 1;
while offset < file_size {
let remaining = file_size - offset;
let this_part_size = std::cmp::min(part_size, remaining) as usize;
parts_plan.push((part_number, offset, this_part_size));
offset += this_part_size as u64;
part_number += 1;
}
// Upload parts concurrently with a semaphore limiting to CONCURRENCY
let semaphore = Arc::new(Semaphore::new(CONCURRENCY));
let client = &self.client;
let bucket = &self.bucket;
let file_path_owned = file_path.to_string();
let progress = Arc::new(std::sync::Mutex::new((0u64, progress_fn)));
let mut handles = Vec::with_capacity(parts_plan.len());
for (pn, off, size) in parts_plan {
let sem = semaphore.clone();
let client = client.clone();
let bucket = bucket.clone();
let key = key.clone();
let upload_id = upload_id.clone();
let fp = file_path_owned.clone();
let progress = progress.clone();
handles.push(tokio::spawn(async move {
let _permit = sem
.acquire()
.await
.map_err(|e| format!("semaphore error: {}", e))?;
// Read this part's data from the file at the correct offset
let mut file = tokio::fs::File::open(&fp)
.await
.map_err(|e| format!("failed to open file {}: {}", fp, e))?;
file.seek(std::io::SeekFrom::Start(off))
.await
.map_err(|e| format!("failed to seek to offset {}: {}", off, e))?;
let mut buf = vec![0u8; size];
file.read_exact(&mut buf)
.await
.map_err(|e| format!("failed to read file at offset {}: {}", off, e))?;
let upload_part_resp = client
.upload_part()
.bucket(&bucket)
.key(&key)
.upload_id(&upload_id)
.part_number(pn)
.body(buf.into())
.send()
.await
.map_err(|e| {
format!("failed to upload part {} at offset {}: {}", pn, off, e)
})?;
let e_tag = upload_part_resp.e_tag().unwrap_or_default().to_string();
// Report progress. The lock is released before the result is
// propagated so an aborting callback cannot poison the mutex
// for the other parts still in flight.
let progress_result = {
let mut guard = progress.lock().unwrap();
guard.0 += size as u64;
let uploaded = guard.0;
let pct = if file_size > 0 {
(uploaded as f32 * 100.0) / file_size as f32
} else {
100.0
};
(guard.1)(uploaded as i64, pct)
};
progress_result?;
Ok::<_, String>(
CompletedPart::builder()
.e_tag(e_tag)
.part_number(pn)
.build(),
)
}));
}
let finish = async {
// Collect results, preserving part order
let mut completed_parts = Vec::with_capacity(handles.len());
for handle in handles {
let part = handle
.await
.map_err(|e| format!("upload task panicked: {}", e))??;
completed_parts.push(part);
}
// Complete multipart upload
let completed_upload = CompletedMultipartUpload::builder()
.set_parts(Some(completed_parts))
.build();
self.client
.complete_multipart_upload()
.bucket(&self.bucket)
.key(&key)
.upload_id(&upload_id)
.multipart_upload(completed_upload)
.send()
.await
.map_err(|e| format!("failed to complete multipart upload: {}", e))?;
Ok::<(), String>(())
}
.await;
if let Err(e) = finish {
// An abandoned multipart upload does not appear in an ordinary
// object listing but still accrues storage charges until a
// lifecycle rule reaps it. Now that a departing caller aborts the
// transfer this is a routine path, not a rare one.
if let Err(abort_err) = self
.client
.abort_multipart_upload()
.bucket(&self.bucket)
.key(&key)
.upload_id(&upload_id)
.send()
.await
{
tracing::warn!(
"failed to abort multipart upload {} for key {}: {}",
upload_id,
key,
abort_err
);
}
return Err(e);
}
Ok((key, file_size))
}
/// Download a file from S3 to a local path with concurrent range requests
/// and progress reporting.
///
/// Returns the file size on success.
/// Uses 64MB part size and 5 concurrent downloads (matches Go s3manager).
/// `progress_fn` returning `Err` aborts the download, mirroring Go's
/// `fn(progressed, percentage) error` in `s3_download.go`.
pub async fn download_file<F>(
&self,
dest_path: &str,
key: &str,
progress_fn: F,
) -> Result<u64, String>
where
F: FnMut(i64, f32) -> Result<(), String> + Send + Sync + 'static,
{
// Get file size first
let head_resp = self
.client
.head_object()
.bucket(&self.bucket)
.key(key)
.send()
.await
.map_err(|e| format!("failed to head object {}: {}", key, e))?;
let file_size = head_resp.content_length().unwrap_or(0) as u64;
// Pre-allocate file to full size so concurrent WriteAt-style writes work
{
let file = tokio::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(dest_path)
.await
.map_err(|e| format!("failed to open dest file {}: {}", dest_path, e))?;
file.set_len(file_size)
.await
.map_err(|e| format!("failed to set file length: {}", e))?;
}
let part_size: u64 = 64 * 1024 * 1024;
// Build list of (offset, size) for all parts
let mut parts_plan: Vec<(u64, u64)> = Vec::new();
let mut offset: u64 = 0;
while offset < file_size {
let remaining = file_size - offset;
let this_part_size = std::cmp::min(part_size, remaining);
parts_plan.push((offset, this_part_size));
offset += this_part_size;
}
// Download parts concurrently with a semaphore limiting to CONCURRENCY
let semaphore = Arc::new(Semaphore::new(CONCURRENCY));
let client = &self.client;
let bucket = &self.bucket;
let dest_path_owned = dest_path.to_string();
let key_owned = key.to_string();
let progress = Arc::new(std::sync::Mutex::new((0u64, progress_fn)));
let mut handles = Vec::with_capacity(parts_plan.len());
for (off, size) in parts_plan {
let sem = semaphore.clone();
let client = client.clone();
let bucket = bucket.clone();
let key = key_owned.clone();
let dp = dest_path_owned.clone();
let progress = progress.clone();
handles.push(tokio::spawn(async move {
let _permit = sem
.acquire()
.await
.map_err(|e| format!("semaphore error: {}", e))?;
let end = off + size - 1;
let range = format!("bytes={}-{}", off, end);
let get_resp = client
.get_object()
.bucket(&bucket)
.key(&key)
.range(&range)
.send()
.await
.map_err(|e| format!("failed to get object {} range {}: {}", key, range, e))?;
let body = get_resp
.body
.collect()
.await
.map_err(|e| format!("failed to read body: {}", e))?;
let bytes = body.into_bytes();
// Write at the correct offset (like Go's WriteAt)
let mut file = tokio::fs::OpenOptions::new()
.write(true)
.open(&dp)
.await
.map_err(|e| format!("failed to open dest file {}: {}", dp, e))?;
file.seek(std::io::SeekFrom::Start(off))
.await
.map_err(|e| format!("failed to seek to offset {}: {}", off, e))?;
file.write_all(&bytes)
.await
.map_err(|e| format!("failed to write to {}: {}", dp, e))?;
// Report progress. The lock is released before the result is
// propagated so an aborting callback cannot poison the mutex
// for the other parts still in flight.
let progress_result = {
let mut guard = progress.lock().unwrap();
guard.0 += bytes.len() as u64;
let downloaded = guard.0;
let pct = if file_size > 0 {
(downloaded as f32 * 100.0) / file_size as f32
} else {
100.0
};
(guard.1)(downloaded as i64, pct)
};
progress_result?;
Ok::<_, String>(())
}));
}
// Wait for all download tasks
for handle in handles {
handle
.await
.map_err(|e| format!("download task panicked: {}", e))??;
}
// fsync the file so its content is durable before the caller trims the .vif
// and deletes the remote object (matches Go's DownloadFile f.Sync()).
let synced = tokio::fs::OpenOptions::new()
.write(true)
.open(dest_path)
.await
.map_err(|e| format!("failed to open {} for fsync: {}", dest_path, e))?;
synced
.sync_all()
.await
.map_err(|e| format!("failed to fsync {}: {}", dest_path, e))?;
Ok(file_size)
}
pub async fn read_range(&self, key: &str, offset: u64, size: usize) -> Result<Vec<u8>, String> {
let end = offset + (size as u64).saturating_sub(1);
let range = format!("bytes={}-{}", offset, end);
let resp = self
.client
.get_object()
.bucket(&self.bucket)
.key(key)
.range(&range)
.send()
.await
.map_err(|e| format!("failed to get object {} range {}: {}", key, range, e))?;
let body = resp
.body
.collect()
.await
.map_err(|e| format!("failed to read object {} body: {}", key, e))?;
Ok(body.into_bytes().to_vec())
}
/// Delete a file from S3.
pub async fn delete_file(&self, key: &str) -> Result<(), String> {
self.client
.delete_object()
.bucket(&self.bucket)
.key(key)
.send()
.await
.map_err(|e| format!("failed to delete object {}: {}", key, e))?;
Ok(())
}
pub fn delete_file_blocking(&self, key: &str) -> Result<(), String> {
let client = self.client.clone();
let bucket = self.bucket.clone();
let key = key.to_string();
block_on_tier_future(async move {
client
.delete_object()
.bucket(&bucket)
.key(&key)
.send()
.await
.map_err(|e| format!("failed to delete object {}: {}", key, e))?;
Ok(())
})
}
pub fn read_range_blocking(
&self,
key: &str,
offset: u64,
size: usize,
) -> Result<Vec<u8>, String> {
let client = self.client.clone();
let bucket = self.bucket.clone();
let key = key.to_string();
block_on_tier_future(async move {
let end = offset + (size as u64).saturating_sub(1);
let range = format!("bytes={}-{}", offset, end);
let resp = client
.get_object()
.bucket(&bucket)
.key(&key)
.range(&range)
.send()
.await
.map_err(|e| format!("failed to get object {} range {}: {}", key, range, e))?;
let body = resp
.body
.collect()
.await
.map_err(|e| format!("failed to read object {} body: {}", key, e))?;
Ok(body.into_bytes().to_vec())
})
}
}
/// Parse a backend name like "s3" or "s3.default" into (backend_type, backend_id).
/// Matches Go's `BackendNameToTypeId`.
pub fn backend_name_to_type_id(backend_name: &str) -> (String, String) {
let parts: Vec<&str> = backend_name.split('.').collect();
match parts.len() {
1 => (backend_name.to_string(), "default".to_string()),
2 => (parts[0].to_string(), parts[1].to_string()),
_ => (String::new(), String::new()),
}
}
/// A registry of configured S3 tier backends, keyed by backend name (e.g., "s3.default").
#[derive(Default)]
pub struct S3TierRegistry {
backends: HashMap<String, Arc<S3TierBackend>>,
}
impl S3TierRegistry {
pub fn new() -> Self {
Self {
backends: HashMap::new(),
}
}
/// Register a backend with the given name.
pub fn register(&mut self, name: String, backend: S3TierBackend) {
self.backends.insert(name, Arc::new(backend));
}
/// Look up a backend by name.
pub fn get(&self, name: &str) -> Option<Arc<S3TierBackend>> {
self.backends.get(name).cloned()
}
/// List all registered backend names.
pub fn names(&self) -> Vec<String> {
self.backends.keys().cloned().collect()
}
/// Remove a backend by name.
pub fn remove(&mut self, name: &str) {
self.backends.remove(name);
}
pub fn clear(&mut self) {
self.backends.clear();
}
}
static GLOBAL_S3_TIER_REGISTRY: OnceLock<RwLock<S3TierRegistry>> = OnceLock::new();
pub fn global_s3_tier_registry() -> &'static RwLock<S3TierRegistry> {
GLOBAL_S3_TIER_REGISTRY.get_or_init(|| RwLock::new(S3TierRegistry::new()))
}
fn block_on_tier_future<F, T>(future: F) -> Result<T, String>
where
F: Future<Output = Result<T, String>> + Send + 'static,
T: Send + 'static,
{
std::thread::spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| format!("failed to build tokio runtime: {}", e))?;
runtime.block_on(future)
})
.join()
.map_err(|_| "tier runtime thread panicked".to_string())?
}

View File

@@ -0,0 +1,547 @@
//! Security: JWT validation and IP whitelist checking.
//!
//! Matches Go's security/guard.go and security/jwt.go.
//! - Guard: combines whitelist IP checking with JWT token validation
//! - JWT: HS256 HMAC signing with file-id claims
pub mod tls;
use std::collections::HashSet;
use std::net::IpAddr;
use std::time::{SystemTime, UNIX_EPOCH};
use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, Validation, decode, encode};
use serde::{Deserialize, Serialize};
// ============================================================================
// JWT Claims
// ============================================================================
/// Claims for volume server file access tokens.
/// Matches Go's `SeaweedFileIdClaims`.
#[derive(Debug, Serialize, Deserialize)]
pub struct FileIdClaims {
/// File ID this token grants access to (e.g., "3,01637037d6").
#[serde(skip_serializing_if = "Option::is_none")]
pub fid: Option<String>,
/// Expiration time (Unix timestamp).
#[serde(skip_serializing_if = "Option::is_none")]
pub exp: Option<u64>,
/// Not before (Unix timestamp).
#[serde(skip_serializing_if = "Option::is_none")]
pub nbf: Option<u64>,
}
/// Signing key wrapper (empty = security disabled).
#[derive(Clone)]
pub struct SigningKey(pub Vec<u8>);
impl SigningKey {
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn from_string(s: &str) -> Self {
SigningKey(s.as_bytes().to_vec())
}
}
/// Generate a JWT token for file access.
pub fn gen_jwt(
signing_key: &SigningKey,
expires_after_sec: i64,
file_id: &str,
) -> Result<String, JwtError> {
if signing_key.is_empty() {
return Err(JwtError::NoSigningKey);
}
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let claims = FileIdClaims {
fid: Some(file_id.to_string()),
exp: if expires_after_sec > 0 {
Some(now + expires_after_sec as u64)
} else {
None
},
nbf: None,
};
let token = encode(
&Header::new(Algorithm::HS256),
&claims,
&EncodingKey::from_secret(&signing_key.0),
)?;
Ok(token)
}
/// Decode and validate a JWT token.
pub fn decode_jwt(signing_key: &SigningKey, token: &str) -> Result<FileIdClaims, JwtError> {
if signing_key.is_empty() {
return Err(JwtError::NoSigningKey);
}
let mut validation = Validation::new(Algorithm::HS256);
// Match Go behavior: tokens without exp are accepted (Go's jwt-go does not require exp)
// But if exp IS present, it must be valid (not expired).
validation.required_spec_claims.clear();
validation.validate_exp = true;
// Go's jwt-go/v5 validates nbf when present
validation.validate_nbf = true;
validation.leeway = 0;
let data = decode::<FileIdClaims>(
token,
&DecodingKey::from_secret(&signing_key.0),
&validation,
)?;
Ok(data.claims)
}
// ============================================================================
// Guard
// ============================================================================
/// Security guard: IP whitelist + JWT token validation.
pub struct Guard {
whitelist_ips: HashSet<String>,
whitelist_cidrs: Vec<(IpAddr, u8)>, // (network, prefix_len)
pub signing_key: SigningKey,
pub expires_after_sec: i64,
pub read_signing_key: SigningKey,
pub read_expires_after_sec: i64,
/// Combined flag: true when whitelist is non-empty OR signing key is present.
/// Matches Go's `isWriteActive = !isEmptyWhiteList || len(SigningKey) != 0`.
is_write_active: bool,
}
impl Guard {
pub fn new(
whitelist: &[String],
signing_key: SigningKey,
expires_after_sec: i64,
read_signing_key: SigningKey,
read_expires_after_sec: i64,
) -> Self {
let mut guard = Guard {
whitelist_ips: HashSet::new(),
whitelist_cidrs: Vec::new(),
signing_key,
expires_after_sec,
read_signing_key,
read_expires_after_sec,
is_write_active: false,
};
guard.update_whitelist(whitelist);
guard
}
/// Update the IP whitelist.
pub fn update_whitelist(&mut self, entries: &[String]) {
self.whitelist_ips.clear();
self.whitelist_cidrs.clear();
for entry in entries {
let entry = entry.trim();
if entry.is_empty() {
continue;
}
if entry.contains('/') {
// CIDR range
if let Some((ip, prefix)) = parse_cidr(entry) {
self.whitelist_cidrs.push((ip, prefix));
} else {
tracing::error!("Parse CIDR {} in whitelist failed", entry);
}
} else {
// Exact IP/hostname
self.whitelist_ips.insert(entry.to_string());
}
}
// Match Go: isWriteActive = !isEmptyWhiteList || len(SigningKey) != 0
let is_empty_whitelist = self.whitelist_ips.is_empty() && self.whitelist_cidrs.is_empty();
self.is_write_active = !is_empty_whitelist || !self.signing_key.is_empty();
}
/// Refresh the JWT signing keys and their expirations in place so operators
/// can rotate keys via SIGHUP without restarting the process. Mirrors Go's
/// `Guard::UpdateSigningKeys`: it swaps the fields and recomputes
/// `is_write_active`.
pub fn update_signing_keys(
&mut self,
signing_key: SigningKey,
expires_after_sec: i64,
read_signing_key: SigningKey,
read_expires_after_sec: i64,
) {
self.signing_key = signing_key;
self.expires_after_sec = expires_after_sec;
self.read_signing_key = read_signing_key;
self.read_expires_after_sec = read_expires_after_sec;
let is_empty_whitelist = self.whitelist_ips.is_empty() && self.whitelist_cidrs.is_empty();
self.is_write_active = !is_empty_whitelist || !self.signing_key.is_empty();
}
/// Whether any write-side control is configured: a non-empty whitelist or a
/// signing key. Mirrors Go's non-nil `guard` -- when this is false, every
/// caller is allowed and admin RPCs need not carry peer info.
pub fn is_write_active(&self) -> bool {
self.is_write_active
}
/// Check if a remote IP is in the whitelist.
/// Returns true if write security is inactive (no whitelist and no signing key),
/// if the whitelist is empty, or if the IP matches.
pub fn check_whitelist(&self, remote_addr: &str) -> bool {
if !self.is_write_active {
return true;
}
if self.whitelist_ips.is_empty() && self.whitelist_cidrs.is_empty() {
return true;
}
let host = extract_host(remote_addr);
// Check exact match
if self.whitelist_ips.contains(&host) {
return true;
}
// Check CIDR ranges
if let Ok(ip) = host.parse::<IpAddr>() {
for &(ref network, prefix_len) in &self.whitelist_cidrs {
if ip_in_cidr(&ip, network, prefix_len) {
return true;
}
}
}
false
}
/// Check if a read signing key is configured.
pub fn has_read_signing_key(&self) -> bool {
!self.read_signing_key.is_empty()
}
/// Validate a request's JWT token.
/// `is_write` determines which signing key to use.
/// Returns Ok(()) if valid, or if security is disabled.
pub fn check_jwt(&self, token: Option<&str>, is_write: bool) -> Result<(), JwtError> {
let key = if is_write {
&self.signing_key
} else {
&self.read_signing_key
};
if key.is_empty() {
return Ok(()); // Security disabled for this operation type
}
let token = token.ok_or(JwtError::MissingToken)?;
decode_jwt(key, token)?;
Ok(())
}
/// Check JWT and validate the file ID claim matches.
pub fn check_jwt_for_file(
&self,
token: Option<&str>,
expected_fid: &str,
is_write: bool,
) -> Result<(), JwtError> {
let key = if is_write {
&self.signing_key
} else {
&self.read_signing_key
};
if key.is_empty() {
return Ok(());
}
let token = token.ok_or(JwtError::MissingToken)?;
let claims = decode_jwt(key, token)?;
match claims.fid {
None => {
return Err(JwtError::MissingFileIdClaim);
}
Some(ref fid) if fid != expected_fid => {
return Err(JwtError::FileIdMismatch {
expected: expected_fid.to_string(),
got: fid.to_string(),
});
}
_ => {}
}
Ok(())
}
}
// ============================================================================
// Helpers
// ============================================================================
/// Extract host from "host:port" or "[::1]:port" format.
fn extract_host(addr: &str) -> String {
// Handle IPv6 with brackets
if addr.starts_with('[')
&& let Some(end) = addr.find(']')
{
return addr[1..end].to_string();
}
// Handle host:port
if let Some(pos) = addr.rfind(':') {
return addr[..pos].to_string();
}
addr.to_string()
}
/// Parse CIDR notation "192.168.1.0/24" into (IpAddr, prefix_len).
fn parse_cidr(cidr: &str) -> Option<(IpAddr, u8)> {
let parts: Vec<&str> = cidr.split('/').collect();
if parts.len() != 2 {
return None;
}
let ip: IpAddr = parts[0].parse().ok()?;
let prefix: u8 = parts[1].parse().ok()?;
Some((ip, prefix))
}
/// Check if an IP is within a CIDR range.
fn ip_in_cidr(ip: &IpAddr, network: &IpAddr, prefix_len: u8) -> bool {
match (ip, network) {
(IpAddr::V4(ip), IpAddr::V4(net)) => {
let ip_bits = u32::from(*ip);
let net_bits = u32::from(*net);
let mask = if prefix_len == 0 {
0
} else if prefix_len >= 32 {
u32::MAX
} else {
u32::MAX << (32 - prefix_len)
};
(ip_bits & mask) == (net_bits & mask)
}
(IpAddr::V6(ip), IpAddr::V6(net)) => {
let ip_bits = u128::from(*ip);
let net_bits = u128::from(*net);
let mask = if prefix_len == 0 {
0
} else if prefix_len >= 128 {
u128::MAX
} else {
u128::MAX << (128 - prefix_len)
};
(ip_bits & mask) == (net_bits & mask)
}
_ => false, // V4/V6 mismatch
}
}
// ============================================================================
// Errors
// ============================================================================
#[derive(Debug, thiserror::Error)]
pub enum JwtError {
#[error("no signing key configured")]
NoSigningKey,
#[error("missing JWT token")]
MissingToken,
#[error("JWT error: {0}")]
Jwt(#[from] jsonwebtoken::errors::Error),
#[error("JWT token missing required fid claim")]
MissingFileIdClaim,
#[error("file ID mismatch: expected {expected}, got {got}")]
FileIdMismatch { expected: String, got: String },
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_jwt_round_trip() {
let key = SigningKey::from_string("test-secret-key");
let token = gen_jwt(&key, 3600, "3,01637037d6").unwrap();
let claims = decode_jwt(&key, &token).unwrap();
assert_eq!(claims.fid, Some("3,01637037d6".to_string()));
}
#[test]
fn test_jwt_no_signing_key() {
let key = SigningKey(vec![]);
assert!(gen_jwt(&key, 3600, "1,abc").is_err());
}
#[test]
fn test_jwt_invalid_token() {
let key = SigningKey::from_string("secret");
let result = decode_jwt(&key, "invalid.token.here");
assert!(result.is_err());
}
#[test]
fn test_jwt_wrong_key() {
let key1 = SigningKey::from_string("secret1");
let key2 = SigningKey::from_string("secret2");
let token = gen_jwt(&key1, 3600, "1,abc").unwrap();
assert!(decode_jwt(&key2, &token).is_err());
}
#[test]
fn test_guard_empty_whitelist() {
let guard = Guard::new(&[], SigningKey(vec![]), 0, SigningKey(vec![]), 0);
assert!(guard.check_whitelist("192.168.1.1:8080"));
}
#[test]
fn test_guard_whitelist_exact() {
let guard = Guard::new(
&["192.168.1.1".to_string(), "10.0.0.1".to_string()],
SigningKey(vec![]),
0,
SigningKey(vec![]),
0,
);
assert!(guard.check_whitelist("192.168.1.1:8080"));
assert!(guard.check_whitelist("10.0.0.1:1234"));
assert!(!guard.check_whitelist("172.16.0.1:8080"));
}
#[test]
fn test_guard_whitelist_cidr() {
let guard = Guard::new(
&["10.0.0.0/8".to_string()],
SigningKey(vec![]),
0,
SigningKey(vec![]),
0,
);
assert!(guard.check_whitelist("10.1.2.3:8080"));
assert!(guard.check_whitelist("10.255.255.255:80"));
assert!(!guard.check_whitelist("11.0.0.1:80"));
}
#[test]
fn test_guard_check_jwt_disabled() {
let guard = Guard::new(&[], SigningKey(vec![]), 0, SigningKey(vec![]), 0);
// No signing key = security disabled
assert!(guard.check_jwt(None, true).is_ok());
assert!(guard.check_jwt(None, false).is_ok());
}
#[test]
fn test_guard_check_jwt_enabled() {
let key = SigningKey::from_string("write-secret");
let read_key = SigningKey::from_string("read-secret");
let guard = Guard::new(&[], key.clone(), 3600, read_key.clone(), 3600);
// Missing token
assert!(guard.check_jwt(None, true).is_err());
// Valid write token
let token = gen_jwt(&key, 3600, "1,abc").unwrap();
assert!(guard.check_jwt(Some(&token), true).is_ok());
// Write token for read should fail (different key)
assert!(guard.check_jwt(Some(&token), false).is_err());
// Valid read token
let read_token = gen_jwt(&read_key, 3600, "1,abc").unwrap();
assert!(guard.check_jwt(Some(&read_token), false).is_ok());
}
#[test]
fn test_guard_check_jwt_file_id() {
let key = SigningKey::from_string("secret");
let guard = Guard::new(&[], key.clone(), 3600, SigningKey(vec![]), 0);
let token = gen_jwt(&key, 3600, "3,01637037d6").unwrap();
// Correct file ID
assert!(
guard
.check_jwt_for_file(Some(&token), "3,01637037d6", true)
.is_ok()
);
// Wrong file ID
let err = guard.check_jwt_for_file(Some(&token), "4,deadbeef", true);
assert!(matches!(err, Err(JwtError::FileIdMismatch { .. })));
}
#[test]
fn test_update_signing_keys_rotates() {
let mut guard = Guard::new(
&[],
SigningKey::from_string("old-write"),
10,
SigningKey::from_string("old-read"),
60,
);
guard.update_signing_keys(
SigningKey::from_string("new-write"),
11,
SigningKey::from_string("new-read"),
61,
);
let token = gen_jwt(&guard.signing_key, 60, "3,01637037d6").unwrap();
// Validates with the rotated key, not the old one.
assert!(decode_jwt(&guard.signing_key, &token).is_ok());
assert!(decode_jwt(&SigningKey::from_string("old-write"), &token).is_err());
assert_eq!(guard.expires_after_sec, 11);
assert_eq!(guard.read_expires_after_sec, 61);
}
#[test]
fn test_update_signing_keys_toggles_write_active() {
let mut guard = Guard::new(&[], SigningKey(vec![]), 0, SigningKey(vec![]), 0);
assert!(guard.check_jwt(None, true).is_ok());
guard.update_signing_keys(SigningKey::from_string("write"), 10, SigningKey(vec![]), 0);
assert!(guard.is_write_active);
guard.update_signing_keys(SigningKey(vec![]), 0, SigningKey(vec![]), 0);
assert!(!guard.is_write_active);
}
#[test]
fn test_extract_host() {
assert_eq!(extract_host("192.168.1.1:8080"), "192.168.1.1");
assert_eq!(extract_host("[::1]:8080"), "::1");
assert_eq!(extract_host("localhost"), "localhost");
}
#[test]
fn test_ip_in_cidr() {
let net: IpAddr = "10.0.0.0".parse().unwrap();
let ip1: IpAddr = "10.1.2.3".parse().unwrap();
let ip2: IpAddr = "11.0.0.1".parse().unwrap();
assert!(ip_in_cidr(&ip1, &net, 8));
assert!(!ip_in_cidr(&ip2, &net, 8));
}
}

View File

@@ -0,0 +1,444 @@
use std::collections::HashSet;
use std::fmt;
use std::sync::Arc;
use rustls::client::danger::HandshakeSignatureValid;
use rustls::crypto::CryptoProvider;
use rustls::crypto::aws_lc_rs;
use rustls::pki_types::UnixTime;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use rustls::server::WebPkiClientVerifier;
use rustls::server::danger::{ClientCertVerified, ClientCertVerifier};
use rustls::{
CipherSuite, DigitallySignedStruct, DistinguishedName, RootCertStore, ServerConfig,
SignatureScheme, SupportedCipherSuite, SupportedProtocolVersion,
};
use x509_parser::prelude::{FromDer, X509Certificate};
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TlsPolicy {
pub min_version: String,
pub max_version: String,
pub cipher_suites: String,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct GrpcClientAuthPolicy {
pub allowed_common_names: Vec<String>,
pub allowed_wildcard_domain: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TlsPolicyError(String);
impl fmt::Display for TlsPolicyError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
impl std::error::Error for TlsPolicyError {}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
enum GoTlsVersion {
Ssl3,
Tls10,
Tls11,
Tls12,
Tls13,
}
#[derive(Debug)]
struct CommonNameVerifier {
inner: Arc<dyn ClientCertVerifier>,
allowed_common_names: HashSet<String>,
allowed_wildcard_domain: String,
}
impl ClientCertVerifier for CommonNameVerifier {
fn offer_client_auth(&self) -> bool {
self.inner.offer_client_auth()
}
fn client_auth_mandatory(&self) -> bool {
self.inner.client_auth_mandatory()
}
fn root_hint_subjects(&self) -> &[DistinguishedName] {
self.inner.root_hint_subjects()
}
fn verify_client_cert(
&self,
end_entity: &CertificateDer<'_>,
intermediates: &[CertificateDer<'_>],
now: UnixTime,
) -> Result<ClientCertVerified, rustls::Error> {
self.inner
.verify_client_cert(end_entity, intermediates, now)?;
let common_name = parse_common_name(end_entity).map_err(|e| {
rustls::Error::General(format!(
"parse client certificate common name failed: {}",
e
))
})?;
if common_name_is_allowed(
&common_name,
&self.allowed_common_names,
&self.allowed_wildcard_domain,
) {
return Ok(ClientCertVerified::assertion());
}
Err(rustls::Error::General(format!(
"Authenticate: invalid subject client common name: {}",
common_name
)))
}
fn verify_tls12_signature(
&self,
message: &[u8],
cert: &CertificateDer<'_>,
dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
self.inner.verify_tls12_signature(message, cert, dss)
}
fn verify_tls13_signature(
&self,
message: &[u8],
cert: &CertificateDer<'_>,
dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
self.inner.verify_tls13_signature(message, cert, dss)
}
fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
self.inner.supported_verify_schemes()
}
}
// aws-lc-rs and ring both get linked transitively, so rustls can't auto-select
// a provider and tonic's client TLS panics on first use. Pin the default to
// aws-lc-rs, matching the server config. Idempotent.
pub fn install_default_crypto_provider() {
let _ = aws_lc_rs::default_provider().install_default();
}
pub fn build_rustls_server_config(
cert_path: &str,
key_path: &str,
ca_path: &str,
policy: &TlsPolicy,
) -> Result<ServerConfig, TlsPolicyError> {
build_rustls_server_config_with_client_auth(cert_path, key_path, ca_path, policy, None)
}
pub fn build_rustls_server_config_with_grpc_client_auth(
cert_path: &str,
key_path: &str,
ca_path: &str,
policy: &TlsPolicy,
client_auth_policy: &GrpcClientAuthPolicy,
) -> Result<ServerConfig, TlsPolicyError> {
build_rustls_server_config_with_client_auth(
cert_path,
key_path,
ca_path,
policy,
Some(client_auth_policy),
)
}
fn build_rustls_server_config_with_client_auth(
cert_path: &str,
key_path: &str,
ca_path: &str,
policy: &TlsPolicy,
client_auth_policy: Option<&GrpcClientAuthPolicy>,
) -> Result<ServerConfig, TlsPolicyError> {
let cert_chain = read_cert_chain(cert_path)?;
let private_key = read_private_key(key_path)?;
let provider = build_crypto_provider(policy)?;
let versions = build_supported_versions(policy)?;
let builder = ServerConfig::builder_with_provider(provider.clone())
.with_protocol_versions(&versions)
.map_err(|e| TlsPolicyError(format!("invalid TLS version policy: {}", e)))?;
let builder = if ca_path.is_empty() {
builder.with_no_client_auth()
} else {
let roots = read_root_store(ca_path)?;
let verifier =
WebPkiClientVerifier::builder_with_provider(Arc::new(roots), provider.clone())
.build()
.map_err(|e| TlsPolicyError(format!("build client verifier failed: {}", e)))?;
let verifier: Arc<dyn ClientCertVerifier> = if let Some(client_auth_policy) =
client_auth_policy.filter(|policy| {
!policy.allowed_common_names.is_empty()
|| !policy.allowed_wildcard_domain.is_empty()
}) {
Arc::new(CommonNameVerifier {
inner: verifier,
allowed_common_names: client_auth_policy
.allowed_common_names
.iter()
.cloned()
.collect(),
allowed_wildcard_domain: client_auth_policy.allowed_wildcard_domain.clone(),
})
} else {
verifier
};
builder.with_client_cert_verifier(verifier)
};
builder
.with_single_cert(cert_chain, private_key)
.map_err(|e| TlsPolicyError(format!("build rustls server config failed: {}", e)))
}
fn read_cert_chain(cert_path: &str) -> Result<Vec<CertificateDer<'static>>, TlsPolicyError> {
let cert_pem = std::fs::read(cert_path).map_err(|e| {
TlsPolicyError(format!(
"Failed to read TLS cert file '{}': {}",
cert_path, e
))
})?;
rustls_pemfile::certs(&mut &cert_pem[..])
.collect::<Result<Vec<_>, _>>()
.map_err(|e| {
TlsPolicyError(format!(
"Failed to parse TLS cert PEM '{}': {}",
cert_path, e
))
})
}
fn read_private_key(key_path: &str) -> Result<PrivateKeyDer<'static>, TlsPolicyError> {
let key_pem = std::fs::read(key_path).map_err(|e| {
TlsPolicyError(format!("Failed to read TLS key file '{}': {}", key_path, e))
})?;
rustls_pemfile::private_key(&mut &key_pem[..])
.map_err(|e| TlsPolicyError(format!("Failed to parse TLS key PEM '{}': {}", key_path, e)))?
.ok_or_else(|| TlsPolicyError(format!("No private key found in '{}'", key_path)))
}
fn read_root_store(ca_path: &str) -> Result<RootCertStore, TlsPolicyError> {
let ca_pem = std::fs::read(ca_path)
.map_err(|e| TlsPolicyError(format!("Failed to read TLS CA file '{}': {}", ca_path, e)))?;
let ca_certs = rustls_pemfile::certs(&mut &ca_pem[..])
.collect::<Result<Vec<_>, _>>()
.map_err(|e| TlsPolicyError(format!("Failed to parse TLS CA PEM '{}': {}", ca_path, e)))?;
let mut roots = RootCertStore::empty();
for cert in ca_certs {
roots
.add(cert)
.map_err(|e| TlsPolicyError(format!("Failed to add CA cert '{}': {}", ca_path, e)))?;
}
Ok(roots)
}
fn build_crypto_provider(policy: &TlsPolicy) -> Result<Arc<CryptoProvider>, TlsPolicyError> {
let mut provider = aws_lc_rs::default_provider();
let cipher_suites = parse_cipher_suites(&provider.cipher_suites, &policy.cipher_suites)?;
if !cipher_suites.is_empty() {
provider.cipher_suites = cipher_suites;
}
Ok(Arc::new(provider))
}
pub fn build_supported_versions(
policy: &TlsPolicy,
) -> Result<Vec<&'static SupportedProtocolVersion>, TlsPolicyError> {
let min_version = parse_go_tls_version(&policy.min_version)?;
let max_version = parse_go_tls_version(&policy.max_version)?;
let versions = [&rustls::version::TLS13, &rustls::version::TLS12]
.into_iter()
.filter(|version| {
let current = go_tls_version_for_supported(version);
min_version.map(|min| current >= min).unwrap_or(true)
&& max_version.map(|max| current <= max).unwrap_or(true)
})
.collect::<Vec<_>>();
if versions.is_empty() {
return Err(TlsPolicyError(format!(
"TLS version range min='{}' max='{}' is unsupported by rustls",
policy.min_version, policy.max_version
)));
}
Ok(versions)
}
fn parse_go_tls_version(value: &str) -> Result<Option<GoTlsVersion>, TlsPolicyError> {
match value.trim() {
"" => Ok(None),
"SSLv3" => Ok(Some(GoTlsVersion::Ssl3)),
"TLS 1.0" => Ok(Some(GoTlsVersion::Tls10)),
"TLS 1.1" => Ok(Some(GoTlsVersion::Tls11)),
"TLS 1.2" => Ok(Some(GoTlsVersion::Tls12)),
"TLS 1.3" => Ok(Some(GoTlsVersion::Tls13)),
other => Err(TlsPolicyError(format!("invalid TLS version {}", other))),
}
}
fn parse_cipher_suites(
available: &[SupportedCipherSuite],
value: &str,
) -> Result<Vec<SupportedCipherSuite>, TlsPolicyError> {
let trimmed = value.trim();
if trimmed.is_empty() {
return Ok(Vec::new());
}
trimmed
.split(',')
.map(|name| {
let suite = parse_cipher_suite_name(name.trim())?;
available
.iter()
.copied()
.find(|candidate| candidate.suite() == suite)
.ok_or_else(|| {
TlsPolicyError(format!(
"TLS cipher suite '{}' is unsupported by the Rust implementation",
name.trim()
))
})
})
.collect()
}
fn parse_cipher_suite_name(value: &str) -> Result<CipherSuite, TlsPolicyError> {
match value {
"TLS_AES_128_GCM_SHA256" => Ok(CipherSuite::TLS13_AES_128_GCM_SHA256),
"TLS_AES_256_GCM_SHA384" => Ok(CipherSuite::TLS13_AES_256_GCM_SHA384),
"TLS_CHACHA20_POLY1305_SHA256" => Ok(CipherSuite::TLS13_CHACHA20_POLY1305_SHA256),
"TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256" => {
Ok(CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256)
}
"TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384" => {
Ok(CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384)
}
"TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256" => {
Ok(CipherSuite::TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256)
}
"TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256" => {
Ok(CipherSuite::TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256)
}
"TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384" => {
Ok(CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384)
}
"TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256" => {
Ok(CipherSuite::TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256)
}
other => Err(TlsPolicyError(format!(
"TLS cipher suite '{}' is unsupported by the Rust implementation",
other
))),
}
}
fn parse_common_name(cert: &CertificateDer<'_>) -> Result<String, TlsPolicyError> {
let (_, certificate) = X509Certificate::from_der(cert.as_ref())
.map_err(|e| TlsPolicyError(format!("parse X.509 certificate failed: {}", e)))?;
let common_name = certificate
.subject()
.iter_common_name()
.next()
.and_then(|common_name| common_name.as_str().ok())
.map(str::to_string);
match common_name {
Some(common_name) => Ok(common_name),
None => Ok(String::new()),
}
}
fn common_name_is_allowed(
common_name: &str,
allowed_common_names: &HashSet<String>,
allowed_wildcard_domain: &str,
) -> bool {
(!allowed_wildcard_domain.is_empty() && common_name.ends_with(allowed_wildcard_domain))
|| allowed_common_names.contains(common_name)
}
fn go_tls_version_for_supported(version: &SupportedProtocolVersion) -> GoTlsVersion {
match version.version {
rustls::ProtocolVersion::TLSv1_2 => GoTlsVersion::Tls12,
rustls::ProtocolVersion::TLSv1_3 => GoTlsVersion::Tls13,
_ => unreachable!("rustls only exposes TLS 1.2 and 1.3"),
}
}
#[cfg(test)]
mod tests {
use super::{TlsPolicy, build_supported_versions, common_name_is_allowed, parse_cipher_suites};
use rustls::crypto::aws_lc_rs;
use std::collections::HashSet;
#[test]
fn test_build_supported_versions_defaults_to_tls12_and_tls13() {
let versions = build_supported_versions(&TlsPolicy::default()).unwrap();
assert_eq!(
versions,
vec![&rustls::version::TLS13, &rustls::version::TLS12]
);
}
#[test]
fn test_build_supported_versions_filters_to_tls13() {
let versions = build_supported_versions(&TlsPolicy {
min_version: "TLS 1.3".to_string(),
max_version: "TLS 1.3".to_string(),
cipher_suites: String::new(),
})
.unwrap();
assert_eq!(versions, vec![&rustls::version::TLS13]);
}
#[test]
fn test_build_supported_versions_rejects_unsupported_legacy_range() {
let err = build_supported_versions(&TlsPolicy {
min_version: "TLS 1.0".to_string(),
max_version: "TLS 1.1".to_string(),
cipher_suites: String::new(),
})
.unwrap_err();
assert!(err.to_string().contains("unsupported by rustls"));
}
#[test]
fn test_parse_cipher_suites_accepts_go_names() {
let cipher_suites = parse_cipher_suites(
&aws_lc_rs::default_provider().cipher_suites,
"TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,TLS_AES_128_GCM_SHA256",
)
.unwrap();
assert_eq!(cipher_suites.len(), 2);
}
#[test]
fn test_common_name_is_allowed_matches_exact_and_wildcard() {
let allowed_common_names =
HashSet::from([String::from("volume-a.internal"), String::from("worker-7")]);
assert!(common_name_is_allowed(
"volume-a.internal",
&allowed_common_names,
"",
));
assert!(common_name_is_allowed(
"node.prod.example.com",
&allowed_common_names,
".example.com",
));
assert!(!common_name_is_allowed(
"node.prod.other.net",
&allowed_common_names,
".example.com",
));
}
}

View File

@@ -0,0 +1,159 @@
use axum::Router;
use axum::body::Body;
use axum::extract::Query;
use axum::http::{StatusCode, header};
use axum::response::{IntoResponse, Response};
use axum::routing::{any, get};
use pprof::protos::Message;
use serde::Deserialize;
#[derive(Deserialize, Default)]
struct ProfileQuery {
seconds: Option<u64>,
}
pub fn build_debug_router() -> Router {
Router::new()
.route("/debug/pprof/", get(pprof_index_handler))
.route("/debug/pprof/cmdline", get(pprof_cmdline_handler))
.route("/debug/pprof/profile", get(pprof_profile_handler))
.route("/debug/pprof/symbol", any(pprof_symbol_handler))
.route("/debug/pprof/trace", get(pprof_trace_handler))
}
async fn pprof_index_handler() -> Response {
let body = concat!(
"<html><head><title>/debug/pprof/</title></head><body>",
"<a href=\"cmdline\">cmdline</a><br>",
"<a href=\"profile\">profile</a><br>",
"<a href=\"symbol\">symbol</a><br>",
"<a href=\"trace\">trace</a><br>",
"</body></html>",
);
(
StatusCode::OK,
[(header::CONTENT_TYPE, "text/html; charset=utf-8")],
body,
)
.into_response()
}
async fn pprof_cmdline_handler() -> Response {
let body = std::env::args().collect::<Vec<_>>().join("\0");
(
StatusCode::OK,
[(header::CONTENT_TYPE, "text/plain; charset=utf-8")],
body,
)
.into_response()
}
async fn pprof_profile_handler(Query(query): Query<ProfileQuery>) -> Response {
let seconds = query.seconds.unwrap_or(30).clamp(1, 300);
let guard = match pprof::ProfilerGuard::new(100) {
Ok(guard) => guard,
Err(e) => {
return (
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to start profiler: {}", e),
)
.into_response();
}
};
tokio::time::sleep(std::time::Duration::from_secs(seconds)).await;
let report = match guard.report().build() {
Ok(report) => report,
Err(e) => {
return (
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to build profile report: {}", e),
)
.into_response();
}
};
let profile = match report.pprof() {
Ok(profile) => profile,
Err(e) => {
return (
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to encode profile: {}", e),
)
.into_response();
}
};
let mut bytes = Vec::new();
if let Err(e) = profile.encode(&mut bytes) {
return (
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to serialize profile: {}", e),
)
.into_response();
}
(
StatusCode::OK,
[(header::CONTENT_TYPE, "application/octet-stream")],
bytes,
)
.into_response()
}
async fn pprof_symbol_handler() -> Response {
(
StatusCode::OK,
[(header::CONTENT_TYPE, "text/plain; charset=utf-8")],
"num_symbols: 0\n",
)
.into_response()
}
async fn pprof_trace_handler(Query(query): Query<ProfileQuery>) -> Response {
let seconds = query.seconds.unwrap_or(1).clamp(1, 30);
tokio::time::sleep(std::time::Duration::from_secs(seconds)).await;
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "application/octet-stream")
.body(Body::from(Vec::<u8>::new()))
.unwrap()
}
#[cfg(test)]
mod tests {
use super::*;
use axum::http::Request;
use tower::ServiceExt;
#[tokio::test]
async fn test_debug_index_route() {
let app = build_debug_router();
let response = app
.oneshot(
Request::builder()
.uri("/debug/pprof/")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn test_debug_cmdline_route() {
let app = build_debug_router();
let response = app
.oneshot(
Request::builder()
.uri("/debug/pprof/cmdline")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
}

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.6 KiB

View File

@@ -0,0 +1,370 @@
use std::error::Error;
use std::fmt;
use std::time::Duration;
use hyper::http::Uri;
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};
use crate::config::VolumeServerConfig;
pub const GRPC_MAX_MESSAGE_SIZE: usize = 1 << 30;
const GRPC_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(60);
const GRPC_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(20);
const GRPC_INITIAL_WINDOW_SIZE: u32 = 16 * 1024 * 1024;
#[derive(Clone, Debug)]
pub struct OutgoingGrpcTlsConfig {
cert_pem: String,
key_pem: String,
ca_pem: String,
}
#[derive(Debug)]
pub struct GrpcClientError(String);
impl fmt::Display for GrpcClientError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
impl Error for GrpcClientError {}
pub fn load_outgoing_grpc_tls(
config: &VolumeServerConfig,
) -> Result<Option<OutgoingGrpcTlsConfig>, GrpcClientError> {
// prefer a dedicated client certificate: CAs may issue certs with only one of the serverAuth/clientAuth EKUs
let (cert_file, key_file) = if !config.grpc_client_cert_file.is_empty()
&& !config.grpc_client_key_file.is_empty()
{
(&config.grpc_client_cert_file, &config.grpc_client_key_file)
} else {
if !config.grpc_client_cert_file.is_empty() || !config.grpc_client_key_file.is_empty() {
tracing::warn!(
"grpc.volume.client_cert and grpc.volume.client_key must both be set, falling back to grpc.volume.cert and grpc.volume.key"
);
}
(&config.grpc_cert_file, &config.grpc_key_file)
};
if cert_file.is_empty() || key_file.is_empty() || config.grpc_ca_file.is_empty() {
return Ok(None);
}
let cert_pem = std::fs::read_to_string(cert_file).map_err(|e| {
GrpcClientError(format!(
"Failed to read outgoing gRPC cert '{}': {}",
cert_file, e
))
})?;
let key_pem = std::fs::read_to_string(key_file).map_err(|e| {
GrpcClientError(format!(
"Failed to read outgoing gRPC key '{}': {}",
key_file, e
))
})?;
let ca_pem = std::fs::read_to_string(&config.grpc_ca_file).map_err(|e| {
GrpcClientError(format!(
"Failed to read outgoing gRPC CA '{}': {}",
config.grpc_ca_file, e
))
})?;
Ok(Some(OutgoingGrpcTlsConfig {
cert_pem,
key_pem,
ca_pem,
}))
}
pub fn grpc_endpoint_uri(grpc_host_port: &str, tls: Option<&OutgoingGrpcTlsConfig>) -> String {
let scheme = if tls.is_some() { "https" } else { "http" };
format!("{}://{}", scheme, grpc_host_port)
}
pub fn build_grpc_endpoint(
grpc_host_port: &str,
tls: Option<&OutgoingGrpcTlsConfig>,
) -> Result<Endpoint, GrpcClientError> {
let uri = grpc_endpoint_uri(grpc_host_port, tls);
let mut endpoint = Channel::from_shared(uri.clone())
.map_err(|e| GrpcClientError(format!("invalid gRPC endpoint {}: {}", uri, e)))?
.http2_keep_alive_interval(GRPC_KEEPALIVE_INTERVAL)
.keep_alive_timeout(GRPC_KEEPALIVE_TIMEOUT)
.keep_alive_while_idle(false)
.initial_stream_window_size(Some(GRPC_INITIAL_WINDOW_SIZE))
.initial_connection_window_size(Some(GRPC_INITIAL_WINDOW_SIZE))
.http2_adaptive_window(false);
if let Some(tls) = tls {
let parsed = uri
.parse::<Uri>()
.map_err(|e| GrpcClientError(format!("invalid gRPC endpoint {}: {}", uri, e)))?;
let host = parsed
.host()
.ok_or_else(|| GrpcClientError(format!("missing host in gRPC endpoint {}", uri)))?;
let tls_config = ClientTlsConfig::new()
.identity(Identity::from_pem(
tls.cert_pem.clone(),
tls.key_pem.clone(),
))
.ca_certificate(Certificate::from_pem(tls.ca_pem.clone()))
.domain_name(host.to_string());
endpoint = endpoint.tls_config(tls_config).map_err(|e| {
GrpcClientError(format!("configure gRPC TLS for {} failed: {}", uri, e))
})?;
}
Ok(endpoint)
}
/// Parse a SeaweedFS server address (`"ip:port.grpcPort"` or
/// `"ip:port"`) into the `host:grpcPort` form `build_grpc_endpoint`
/// expects. With the trailing `.grpcPort` segment, that segment IS
/// the gRPC port; without it, the gRPC port is `port + 10000`
/// (SeaweedFS's HTTP↔gRPC port-offset convention).
///
/// Shared between `grpc_server.rs` and the distributed-EC-read path
/// in `store_ec.rs` — keep this as the single source of truth so the
/// HTTP↔gRPC port translation can't drift between callers.
pub fn parse_grpc_address(source: &str) -> Result<String, String> {
let colon_idx = source
.rfind(':')
.ok_or_else(|| format!("cannot parse address: {}", source))?;
let host = &source[..colon_idx];
let port_part = &source[colon_idx + 1..];
if let Some(dot_idx) = port_part.rfind('.') {
// Format: "ip:port.grpcPort". Validate BOTH ports as u16
// so a malformed HTTP port (e.g. `host:abc.18080`) is
// rejected here rather than tripping a downstream
// `build_grpc_endpoint` URI parse failure with a less
// useful error.
let http_port = &port_part[..dot_idx];
let grpc_port = &port_part[dot_idx + 1..];
http_port
.parse::<u16>()
.map_err(|e| format!("invalid http port {:?}: {}", http_port, e))?;
grpc_port
.parse::<u16>()
.map_err(|e| format!("invalid grpc port {:?}: {}", grpc_port, e))?;
return Ok(format!("{}:{}", host, grpc_port));
}
// Format: "ip:port" → grpc = port + 10000. Reject inputs whose
// implicit grpc port would overflow the TCP port range (e.g.
// `host:60000` produces 70000 — invalid). Without this check
// the cast silently wraps and the endpoint call later fails
// with an opaque connection error.
let port: u16 = port_part
.parse()
.map_err(|e| format!("invalid port {:?}: {}", port_part, e))?;
let grpc_port = port as u32 + 10000;
if grpc_port > u16::MAX as u32 {
return Err(format!(
"implicit grpc port out of range: {} + 10000 = {}",
port, grpc_port
));
}
Ok(format!("{}:{}", host, grpc_port))
}
#[cfg(test)]
mod tests {
use super::{build_grpc_endpoint, grpc_endpoint_uri, load_outgoing_grpc_tls};
use crate::config::{NeedleMapKind, ReadMode, VolumeServerConfig};
use crate::security::tls::TlsPolicy;
const TEST_CERT_PEM: &str = "-----BEGIN CERTIFICATE-----\nMIIBPDCB76ADAgECAhRuRPQgeAu43BT/M7EfAWSdapVdYDAFBgMrZXAwFDESMBAG\nA1UEAwwJbG9jYWxob3N0MB4XDTI2MDcwNTE2MTUyOVoXDTM2MDcwMjE2MTUyOVow\nFDESMBAGA1UEAwwJbG9jYWxob3N0MCowBQYDK2VwAyEAr/3bNIFI+8V32oCiY6y+\nXRFmZpdNQ2g//VtRkT+nQg+jUzBRMB0GA1UdDgQWBBTsy9tLf1zPiXCQfgci6zNi\ndEzRSjAfBgNVHSMEGDAWgBTsy9tLf1zPiXCQfgci6zNidEzRSjAPBgNVHRMBAf8E\nBTADAQH/MAUGAytlcANBAIvsdw0IbvOBBkb9cd7BfMJfIP9pQQrAL03pCRWJFnFh\nSysaLVgFXI4T078IiaM874oO+iB+5vNbWEpc7CkGow4=\n-----END CERTIFICATE-----\n";
const TEST_KEY_PEM: &str = "-----BEGIN PRIVATE KEY-----\nMC4CAQAwBQYDK2VwBCIEIHbyn71Kk+Y7KT3sBctit7uZpErpoH6qDbFj6P8qGaZH\n-----END PRIVATE KEY-----\n";
#[test]
fn test_build_grpc_endpoint_with_tls_resolves_crypto_provider() {
crate::security::tls::install_default_crypto_provider();
let tls = super::OutgoingGrpcTlsConfig {
cert_pem: TEST_CERT_PEM.to_string(),
key_pem: TEST_KEY_PEM.to_string(),
ca_pem: TEST_CERT_PEM.to_string(),
};
let endpoint = build_grpc_endpoint("127.0.0.1:19333", Some(&tls)).unwrap();
assert_eq!(endpoint.uri().scheme_str(), Some("https"));
}
fn sample_config() -> VolumeServerConfig {
VolumeServerConfig {
port: 8080,
grpc_port: 18080,
public_port: 8080,
ip: "127.0.0.1".to_string(),
bind_ip: String::new(),
public_url: "127.0.0.1:8080".to_string(),
id: String::new(),
masters: vec![],
pre_stop_seconds: 0,
idle_timeout: 0,
data_center: String::new(),
rack: String::new(),
index_type: NeedleMapKind::InMemory,
disk_type: String::new(),
folders: vec![],
folder_max_limits: vec![],
folder_tags: vec![],
min_free_spaces: vec![],
disk_types: vec![],
idx_folder: String::new(),
white_list: vec![],
fix_jpg_orientation: false,
read_mode: ReadMode::Local,
allow_untrusted_remote_endpoints: false,
cpu_profile: String::new(),
mem_profile: String::new(),
compaction_byte_per_second: 0,
maintenance_byte_per_second: 0,
file_size_limit_bytes: 0,
concurrent_upload_limit: 0,
concurrent_download_limit: 0,
inflight_upload_data_timeout: std::time::Duration::from_secs(0),
inflight_download_data_timeout: std::time::Duration::from_secs(0),
has_slow_read: false,
read_buffer_size_mb: 0,
ldb_timeout: 0,
pprof: false,
metrics_port: 0,
metrics_ip: String::new(),
debug: false,
debug_port: 0,
ui_enabled: false,
jwt_signing_key: vec![],
jwt_signing_expires_seconds: 0,
jwt_read_signing_key: vec![],
jwt_read_signing_expires_seconds: 0,
https_cert_file: String::new(),
https_key_file: String::new(),
https_ca_file: String::new(),
https_client_enabled: false,
https_client_cert_file: String::new(),
https_client_key_file: String::new(),
https_client_ca_file: String::new(),
grpc_cert_file: String::new(),
grpc_key_file: String::new(),
grpc_client_cert_file: String::new(),
grpc_client_key_file: String::new(),
grpc_ca_file: String::new(),
grpc_allowed_wildcard_domain: String::new(),
grpc_volume_allowed_common_names: vec![],
tls_policy: TlsPolicy::default(),
enable_write_queue: false,
security_file: String::new(),
}
}
#[test]
fn test_grpc_endpoint_uri_uses_https_when_tls_enabled() {
let tls = super::OutgoingGrpcTlsConfig {
cert_pem: "cert".to_string(),
key_pem: "key".to_string(),
ca_pem: "ca".to_string(),
};
assert_eq!(
grpc_endpoint_uri("master.example.com:19333", Some(&tls)),
"https://master.example.com:19333"
);
}
#[test]
fn test_load_outgoing_grpc_tls_requires_cert_key_and_ca() {
let mut config = sample_config();
config.grpc_cert_file = "/tmp/client.pem".to_string();
assert!(load_outgoing_grpc_tls(&config).unwrap().is_none());
}
fn write_pem_files(dir: &tempfile::TempDir, config: &mut VolumeServerConfig) {
let write = |name: &str, content: &str| {
let path = dir.path().join(name);
std::fs::write(&path, content).unwrap();
path.to_str().unwrap().to_string()
};
config.grpc_cert_file = write("server.pem", "server-cert");
config.grpc_key_file = write("server.key", "server-key");
config.grpc_ca_file = write("ca.pem", "ca");
}
#[test]
fn test_load_outgoing_grpc_tls_prefers_client_cert() {
let dir = tempfile::TempDir::new().unwrap();
let mut config = sample_config();
write_pem_files(&dir, &mut config);
let client_cert = dir.path().join("client.pem");
let client_key = dir.path().join("client.key");
std::fs::write(&client_cert, "client-cert").unwrap();
std::fs::write(&client_key, "client-key").unwrap();
config.grpc_client_cert_file = client_cert.to_str().unwrap().to_string();
config.grpc_client_key_file = client_key.to_str().unwrap().to_string();
let tls = load_outgoing_grpc_tls(&config).unwrap().unwrap();
assert_eq!(tls.cert_pem, "client-cert");
assert_eq!(tls.key_pem, "client-key");
}
#[test]
fn test_load_outgoing_grpc_tls_falls_back_to_server_cert() {
let dir = tempfile::TempDir::new().unwrap();
let mut config = sample_config();
write_pem_files(&dir, &mut config);
let tls = load_outgoing_grpc_tls(&config).unwrap().unwrap();
assert_eq!(tls.cert_pem, "server-cert");
assert_eq!(tls.key_pem, "server-key");
}
#[test]
fn test_build_grpc_endpoint_without_tls_uses_http_scheme() {
let endpoint = build_grpc_endpoint("127.0.0.1:19333", None).unwrap();
assert_eq!(endpoint.uri().scheme_str(), Some("http"));
}
#[test]
fn test_parse_grpc_address_dotted_form() {
use super::parse_grpc_address;
assert_eq!(
parse_grpc_address("127.0.0.1:8080.18080").unwrap(),
"127.0.0.1:18080"
);
}
#[test]
fn test_parse_grpc_address_implicit_form_adds_10000() {
use super::parse_grpc_address;
assert_eq!(
parse_grpc_address("127.0.0.1:8080").unwrap(),
"127.0.0.1:18080"
);
}
#[test]
fn test_parse_grpc_address_rejects_non_numeric_http_port_in_dotted_form() {
use super::parse_grpc_address;
let err = parse_grpc_address("host:abc.18080").unwrap_err();
assert!(err.contains("invalid http port"), "{}", err);
}
#[test]
fn test_parse_grpc_address_rejects_non_numeric_grpc_port_in_dotted_form() {
use super::parse_grpc_address;
let err = parse_grpc_address("host:8080.xyz").unwrap_err();
assert!(err.contains("invalid grpc port"), "{}", err);
}
#[test]
fn test_parse_grpc_address_rejects_implicit_port_that_overflows() {
use super::parse_grpc_address;
let err = parse_grpc_address("127.0.0.1:60000").unwrap_err();
assert!(err.contains("out of range"), "{}", err);
}
#[test]
fn test_parse_grpc_address_rejects_input_without_colon() {
use super::parse_grpc_address;
let err = parse_grpc_address("hostname").unwrap_err();
assert!(err.contains("cannot parse"), "{}", err);
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,102 @@
use crate::pb::volume_server_pb;
pub fn collect_mem_status() -> volume_server_pb::MemStatus {
#[allow(unused_mut)]
let mut mem = volume_server_pb::MemStatus {
goroutines: 1,
..Default::default()
};
#[cfg(target_os = "linux")]
{
if let Some((all, free)) = get_system_memory_linux() {
mem.all = all;
mem.free = free;
mem.used = all.saturating_sub(free);
}
if let Some(status) = read_process_status_linux() {
if status.threads > 0 {
mem.goroutines = status.threads as i32;
}
if let Some(rss) = status.rss {
mem.self_ = rss;
}
if let Some(heap) = status.data.or(status.rss) {
mem.heap = heap;
}
if let Some(stack) = status.stack {
mem.stack = stack;
}
}
}
mem
}
#[cfg(target_os = "linux")]
fn get_system_memory_linux() -> Option<(u64, u64)> {
unsafe {
let mut info: libc::sysinfo = std::mem::zeroed();
if libc::sysinfo(&mut info) == 0 {
let unit = info.mem_unit as u64;
let total = info.totalram as u64 * unit;
let free = info.freeram as u64 * unit;
return Some((total, free));
}
}
None
}
#[cfg(target_os = "linux")]
#[derive(Default)]
struct ProcessStatus {
threads: u64,
rss: Option<u64>,
data: Option<u64>,
stack: Option<u64>,
}
#[cfg(target_os = "linux")]
fn read_process_status_linux() -> Option<ProcessStatus> {
let status = std::fs::read_to_string("/proc/self/status").ok()?;
let mut out = ProcessStatus::default();
for line in status.lines() {
if let Some(value) = line.strip_prefix("Threads:") {
out.threads = value.trim().parse().ok()?;
continue;
}
if let Some(value) = parse_proc_status_kib_field(line, "VmRSS:") {
out.rss = Some(value);
continue;
}
if let Some(value) = parse_proc_status_kib_field(line, "VmData:") {
out.data = Some(value);
continue;
}
if let Some(value) = parse_proc_status_kib_field(line, "VmStk:") {
out.stack = Some(value);
}
}
Some(out)
}
#[cfg(target_os = "linux")]
fn parse_proc_status_kib_field(line: &str, prefix: &str) -> Option<u64> {
let raw = line.strip_prefix(prefix)?.trim();
let value = raw.strip_suffix(" kB").unwrap_or(raw).trim();
value.parse::<u64>().ok().map(|kib| kib * 1024)
}
#[cfg(test)]
mod tests {
use super::collect_mem_status;
#[test]
fn test_collect_mem_status_reports_live_process_state() {
let mem = collect_mem_status();
assert!(mem.goroutines > 0);
}
}

View File

@@ -0,0 +1,15 @@
#[cfg(unix)]
pub mod debug;
pub mod grpc_client;
pub mod grpc_server;
pub mod handlers;
pub mod heartbeat;
pub mod memory_status;
#[cfg(unix)]
pub mod profiling;
pub mod request_id;
pub mod server_stats;
pub mod store_ec;
pub mod ui;
pub mod volume_server;
pub mod write_queue;

View File

@@ -0,0 +1,190 @@
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use pprof::protos::Message;
use crate::config::VolumeServerConfig;
const GO_CPU_PROFILE_FREQUENCY: i32 = 100;
const GO_PPROF_BLOCKLIST: [&str; 4] = ["libc", "libgcc", "pthread", "vdso"];
pub struct CpuProfileSession {
output_path: PathBuf,
guard: pprof::ProfilerGuard<'static>,
}
impl CpuProfileSession {
pub fn start(config: &VolumeServerConfig) -> Result<Option<Self>, String> {
if config.cpu_profile.is_empty() {
if !config.mem_profile.is_empty() && !config.pprof {
tracing::warn!(
"--memprofile is not yet supported in the Rust volume server; ignoring '{}'",
config.mem_profile
);
}
return Ok(None);
}
if config.pprof {
tracing::info!(
"--pprof is enabled; ignoring --cpuprofile '{}' and --memprofile '{}'",
config.cpu_profile,
config.mem_profile
);
return Ok(None);
}
if !config.mem_profile.is_empty() {
tracing::warn!(
"--memprofile is not yet supported in the Rust volume server; only --cpuprofile '{}' will be written",
config.cpu_profile
);
}
let guard = pprof::ProfilerGuardBuilder::default()
.frequency(GO_CPU_PROFILE_FREQUENCY)
.blocklist(&GO_PPROF_BLOCKLIST)
.build()
.map_err(|e| {
format!(
"Failed to start CPU profiler '{}': {}",
config.cpu_profile, e
)
})?;
Ok(Some(Self {
output_path: PathBuf::from(&config.cpu_profile),
guard,
}))
}
pub fn finish(self) -> Result<(), String> {
let report = self
.guard
.report()
.build()
.map_err(|e| format!("Failed to build CPU profile report: {}", e))?;
let profile = report
.pprof()
.map_err(|e| format!("Failed to encode CPU profile report: {}", e))?;
let mut bytes = Vec::new();
profile
.encode(&mut bytes)
.map_err(|e| format!("Failed to serialize CPU profile report: {}", e))?;
let mut file = File::create(&self.output_path).map_err(|e| {
format!(
"Failed to create CPU profile '{}': {}",
self.output_path.display(),
e
)
})?;
file.write_all(&bytes).map_err(|e| {
format!(
"Failed to write CPU profile '{}': {}",
self.output_path.display(),
e
)
})?;
file.flush().map_err(|e| {
format!(
"Failed to flush CPU profile '{}': {}",
self.output_path.display(),
e
)
})?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::CpuProfileSession;
use crate::config::{NeedleMapKind, ReadMode, VolumeServerConfig};
use crate::security::tls::TlsPolicy;
fn sample_config() -> VolumeServerConfig {
VolumeServerConfig {
port: 8080,
grpc_port: 18080,
public_port: 8080,
ip: "127.0.0.1".to_string(),
bind_ip: "127.0.0.1".to_string(),
public_url: "127.0.0.1:8080".to_string(),
id: "127.0.0.1:8080".to_string(),
masters: vec![],
pre_stop_seconds: 0,
idle_timeout: 0,
data_center: String::new(),
rack: String::new(),
index_type: NeedleMapKind::InMemory,
disk_type: String::new(),
folders: vec!["/tmp".to_string()],
folder_max_limits: vec![8],
folder_tags: vec![vec![]],
min_free_spaces: vec![],
disk_types: vec![String::new()],
idx_folder: String::new(),
white_list: vec![],
fix_jpg_orientation: false,
read_mode: ReadMode::Local,
allow_untrusted_remote_endpoints: false,
cpu_profile: String::new(),
mem_profile: String::new(),
compaction_byte_per_second: 0,
maintenance_byte_per_second: 0,
file_size_limit_bytes: 0,
concurrent_upload_limit: 0,
concurrent_download_limit: 0,
inflight_upload_data_timeout: std::time::Duration::from_secs(0),
inflight_download_data_timeout: std::time::Duration::from_secs(0),
has_slow_read: false,
read_buffer_size_mb: 4,
ldb_timeout: 0,
pprof: false,
metrics_port: 0,
metrics_ip: String::new(),
debug: false,
debug_port: 0,
ui_enabled: false,
jwt_signing_key: vec![],
jwt_signing_expires_seconds: 0,
jwt_read_signing_key: vec![],
jwt_read_signing_expires_seconds: 0,
https_cert_file: String::new(),
https_key_file: String::new(),
https_ca_file: String::new(),
https_client_enabled: false,
https_client_cert_file: String::new(),
https_client_key_file: String::new(),
https_client_ca_file: String::new(),
grpc_cert_file: String::new(),
grpc_key_file: String::new(),
grpc_client_cert_file: String::new(),
grpc_client_key_file: String::new(),
grpc_ca_file: String::new(),
grpc_allowed_wildcard_domain: String::new(),
grpc_volume_allowed_common_names: vec![],
tls_policy: TlsPolicy::default(),
enable_write_queue: false,
security_file: String::new(),
}
}
#[test]
fn test_cpu_profile_session_skips_when_disabled() {
let config = sample_config();
assert!(CpuProfileSession::start(&config).unwrap().is_none());
}
#[test]
fn test_cpu_profile_session_skips_when_pprof_enabled() {
let mut config = sample_config();
config.cpu_profile = "/tmp/cpu.pb".to_string();
config.pprof = true;
assert!(CpuProfileSession::start(&config).unwrap().is_none());
}
}

View File

@@ -0,0 +1,137 @@
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use hyper::http::{self, HeaderValue};
use tonic::metadata::MetadataValue;
use tonic::{Request, Status};
use tower::{Layer, Service};
tokio::task_local! {
static CURRENT_REQUEST_ID: String;
}
#[derive(Clone, Debug, Default)]
pub struct GrpcRequestIdLayer;
#[derive(Clone, Debug)]
pub struct GrpcRequestIdService<S> {
inner: S,
}
impl<S> Layer<S> for GrpcRequestIdLayer {
type Service = GrpcRequestIdService<S>;
fn layer(&self, inner: S) -> Self::Service {
GrpcRequestIdService { inner }
}
}
impl<S, B> Service<http::Request<B>> for GrpcRequestIdService<S>
where
S: Service<http::Request<B>, Response = http::Response<tonic::body::Body>> + Send + 'static,
S::Future: Send + 'static,
B: Send + 'static,
{
type Response = http::Response<tonic::body::Body>;
type Error = S::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, mut request: http::Request<B>) -> Self::Future {
let request_id = match request.headers().get("x-amz-request-id") {
Some(value) => match value.to_str() {
Ok(value) if !value.is_empty() => value.to_owned(),
_ => generate_grpc_request_id(),
},
None => generate_grpc_request_id(),
};
if let Ok(value) = HeaderValue::from_str(&request_id) {
request.headers_mut().insert("x-amz-request-id", value);
}
let future = self.inner.call(request);
Box::pin(async move {
let mut response: http::Response<tonic::body::Body> =
scope_request_id(request_id.clone(), future).await?;
if let Ok(value) = HeaderValue::from_str(&request_id) {
response.headers_mut().insert("x-amz-request-id", value);
}
Ok(response)
})
}
}
pub async fn scope_request_id<F, T>(request_id: String, future: F) -> T
where
F: Future<Output = T>,
{
CURRENT_REQUEST_ID.scope(request_id, future).await
}
pub fn current_request_id() -> Option<String> {
CURRENT_REQUEST_ID.try_with(Clone::clone).ok()
}
pub fn outgoing_request_id_interceptor(mut request: Request<()>) -> Result<Request<()>, Status> {
if let Some(request_id) = current_request_id() {
let value = MetadataValue::try_from(request_id.as_str())
.map_err(|_| Status::internal("invalid scoped request id"))?;
request.metadata_mut().insert("x-amz-request-id", value);
}
Ok(request)
}
pub fn generate_http_request_id() -> String {
use rand::RngExt;
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_nanos() as u64;
let rand_val: u32 = rand::rng().random();
format!("{:X}{:08X}", nanos, rand_val)
}
fn generate_grpc_request_id() -> String {
uuid::Uuid::new_v4().to_string()
}
#[cfg(test)]
mod tests {
use super::{current_request_id, outgoing_request_id_interceptor, scope_request_id};
use tonic::Request;
#[tokio::test]
async fn test_scope_request_id_exposes_current_value() {
let request_id = "req-123".to_string();
let current = scope_request_id(
request_id.clone(),
async move { current_request_id().unwrap() },
)
.await;
assert_eq!(current, request_id);
}
#[tokio::test]
async fn test_outgoing_request_id_interceptor_propagates_scope() {
let request = scope_request_id("req-456".to_string(), async move {
outgoing_request_id_interceptor(Request::new(())).unwrap()
})
.await;
assert_eq!(
request
.metadata()
.get("x-amz-request-id")
.unwrap()
.to_str()
.unwrap(),
"req-456"
);
}
}

View File

@@ -0,0 +1,248 @@
use chrono::{Datelike, Local, Timelike};
use serde::Serialize;
use std::sync::{LazyLock, Mutex};
use std::time::Instant;
static START_TIME: LazyLock<Instant> = LazyLock::new(Instant::now);
static SERVER_STATS: LazyLock<ServerStats> = LazyLock::new(ServerStats::default);
#[derive(Default)]
pub struct ServerStats {
inner: Mutex<ServerStatsInner>,
}
#[derive(Default)]
struct ServerStatsInner {
requests: DurationCounter,
connections: DurationCounter,
assign_requests: DurationCounter,
read_requests: DurationCounter,
write_requests: DurationCounter,
delete_requests: DurationCounter,
bytes_in: DurationCounter,
bytes_out: DurationCounter,
}
#[derive(Clone, Serialize)]
#[serde(rename_all = "PascalCase")]
pub struct ServerStatsSnapshot {
pub requests: DurationCounterSnapshot,
pub connections: DurationCounterSnapshot,
pub assign_requests: DurationCounterSnapshot,
pub read_requests: DurationCounterSnapshot,
pub write_requests: DurationCounterSnapshot,
pub delete_requests: DurationCounterSnapshot,
pub bytes_in: DurationCounterSnapshot,
pub bytes_out: DurationCounterSnapshot,
}
#[derive(Clone, Serialize)]
#[serde(rename_all = "PascalCase")]
pub struct DurationCounterSnapshot {
pub minute_counter: RoundRobinCounterSnapshot,
pub hour_counter: RoundRobinCounterSnapshot,
pub day_counter: RoundRobinCounterSnapshot,
pub week_counter: RoundRobinCounterSnapshot,
}
#[derive(Clone, Serialize)]
#[serde(rename_all = "PascalCase")]
pub struct RoundRobinCounterSnapshot {
pub last_index: i32,
pub values: Vec<i64>,
pub counts: Vec<i64>,
}
#[derive(Clone)]
struct DurationCounter {
minute_counter: RoundRobinCounter,
hour_counter: RoundRobinCounter,
day_counter: RoundRobinCounter,
week_counter: RoundRobinCounter,
}
#[derive(Clone)]
struct RoundRobinCounter {
last_index: i32,
values: Vec<i64>,
counts: Vec<i64>,
}
impl Default for DurationCounter {
fn default() -> Self {
Self {
minute_counter: RoundRobinCounter::new(60),
hour_counter: RoundRobinCounter::new(60),
day_counter: RoundRobinCounter::new(24),
week_counter: RoundRobinCounter::new(7),
}
}
}
impl RoundRobinCounter {
fn new(slots: usize) -> Self {
Self {
last_index: -1,
values: vec![0; slots],
counts: vec![0; slots],
}
}
fn add(&mut self, index: usize, val: i64) {
if index >= self.values.len() {
return;
}
while self.last_index != index as i32 {
self.last_index = (self.last_index + 1).rem_euclid(self.values.len() as i32);
self.values[self.last_index as usize] = 0;
self.counts[self.last_index as usize] = 0;
}
self.values[index] += val;
self.counts[index] += 1;
}
fn snapshot(&self) -> RoundRobinCounterSnapshot {
RoundRobinCounterSnapshot {
last_index: self.last_index,
values: self.values.clone(),
counts: self.counts.clone(),
}
}
}
impl DurationCounter {
fn add_now(&mut self, val: i64) {
let now = Local::now();
self.minute_counter.add(now.second() as usize, val);
self.hour_counter.add(now.minute() as usize, val);
self.day_counter.add(now.hour() as usize, val);
self.week_counter
.add(now.weekday().num_days_from_sunday() as usize, val);
}
fn snapshot(&self) -> DurationCounterSnapshot {
DurationCounterSnapshot {
minute_counter: self.minute_counter.snapshot(),
hour_counter: self.hour_counter.snapshot(),
day_counter: self.day_counter.snapshot(),
week_counter: self.week_counter.snapshot(),
}
}
}
impl ServerStatsInner {
fn snapshot(&self) -> ServerStatsSnapshot {
ServerStatsSnapshot {
requests: self.requests.snapshot(),
connections: self.connections.snapshot(),
assign_requests: self.assign_requests.snapshot(),
read_requests: self.read_requests.snapshot(),
write_requests: self.write_requests.snapshot(),
delete_requests: self.delete_requests.snapshot(),
bytes_in: self.bytes_in.snapshot(),
bytes_out: self.bytes_out.snapshot(),
}
}
}
impl ServerStats {
fn update<F>(&self, update: F)
where
F: FnOnce(&mut ServerStatsInner),
{
let mut inner = self.inner.lock().unwrap();
update(&mut inner);
}
fn snapshot(&self) -> ServerStatsSnapshot {
self.inner.lock().unwrap().snapshot()
}
}
impl RoundRobinCounterSnapshot {
pub fn to_list(&self) -> Vec<i64> {
if self.values.is_empty() {
return Vec::new();
}
let mut ret = Vec::with_capacity(self.values.len());
let mut index = self.last_index;
let mut step = self.values.len();
while step > 0 {
step -= 1;
index += 1;
if index >= self.values.len() as i32 {
index = 0;
}
ret.push(self.values[index as usize]);
}
ret
}
}
pub fn init_process_start() {
LazyLock::force(&START_TIME);
LazyLock::force(&SERVER_STATS);
}
pub fn uptime_string() -> String {
let secs = START_TIME.elapsed().as_secs();
let hours = secs / 3600;
let minutes = (secs % 3600) / 60;
let seconds = secs % 60;
let mut out = String::new();
if hours > 0 {
out.push_str(&format!("{}h", hours));
}
if hours > 0 || minutes > 0 {
out.push_str(&format!("{}m", minutes));
}
out.push_str(&format!("{}s", seconds));
out
}
pub fn snapshot() -> ServerStatsSnapshot {
SERVER_STATS.snapshot()
}
pub fn record_request_open() {
SERVER_STATS.update(|inner| inner.requests.add_now(1));
}
pub fn record_request_close() {
SERVER_STATS.update(|inner| inner.requests.add_now(-1));
}
pub fn record_connection_open() {
SERVER_STATS.update(|inner| inner.connections.add_now(1));
}
pub fn record_connection_close() {
SERVER_STATS.update(|inner| inner.connections.add_now(-1));
}
pub fn record_read_request() {
SERVER_STATS.update(|inner| inner.read_requests.add_now(1));
}
pub fn record_write_request() {
SERVER_STATS.update(|inner| inner.write_requests.add_now(1));
}
pub fn record_delete_request() {
SERVER_STATS.update(|inner| inner.delete_requests.add_now(1));
}
pub fn record_bytes_in(bytes: i64) {
SERVER_STATS.update(|inner| inner.bytes_in.add_now(bytes));
}
pub fn record_bytes_out(bytes: i64) {
SERVER_STATS.update(|inner| inner.bytes_out.add_now(bytes));
}
#[cfg(test)]
pub fn reset_for_tests() {
LazyLock::force(&START_TIME);
let mut inner = SERVER_STATS.inner.lock().unwrap();
*inner = ServerStatsInner::default();
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,506 @@
use std::fmt::Write as _;
use crate::server::server_stats;
use crate::server::volume_server::VolumeServerState;
use crate::storage::store::Store;
pub struct EmbeddedAsset {
pub content_type: &'static str,
pub bytes: &'static [u8],
}
struct UiDiskRow {
dir: String,
disk_type: String,
all: u64,
free: u64,
used: u64,
}
struct UiVolumeRow {
id: u32,
collection: String,
disk_type: String,
size: u64,
file_count: i64,
delete_count: i64,
deleted_byte_count: u64,
ttl: String,
read_only: bool,
version: u32,
remote_storage_name: String,
remote_storage_key: String,
}
struct UiEcShardRow {
shard_id: u8,
size: u64,
}
struct UiEcVolumeRow {
volume_id: u32,
collection: String,
size: u64,
shards: Vec<UiEcShardRow>,
created_at: String,
}
pub fn favicon_asset() -> EmbeddedAsset {
EmbeddedAsset {
content_type: "image/x-icon",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/favicon.ico"
)),
}
}
pub fn lookup_static_asset(path: &str) -> Option<EmbeddedAsset> {
let path = path.trim_start_matches('/');
let asset = match path {
"bootstrap/3.3.1/css/bootstrap.min.css" => EmbeddedAsset {
content_type: "text/css; charset=utf-8",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/bootstrap/3.3.1/css/bootstrap.min.css"
)),
},
"bootstrap/3.3.1/fonts/glyphicons-halflings-regular.eot" => EmbeddedAsset {
content_type: "application/vnd.ms-fontobject",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/bootstrap/3.3.1/fonts/glyphicons-halflings-regular.eot"
)),
},
"bootstrap/3.3.1/fonts/glyphicons-halflings-regular.svg" => EmbeddedAsset {
content_type: "image/svg+xml",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/bootstrap/3.3.1/fonts/glyphicons-halflings-regular.svg"
)),
},
"bootstrap/3.3.1/fonts/glyphicons-halflings-regular.ttf" => EmbeddedAsset {
content_type: "font/ttf",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/bootstrap/3.3.1/fonts/glyphicons-halflings-regular.ttf"
)),
},
"bootstrap/3.3.1/fonts/glyphicons-halflings-regular.woff" => EmbeddedAsset {
content_type: "font/woff",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/bootstrap/3.3.1/fonts/glyphicons-halflings-regular.woff"
)),
},
"images/folder.gif" => EmbeddedAsset {
content_type: "image/gif",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/images/folder.gif"
)),
},
"javascript/jquery-3.6.0.min.js" => EmbeddedAsset {
content_type: "application/javascript; charset=utf-8",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/javascript/jquery-3.6.0.min.js"
)),
},
"javascript/jquery-sparklines/2.1.2/jquery.sparkline.min.js" => EmbeddedAsset {
content_type: "application/javascript; charset=utf-8",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/javascript/jquery-sparklines/2.1.2/jquery.sparkline.min.js"
)),
},
"seaweed50x50.png" => EmbeddedAsset {
content_type: "image/png",
bytes: include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/static/seaweed50x50.png"
)),
},
_ => return None,
};
Some(asset)
}
pub fn render_volume_server_html(state: &VolumeServerState) -> String {
let counters = server_stats::snapshot();
let (disk_rows, volume_rows, remote_volume_rows, ec_volume_rows) = {
let store = state.store.read().unwrap();
collect_ui_data(&store)
};
let masters = if state.master_urls.is_empty() {
"[]".to_string()
} else {
format!("[{}]", state.master_urls.join(" "))
};
let uptime = server_stats::uptime_string();
let read_week = join_i64(&counters.read_requests.week_counter.to_list());
let read_day = join_i64(&counters.read_requests.day_counter.to_list());
let read_hour = join_i64(&counters.read_requests.hour_counter.to_list());
let read_minute = join_i64(&counters.read_requests.minute_counter.to_list());
let mut disk_rows_html = String::new();
for disk in &disk_rows {
let _ = write!(
disk_rows_html,
"<tr><td>{}</td><td>{}</td><td>{}</td><td>{}</td><td>{:.2}%</td></tr>",
escape_html(&disk.dir),
escape_html(&disk.disk_type),
bytes_to_human_readable(disk.all),
bytes_to_human_readable(disk.free),
percent_from(disk.all, disk.used),
);
}
let mut volume_rows_html = String::new();
for volume in &volume_rows {
let _ = write!(
volume_rows_html,
"<tr><td><code>{}</code></td><td>{}</td><td>{}</td><td>{}</td><td>{}</td><td>{} / {}</td><td>{}</td><td>{}</td><td>{}</td></tr>",
volume.id,
escape_html(&volume.collection),
escape_html(&volume.disk_type),
bytes_to_human_readable(volume.size),
volume.file_count,
volume.delete_count,
bytes_to_human_readable(volume.deleted_byte_count),
escape_html(&volume.ttl),
volume.read_only,
volume.version,
);
}
let remote_section = if remote_volume_rows.is_empty() {
String::new()
} else {
let mut remote_rows_html = String::new();
for volume in &remote_volume_rows {
let _ = write!(
remote_rows_html,
"<tr><td><code>{}</code></td><td>{}</td><td>{}</td><td>{}</td><td>{} / {}</td><td>{}</td><td>{}</td></tr>",
volume.id,
escape_html(&volume.collection),
bytes_to_human_readable(volume.size),
volume.file_count,
volume.delete_count,
bytes_to_human_readable(volume.deleted_byte_count),
escape_html(&volume.remote_storage_name),
escape_html(&volume.remote_storage_key),
);
}
format!(
r#"<div class="row">
<h2>Remote Volumes</h2>
<table class="table table-striped">
<thead>
<tr>
<th>Id</th>
<th>Collection</th>
<th>Size</th>
<th>Files</th>
<th>Trash</th>
<th>Remote</th>
<th>Key</th>
</tr>
</thead>
<tbody>{}</tbody>
</table>
</div>"#,
remote_rows_html
)
};
let ec_section = if ec_volume_rows.is_empty() {
String::new()
} else {
let mut ec_rows_html = String::new();
for ec in &ec_volume_rows {
let mut shard_labels = String::new();
for shard in &ec.shards {
let _ = write!(
shard_labels,
"<span class=\"label label-info\" style=\"margin-right: 5px;\">{}: {}</span>",
shard.shard_id,
bytes_to_human_readable(shard.size)
);
}
let _ = write!(
ec_rows_html,
"<tr><td><code>{}</code></td><td>{}</td><td>{}</td><td>{}</td><td>{}</td></tr>",
ec.volume_id,
escape_html(&ec.collection),
bytes_to_human_readable(ec.size),
shard_labels,
escape_html(&ec.created_at),
);
}
format!(
r#"<div class="row">
<h2>Erasure Coding Shards</h2>
<table class="table table-striped">
<thead>
<tr>
<th>Id</th>
<th>Collection</th>
<th>Total Size</th>
<th>Shard Details</th>
<th>CreatedAt</th>
</tr>
</thead>
<tbody>{}</tbody>
</table>
</div>"#,
ec_rows_html
)
};
format!(
r#"<!DOCTYPE html>
<html>
<head>
<title>SeaweedFS {version}</title>
<link rel="stylesheet" href="/seaweedfsstatic/bootstrap/3.3.1/css/bootstrap.min.css">
<script type="text/javascript" src="/seaweedfsstatic/javascript/jquery-3.6.0.min.js"></script>
<script type="text/javascript" src="/seaweedfsstatic/javascript/jquery-sparklines/2.1.2/jquery.sparkline.min.js"></script>
<script type="text/javascript">
$(function () {{
var periods = ['second', 'minute', 'hour', 'day'];
for (var i = 0; i < periods.length; i++) {{
var period = periods[i];
$('.inlinesparkline-' + period).sparkline('html', {{
type: 'line',
barColor: 'red',
tooltipSuffix: ' request per ' + period
}});
}}
}});
</script>
<style>
#jqstooltip {{
height: 28px !important;
width: 150px !important;
}}
</style>
</head>
<body>
<div class="container">
<div class="page-header">
<h1>
<a href="https://github.com/seaweedfs/seaweedfs"><img src="/seaweedfsstatic/seaweed50x50.png" alt="SeaweedFS"></a>
SeaweedFS <small>{version}</small>
</h1>
</div>
<div class="row">
<div class="col-sm-6">
<h2>Disk Stats</h2>
<table class="table table-striped">
<thead>
<tr>
<th>Path</th>
<th>Disk</th>
<th>Total</th>
<th>Free</th>
<th>Usage</th>
</tr>
</thead>
<tbody>{disk_rows_html}</tbody>
</table>
</div>
<div class="col-sm-6">
<h2>System Stats</h2>
<table class="table table-condensed table-striped">
<tr><th>Masters</th><td>{masters}</td></tr>
<tr><th>Weekly # ReadRequests</th><td><span class="inlinesparkline-day">{read_week}</span></td></tr>
<tr><th>Daily # ReadRequests</th><td><span class="inlinesparkline-hour">{read_day}</span></td></tr>
<tr><th>Hourly # ReadRequests</th><td><span class="inlinesparkline-minute">{read_hour}</span></td></tr>
<tr><th>Last Minute # ReadRequests</th><td><span class="inlinesparkline-second">{read_minute}</span></td></tr>
<tr><th>Up Time</th><td>{uptime}</td></tr>
</table>
</div>
</div>
<div class="row">
<h2>Volumes</h2>
<table class="table table-striped">
<thead>
<tr>
<th>Id</th>
<th>Collection</th>
<th>Disk</th>
<th>Data Size</th>
<th>Files</th>
<th>Trash</th>
<th>TTL</th>
<th>ReadOnly</th>
<th>Version</th>
</tr>
</thead>
<tbody>{volume_rows_html}</tbody>
</table>
</div>
{remote_section}
{ec_section}
</div>
</body>
</html>"#,
version = escape_html(crate::version::version()),
disk_rows_html = disk_rows_html,
masters = escape_html(&masters),
read_week = read_week,
read_day = read_day,
read_hour = read_hour,
read_minute = read_minute,
uptime = escape_html(&uptime),
volume_rows_html = volume_rows_html,
remote_section = remote_section,
ec_section = ec_section,
)
}
fn collect_ui_data(
store: &Store,
) -> (
Vec<UiDiskRow>,
Vec<UiVolumeRow>,
Vec<UiVolumeRow>,
Vec<UiEcVolumeRow>,
) {
let mut disk_rows = Vec::new();
let mut volumes = Vec::new();
let mut remote_volumes = Vec::new();
let mut ec_volumes = Vec::new();
for loc in &store.locations {
let dir = absolute_display_path(&loc.directory);
let (all, free) = crate::storage::disk_location::get_disk_stats(&dir);
disk_rows.push(UiDiskRow {
dir,
disk_type: loc.disk_type.to_string(),
all,
free,
used: all.saturating_sub(free),
});
for (_, volume) in loc.volumes() {
let (remote_storage_name, remote_storage_key) = volume.remote_storage_name_key();
let row = UiVolumeRow {
id: volume.id.0,
collection: volume.collection.clone(),
disk_type: loc.disk_type.to_string(),
size: volume.content_size(),
file_count: volume.file_count(),
delete_count: volume.deleted_count(),
deleted_byte_count: volume.deleted_size(),
ttl: volume.super_block.ttl.to_string(),
read_only: volume.is_read_only(),
version: volume.version().0 as u32,
remote_storage_name,
remote_storage_key,
};
if row.remote_storage_name.is_empty() {
volumes.push(row);
} else {
remote_volumes.push(row);
}
}
for (_, ec_volume) in loc.ec_volumes() {
let mut shards = Vec::new();
let mut total_size = 0u64;
let mut created_at = String::from("-");
for shard in ec_volume.shards.iter().flatten() {
let shard_size = shard.file_size().max(0) as u64;
total_size = total_size.saturating_add(shard_size);
shards.push(UiEcShardRow {
shard_id: shard.shard_id,
size: shard_size,
});
if created_at == "-"
&& let Ok(metadata) = std::fs::metadata(shard.file_name())
&& let Ok(modified) = metadata.modified()
{
let ts: chrono::DateTime<chrono::Local> = modified.into();
created_at = ts.format("%Y-%m-%d %H:%M").to_string();
}
}
let preferred_size = ec_volume.dat_file_size.max(0) as u64;
ec_volumes.push(UiEcVolumeRow {
volume_id: ec_volume.volume_id.0,
collection: ec_volume.collection.clone(),
size: preferred_size.max(total_size),
shards,
created_at,
});
}
}
disk_rows.sort_by(|left, right| left.dir.cmp(&right.dir));
volumes.sort_by_key(|row| row.id);
remote_volumes.sort_by_key(|row| row.id);
ec_volumes.sort_by_key(|row| row.volume_id);
(disk_rows, volumes, remote_volumes, ec_volumes)
}
fn absolute_display_path(path: &str) -> String {
let p = std::path::Path::new(path);
if p.is_absolute() {
return path.to_string();
}
std::env::current_dir()
.map(|cwd| cwd.join(p).to_string_lossy().to_string())
.unwrap_or_else(|_| path.to_string())
}
fn join_i64(values: &[i64]) -> String {
values
.iter()
.map(std::string::ToString::to_string)
.collect::<Vec<_>>()
.join(",")
}
fn percent_from(total: u64, part: u64) -> f64 {
if total == 0 {
return 0.0;
}
(part as f64 / total as f64) * 100.0
}
fn bytes_to_human_readable(bytes: u64) -> String {
const UNIT: u64 = 1024;
if bytes < UNIT {
return format!("{} B", bytes);
}
let mut div = UNIT;
let mut exp = 0usize;
let mut n = bytes / UNIT;
while n >= UNIT {
div *= UNIT;
n /= UNIT;
exp += 1;
}
format!(
"{:.2} {}iB",
bytes as f64 / div as f64,
["K", "M", "G", "T", "P", "E"][exp]
)
}
fn escape_html(input: &str) -> String {
input
.replace('&', "&amp;")
.replace('<', "&lt;")
.replace('>', "&gt;")
.replace('"', "&quot;")
}

View File

@@ -0,0 +1,542 @@
//! VolumeServer: the main HTTP server for volume operations.
//!
//! Routes:
//! GET/HEAD /{vid},{fid} — read a file
//! POST/PUT /{vid},{fid} — write a file
//! DELETE /{vid},{fid} — delete a file
//! GET /status — server status
//! GET /healthz — health check
//!
//! Matches Go's server/volume_server.go.
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU32, Ordering};
use std::sync::{Arc, RwLock};
use axum::{
Router,
extract::{Request, State, connect_info::ConnectInfo},
http::{HeaderValue, Method, StatusCode, header},
middleware::{self, Next},
response::{IntoResponse, Response},
routing::{any, get},
};
use crate::config::ReadMode;
use crate::security::Guard;
use crate::storage::store::Store;
use super::grpc_client::OutgoingGrpcTlsConfig;
use super::handlers;
use super::write_queue::WriteQueue;
#[derive(Clone, Debug, Default)]
pub struct RuntimeMetricsConfig {
pub push_gateway: crate::metrics::PushGatewayConfig,
}
/// Shared state for the volume server.
pub struct VolumeServerState {
pub store: RwLock<Store>,
pub guard: RwLock<Guard>,
pub is_stopping: RwLock<bool>,
/// Maintenance mode flag.
pub maintenance: AtomicBool,
/// State version — incremented on each SetState call.
pub state_version: AtomicU32,
/// Throttling: concurrent upload/download limits (in bytes, 0 = disabled).
pub concurrent_upload_limit: i64,
pub concurrent_download_limit: i64,
pub inflight_upload_data_timeout: std::time::Duration,
pub inflight_download_data_timeout: std::time::Duration,
/// Current in-flight upload/download bytes.
pub inflight_upload_bytes: AtomicI64,
pub inflight_download_bytes: AtomicI64,
/// Notify waiters when inflight bytes decrease.
pub upload_notify: tokio::sync::Notify,
pub download_notify: tokio::sync::Notify,
/// Data center name from config.
pub data_center: String,
/// Rack name from config.
pub rack: String,
/// File size limit in bytes (0 = no limit).
pub file_size_limit_bytes: i64,
/// Default IO rate limit for maintenance copy/replication work.
pub maintenance_byte_per_second: i64,
/// Whether the server is connected to master (heartbeat active).
pub is_heartbeating: AtomicBool,
/// Whether master addresses are configured.
pub has_master: bool,
/// Seconds to wait before shutting down servers (graceful drain).
pub pre_stop_seconds: u32,
/// Notify heartbeat to send an immediate update when volume state changes.
pub volume_state_notify: tokio::sync::Notify,
/// Optional batched write queue for improved throughput under load.
pub write_queue: std::sync::OnceLock<WriteQueue>,
/// Registry of S3 tier backends for tiered storage operations.
pub s3_tier_registry: std::sync::RwLock<crate::remote_storage::s3_tier::S3TierRegistry>,
/// Read mode: local, proxy, or redirect for non-local volumes.
pub read_mode: ReadMode,
/// If true, FetchAndWriteNeedle skips remote S3 endpoint validation,
/// allowing arbitrary (incl. loopback / link-local / metadata) hosts.
pub allow_untrusted_remote_endpoints: bool,
/// First master address for volume lookups (e.g., "localhost:9333").
pub master_url: String,
/// Seed master addresses for UI rendering.
pub master_urls: Vec<String>,
/// Canonical http `host:port` form of every configured seed master.
/// Built once at construction so Ping admission stays O(1). Mirrors
/// Go's `seedMasterSet` on `VolumeServer`.
pub seed_master_set: std::collections::HashSet<String>,
/// Current master this server is heartbeating with, in canonical http
/// `host:port` form. Empty when no heartbeat connection is active. The
/// heartbeat goroutine writes; admission reads — the lock keeps them
/// from racing on a leader change. Mirrors Go's `currentMaster` plus
/// `currentMasterLock`.
pub current_master_url: tokio::sync::RwLock<String>,
/// This server's own address (ip:port) for filtering self from lookup results.
pub self_url: String,
/// HTTP client for proxy requests and master lookups.
pub http_client: reqwest::Client,
/// Scheme used for outgoing master and peer HTTP requests ("http" or "https").
pub outgoing_http_scheme: String,
/// Optional client TLS material for outgoing gRPC connections.
pub outgoing_grpc_tls: Option<OutgoingGrpcTlsConfig>,
/// Metrics push settings learned from master heartbeat responses.
pub metrics_runtime: std::sync::RwLock<RuntimeMetricsConfig>,
pub metrics_notify: tokio::sync::Notify,
/// Whether JPEG uploads should be normalized using EXIF orientation.
pub fix_jpg_orientation: bool,
/// Read tuning flags for large-file streaming.
pub has_slow_read: bool,
pub read_buffer_size_bytes: usize,
/// Path to security.toml — stored for SIGHUP reload.
pub security_file: String,
/// Original CLI whitelist entries — stored for SIGHUP reload.
pub cli_white_list: Vec<String>,
/// Path to state.pb file for persisting VolumeServerState across restarts.
pub state_file_path: String,
}
impl VolumeServerState {
/// Check if the server is in maintenance mode; return gRPC error if so.
pub fn check_maintenance(&self) -> Result<(), tonic::Status> {
if self.maintenance.load(Ordering::Relaxed) {
let id = self.store.read().unwrap().id.clone();
return Err(tonic::Status::unavailable(format!(
"volume server {} is in maintenance mode",
id
)));
}
Ok(())
}
/// Build the seed master set from a list of raw `host:port[.grpcPort]`
/// addresses, normalised the same way Go's `pb.ServerAddress.ToHttpAddress`
/// does (drop the `.grpcPort` suffix, preserve everything else).
pub fn build_seed_master_set(master_urls: &[String]) -> std::collections::HashSet<String> {
master_urls
.iter()
.map(|m| to_http_address(m).into_owned())
.collect()
}
/// Returns true iff `target` (normalised to canonical http `host:port`)
/// is a master this server already knows about. Volume servers do not
/// keep a peer-volume or peer-filer list, so Ping is scoped to masters.
/// Mirrors Go's `VolumeServer.isKnownPingTarget`.
pub async fn is_known_ping_target(&self, target: &str, target_type: &str) -> bool {
if target_type != "master" {
return false;
}
let key = to_http_address(target).into_owned();
if key.is_empty() {
return false;
}
let current = self.current_master_url.read().await.clone();
if !current.is_empty() && current == key {
return true;
}
self.seed_master_set.contains(&key)
}
}
pub fn build_metrics_router() -> Router {
Router::new().route("/metrics", get(handlers::metrics_handler))
}
pub fn normalize_outgoing_http_url(scheme: &str, raw_target: &str) -> Result<String, String> {
if raw_target.starts_with("http://") || raw_target.starts_with("https://") {
let mut url = reqwest::Url::parse(raw_target)
.map_err(|e| format!("invalid url {}: {}", raw_target, e))?;
url.set_scheme(scheme)
.map_err(|_| format!("invalid scheme {}", scheme))?;
return Ok(url.to_string());
}
Ok(format!("{}://{}", scheme, raw_target))
}
/// Convert a SeaweedFS server address to its HTTP `host:port` form.
///
/// Mirrors Go's `pb.ServerAddress.ToHttpAddress()`: SeaweedFS encodes a
/// server's gRPC port by appending `.grpcPort` to the HTTP port (e.g.
/// `host:9333.19333`). For HTTP requests we want only the HTTP `host:port`.
/// - `host:port.grpcPort` -> `host:port`
/// - `host:port` -> `host:port` (unchanged)
/// - Anything that does not look like `host:port[.grpcPort]` is returned unchanged.
///
/// Returns a `Cow<str>` so the common (no-suffix) case borrows from `addr`
/// without allocating; only the rewrite branch produces a new `String`.
pub fn to_http_address(addr: &str) -> std::borrow::Cow<'_, str> {
let Some(ports_sep_index) = addr.rfind(':') else {
return std::borrow::Cow::Borrowed(addr);
};
let ports = &addr[ports_sep_index + 1..];
if let Some(dot_idx) = ports.rfind('.') {
let http_port = &ports[..dot_idx];
let grpc_port = &ports[dot_idx + 1..];
// Only strip the suffix when both parts parse as real ports — leave
// anything else (e.g. "host:abc.def") untouched so bad config surfaces
// rather than being silently rewritten. Mirrors the validation already
// done in `to_grpc_address` for the inverse direction.
if let (Ok(_), Ok(_)) = (http_port.parse::<u16>(), grpc_port.parse::<u16>()) {
return std::borrow::Cow::Owned(addr[..ports_sep_index + 1 + dot_idx].to_string());
}
}
std::borrow::Cow::Borrowed(addr)
}
fn request_remote_addr(request: &Request) -> Option<SocketAddr> {
request
.extensions()
.get::<ConnectInfo<SocketAddr>>()
.map(|info| info.0)
}
fn request_is_whitelisted(state: &VolumeServerState, request: &Request) -> bool {
request_remote_addr(request)
.map(|remote_addr| {
state
.guard
.read()
.unwrap()
.check_whitelist(&remote_addr.to_string())
})
.unwrap_or(true)
}
/// Middleware: set Server header, echo x-amz-request-id, set CORS if Origin present.
async fn common_headers_middleware(request: Request, next: Next) -> Response {
let origin = request.headers().get("origin").cloned();
let request_id = super::request_id::generate_http_request_id();
let mut response =
super::request_id::scope_request_id(
request_id.clone(),
async move { next.run(request).await },
)
.await;
let headers = response.headers_mut();
if let Ok(val) = HeaderValue::from_str(crate::version::server_header()) {
headers.insert("Server", val);
}
if let Ok(val) = HeaderValue::from_str(&request_id) {
headers.insert("X-Request-Id", val.clone());
headers.insert("x-amz-request-id", val);
}
if origin.is_some() {
headers.insert("Access-Control-Allow-Origin", HeaderValue::from_static("*"));
headers.insert(
"Access-Control-Allow-Credentials",
HeaderValue::from_static("true"),
);
}
response
}
/// Admin store handler — dispatches based on HTTP method.
/// Matches Go's privateStoreHandler: GET/HEAD → read, POST/PUT → write,
/// DELETE → delete, OPTIONS → CORS headers, anything else → 400.
async fn admin_store_handler(state: State<Arc<VolumeServerState>>, request: Request) -> Response {
let start = std::time::Instant::now();
let method = request.method().clone();
let mut method_str = method.as_str().to_string();
let request_bytes = request
.headers()
.get(header::CONTENT_LENGTH)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<i64>().ok())
.filter(|value| *value > 0)
.unwrap_or(0);
super::server_stats::record_request_open();
crate::metrics::INFLIGHT_REQUESTS_GAUGE
.with_label_values(&[&method_str])
.inc();
let whitelist_rejected = matches!(method, Method::POST | Method::PUT | Method::DELETE)
&& !request_is_whitelisted(&state, &request);
let response = match method.clone() {
_ if whitelist_rejected => StatusCode::UNAUTHORIZED.into_response(),
Method::GET | Method::HEAD => {
super::server_stats::record_read_request();
handlers::get_or_head_handler_from_request(state, request).await
}
Method::POST | Method::PUT => {
super::server_stats::record_write_request();
if request_bytes > 0 {
super::server_stats::record_bytes_in(request_bytes);
}
handlers::post_handler(state, request).await
}
Method::DELETE => {
super::server_stats::record_delete_request();
handlers::delete_handler(state, request).await
}
Method::OPTIONS => {
super::server_stats::record_read_request();
admin_options_response()
}
_ => {
let method_name = request.method().to_string();
let query = request.uri().query().map(|q| q.to_string());
method_str = "INVALID".to_string();
handlers::json_error_with_query(
StatusCode::BAD_REQUEST,
format!("unsupported method {}", method_name),
query.as_deref(),
)
}
};
if method == Method::GET
&& let Some(response_bytes) = response
.headers()
.get(header::CONTENT_LENGTH)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<i64>().ok())
.filter(|value| *value > 0)
{
super::server_stats::record_bytes_out(response_bytes);
}
super::server_stats::record_request_close();
crate::metrics::INFLIGHT_REQUESTS_GAUGE
.with_label_values(&[&method_str])
.dec();
crate::metrics::REQUEST_COUNTER
.with_label_values(&[&method_str, response.status().as_str()])
.inc();
crate::metrics::REQUEST_DURATION
.with_label_values(&[&method_str])
.observe(start.elapsed().as_secs_f64());
response
}
/// Public store handler — dispatches based on HTTP method.
/// Matches Go's publicReadOnlyHandler: GET/HEAD → read, OPTIONS → CORS,
/// anything else → 200 (passthrough no-op).
async fn public_store_handler(state: State<Arc<VolumeServerState>>, request: Request) -> Response {
let start = std::time::Instant::now();
let method = request.method().clone();
let method_str = method.as_str().to_string();
super::server_stats::record_request_open();
crate::metrics::INFLIGHT_REQUESTS_GAUGE
.with_label_values(&[&method_str])
.inc();
let response = match method.clone() {
Method::GET | Method::HEAD => {
super::server_stats::record_read_request();
handlers::get_or_head_handler_from_request(state, request).await
}
Method::OPTIONS => {
super::server_stats::record_read_request();
public_options_response()
}
_ => StatusCode::OK.into_response(),
};
if method == Method::GET
&& let Some(response_bytes) = response
.headers()
.get(header::CONTENT_LENGTH)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<i64>().ok())
.filter(|value| *value > 0)
{
super::server_stats::record_bytes_out(response_bytes);
}
super::server_stats::record_request_close();
crate::metrics::INFLIGHT_REQUESTS_GAUGE
.with_label_values(&[&method_str])
.dec();
crate::metrics::REQUEST_COUNTER
.with_label_values(&[&method_str, response.status().as_str()])
.inc();
crate::metrics::REQUEST_DURATION
.with_label_values(&[&method_str])
.observe(start.elapsed().as_secs_f64());
response
}
/// Build OPTIONS response for admin port.
fn admin_options_response() -> Response {
let mut response = StatusCode::OK.into_response();
let headers = response.headers_mut();
headers.insert(
"Access-Control-Allow-Methods",
HeaderValue::from_static("PUT, POST, GET, DELETE, OPTIONS"),
);
headers.insert(
"Access-Control-Allow-Headers",
HeaderValue::from_static("*"),
);
response
}
/// Build OPTIONS response for public port.
fn public_options_response() -> Response {
let mut response = StatusCode::OK.into_response();
let headers = response.headers_mut();
headers.insert(
"Access-Control-Allow-Methods",
HeaderValue::from_static("GET, OPTIONS"),
);
headers.insert(
"Access-Control-Allow-Headers",
HeaderValue::from_static("*"),
);
response
}
/// Build the admin (private) HTTP router — supports all operations.
/// UI route is only registered when no signing keys are configured,
/// matching Go's `if signingKey == "" || enableUiAccess` check.
pub fn build_admin_router(state: Arc<VolumeServerState>) -> Router {
let guard = state.guard.read().unwrap();
// This helper can only derive the default Go behavior from the guard state:
// UI stays enabled when the write signing key is empty. The explicit
// `access.ui` override is handled by `build_admin_router_with_ui(...)`.
let ui_enabled = guard.signing_key.0.is_empty();
drop(guard);
build_admin_router_with_ui(state, ui_enabled)
}
/// Build the admin router with an explicit UI exposure flag.
pub fn build_admin_router_with_ui(state: Arc<VolumeServerState>, ui_enabled: bool) -> Router {
let mut router = Router::new()
.route("/status", get(handlers::status_handler))
.route("/healthz", get(handlers::healthz_handler))
.route("/favicon.ico", get(handlers::favicon_handler))
.route(
"/seaweedfsstatic/{*path}",
get(handlers::static_asset_handler),
)
.route("/", any(admin_store_handler))
.route("/{path}", any(admin_store_handler))
.route("/{vid}/{fid}", any(admin_store_handler))
.route("/{vid}/{fid}/{filename}", any(admin_store_handler))
.fallback(admin_store_handler);
if ui_enabled {
// Note: /stats/* endpoints are commented out in Go's volume_server.go (L130-134).
// Only the UI endpoint is registered when UI access is enabled.
router = router.route("/ui/index.html", get(handlers::ui_handler));
}
router
.layer(middleware::from_fn(common_headers_middleware))
.with_state(state)
}
/// Build the public (read-only) HTTP router — only GET/HEAD.
pub fn build_public_router(state: Arc<VolumeServerState>) -> Router {
Router::new()
.route("/favicon.ico", get(handlers::favicon_handler))
.route(
"/seaweedfsstatic/{*path}",
get(handlers::static_asset_handler),
)
.route("/", any(public_store_handler))
.route("/{path}", any(public_store_handler))
.route("/{vid}/{fid}", any(public_store_handler))
.route("/{vid}/{fid}/{filename}", any(public_store_handler))
.fallback(public_store_handler)
.layer(middleware::from_fn(common_headers_middleware))
.with_state(state)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_to_http_address_strips_grpc_port_suffix() {
assert_eq!(to_http_address("10.0.0.1:9333.19333"), "10.0.0.1:9333");
assert_eq!(
to_http_address("master.local:9333.19333"),
"master.local:9333"
);
assert_eq!(to_http_address("10.85.183.6:5300.6300"), "10.85.183.6:5300");
}
#[test]
fn test_to_http_address_passthrough_without_grpc_suffix() {
assert_eq!(to_http_address("10.0.0.1:9333"), "10.0.0.1:9333");
assert_eq!(to_http_address("master.local:9333"), "master.local:9333");
}
#[test]
fn test_to_http_address_returns_input_when_unparseable() {
assert_eq!(to_http_address(""), "");
assert_eq!(to_http_address("no-port"), "no-port");
// Trailing colon: nothing after the separator, treat as unparseable.
assert_eq!(to_http_address("host:"), "host:");
}
#[test]
fn test_to_http_address_borrows_when_unchanged_and_owns_when_stripped() {
// The common case (no suffix) must not allocate.
let result = to_http_address("10.0.0.1:9333");
assert!(matches!(result, std::borrow::Cow::Borrowed(_)));
// Stripping requires a new string.
let result = to_http_address("10.0.0.1:9333.19333");
assert!(matches!(result, std::borrow::Cow::Owned(_)));
assert_eq!(result, "10.0.0.1:9333");
// Unparseable / passthrough also borrows.
let result = to_http_address("host:abc.def");
assert!(matches!(result, std::borrow::Cow::Borrowed(_)));
}
#[test]
fn test_to_http_address_keeps_non_numeric_dotted_suffix() {
// The dotted form is only valid when both sides are real port numbers.
// Otherwise the address is malformed config (e.g. a hostname like
// "host:abc.def"), and silently rewriting it would just hide the bug.
assert_eq!(to_http_address("host:abc.def"), "host:abc.def");
assert_eq!(to_http_address("host:9333.notaport"), "host:9333.notaport");
assert_eq!(
to_http_address("host:notaport.19333"),
"host:notaport.19333"
);
// Out-of-range ports must not be silently truncated either.
assert_eq!(to_http_address("host:99999.19333"), "host:99999.19333");
}
#[test]
fn test_to_http_address_handles_bracketed_ipv6_literals() {
// The function uses `rfind(':')`, so for bracketed IPv6 the port
// separator is correctly identified as the colon AFTER the closing
// bracket — making IPv4 and IPv6 behave the same.
assert_eq!(to_http_address("[::1]:9333.19333"), "[::1]:9333");
assert_eq!(
to_http_address("[2001:db8::10]:5300.6300"),
"[2001:db8::10]:5300"
);
// Plain bracketed IPv6 without a dotted suffix is borrowed unchanged.
let result = to_http_address("[2001:db8::1]:9333");
assert!(matches!(result, std::borrow::Cow::Borrowed(_)));
assert_eq!(result, "[2001:db8::1]:9333");
// Non-numeric / out-of-range / missing suffix all preserve the input.
assert_eq!(to_http_address("[::1]:"), "[::1]:");
assert_eq!(to_http_address("[::1]:abc.def"), "[::1]:abc.def");
assert_eq!(to_http_address("[::1]:99999.19333"), "[::1]:99999.19333");
}
}

View File

@@ -0,0 +1,339 @@
//! Async batched write processing for the volume server.
//!
//! Instead of each upload handler directly calling `write_needle`, writes are
//! submitted to a queue. A background worker drains the queue in batches (up to
//! 128 entries), groups them by volume ID, and processes them together under a
//! single store lock. Requests that asked for `fsync` are flushed by
//! `write_needle` itself, one flush per durable write.
use std::sync::Arc;
use tokio::sync::{mpsc, oneshot};
use tracing::debug;
use crate::storage::needle::needle::Needle;
use crate::storage::types::{Size, VolumeId};
use crate::storage::volume::VolumeError;
use super::volume_server::VolumeServerState;
/// Result of a single write operation: (offset, size, is_unchanged).
pub type WriteResult = Result<(u64, Size, bool), VolumeError>;
/// The needles queued for one volume, each with the durability it asked for.
type VolumeBatch = Vec<(Needle, bool, oneshot::Sender<WriteResult>)>;
/// A request to write a needle, submitted to the write queue.
pub struct WriteRequest {
pub volume_id: VolumeId,
pub needle: Needle,
pub fsync: bool,
pub response_tx: oneshot::Sender<WriteResult>,
}
/// Maximum number of write requests to batch together.
const MAX_BATCH_SIZE: usize = 128;
/// Maximum bytes to accumulate per batch before breaking (matches Go's 4MB limit).
/// This prevents large writes from accumulating unbounded latency.
const MAX_BATCH_BYTES: usize = 4 * 1024 * 1024;
/// Handle for submitting write requests to the background worker.
#[derive(Clone)]
pub struct WriteQueue {
tx: mpsc::Sender<WriteRequest>,
}
impl WriteQueue {
/// Create a new write queue and spawn the background worker.
///
/// `capacity` controls the channel buffer size (backpressure kicks in when full).
/// The worker holds a reference to `state` for accessing the store.
pub fn new(state: Arc<VolumeServerState>, capacity: usize) -> Self {
let (tx, rx) = mpsc::channel(capacity);
let worker = WriteQueueWorker { rx, state };
tokio::spawn(worker.run());
WriteQueue { tx }
}
/// Submit a write request and wait for the result.
///
/// Returns `Err` if the worker has shut down or the response channel was dropped.
pub async fn submit(&self, volume_id: VolumeId, needle: Needle, fsync: bool) -> WriteResult {
let (response_tx, response_rx) = oneshot::channel();
let request = WriteRequest {
volume_id,
needle,
fsync,
response_tx,
};
// Send to queue; this awaits if the channel is full (backpressure).
if self.tx.send(request).await.is_err() {
return Err(VolumeError::Io(std::io::Error::new(
std::io::ErrorKind::BrokenPipe,
"write queue worker has shut down",
)));
}
// Wait for the worker to process our request.
match response_rx.await {
Ok(result) => result,
Err(_) => Err(VolumeError::Io(std::io::Error::new(
std::io::ErrorKind::BrokenPipe,
"write queue worker dropped response channel",
))),
}
}
}
/// Background worker that drains write requests and processes them in batches.
struct WriteQueueWorker {
rx: mpsc::Receiver<WriteRequest>,
state: Arc<VolumeServerState>,
}
impl WriteQueueWorker {
async fn run(mut self) {
debug!("write queue worker started");
loop {
// Wait for the first request (blocks until one arrives or channel closes).
let first = match self.rx.recv().await {
Some(req) => req,
None => {
debug!("write queue channel closed, worker exiting");
return;
}
};
// Drain as many additional requests as available, up to MAX_BATCH_SIZE
// or MAX_BATCH_BYTES (matches Go: 128 requests or 4MB, whichever comes first).
let mut batch = Vec::with_capacity(MAX_BATCH_SIZE);
let mut batch_bytes: usize = first.needle.data.len();
batch.push(first);
while batch.len() < MAX_BATCH_SIZE && batch_bytes < MAX_BATCH_BYTES {
match self.rx.try_recv() {
Ok(req) => {
batch_bytes += req.needle.data.len();
batch.push(req);
}
Err(_) => break,
}
}
let batch_size = batch.len();
debug!("processing write batch of {} requests", batch_size);
// Process the batch in spawn_blocking since write_needle does file I/O.
let state = self.state.clone();
let _ = tokio::task::spawn_blocking(move || {
process_batch(state, batch);
})
.await;
}
}
}
/// Process a batch of write requests, grouped by volume ID.
///
/// Groups writes by volume to minimize the number of store lock acquisitions,
/// then sends results back via each request's oneshot channel.
fn process_batch(state: Arc<VolumeServerState>, batch: Vec<WriteRequest>) {
// Group requests by volume ID for efficient processing.
// We use a Vec of (VolumeId, Vec<(Needle, Sender)>) to preserve order
// and avoid requiring Hash on VolumeId.
let mut groups: Vec<(VolumeId, VolumeBatch)> = Vec::new();
for req in batch {
let vid = req.volume_id;
if let Some(group) = groups.iter_mut().find(|(v, _)| *v == vid) {
group.1.push((req.needle, req.fsync, req.response_tx));
} else {
groups.push((vid, vec![(req.needle, req.fsync, req.response_tx)]));
}
}
// Process each volume group under a single store lock.
let mut store = state.store.write().unwrap();
for (vid, entries) in groups {
for (mut needle, fsync, response_tx) in entries {
let result = store.write_volume_needle(vid, &mut needle, fsync);
// Send result back; ignore error if receiver dropped.
let _ = response_tx.send(result);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::types::VolumeId;
/// Helper to create a minimal VolumeServerState for testing.
fn make_test_state() -> Arc<VolumeServerState> {
use crate::security::{Guard, SigningKey};
use crate::server::volume_server::RuntimeMetricsConfig;
use crate::storage::needle_map::NeedleMapKind;
use crate::storage::store::Store;
use std::sync::RwLock;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU32};
let store = Store::new(NeedleMapKind::InMemory);
let guard = Guard::new(&[], SigningKey(vec![]), 0, SigningKey(vec![]), 0);
Arc::new(VolumeServerState {
store: RwLock::new(store),
guard: RwLock::new(guard),
is_stopping: RwLock::new(false),
maintenance: AtomicBool::new(false),
state_version: AtomicU32::new(0),
concurrent_upload_limit: 0,
concurrent_download_limit: 0,
inflight_upload_data_timeout: std::time::Duration::ZERO,
inflight_download_data_timeout: std::time::Duration::ZERO,
inflight_upload_bytes: AtomicI64::new(0),
inflight_download_bytes: AtomicI64::new(0),
upload_notify: tokio::sync::Notify::new(),
download_notify: tokio::sync::Notify::new(),
data_center: String::new(),
rack: String::new(),
file_size_limit_bytes: 0,
maintenance_byte_per_second: 0,
is_heartbeating: AtomicBool::new(false),
has_master: false,
pre_stop_seconds: 0,
volume_state_notify: tokio::sync::Notify::new(),
write_queue: std::sync::OnceLock::new(),
s3_tier_registry: std::sync::RwLock::new(
crate::remote_storage::s3_tier::S3TierRegistry::new(),
),
read_mode: crate::config::ReadMode::Local,
allow_untrusted_remote_endpoints: false,
master_url: String::new(),
master_urls: Vec::new(),
seed_master_set: std::collections::HashSet::new(),
current_master_url: tokio::sync::RwLock::new(String::new()),
self_url: String::new(),
http_client: reqwest::Client::new(),
outgoing_http_scheme: "http".to_string(),
outgoing_grpc_tls: None,
metrics_runtime: std::sync::RwLock::new(RuntimeMetricsConfig::default()),
metrics_notify: tokio::sync::Notify::new(),
fix_jpg_orientation: false,
has_slow_read: true,
read_buffer_size_bytes: 4 * 1024 * 1024,
security_file: String::new(),
cli_white_list: vec![],
state_file_path: String::new(),
})
}
#[tokio::test]
async fn test_write_queue_submit_no_volume() {
// Submit a write to a non-existent volume -- should return VolumeError::NotFound.
let state = make_test_state();
let queue = WriteQueue::new(state, MAX_BATCH_SIZE);
let needle = Needle {
id: 1.into(),
cookie: 0x12345678.into(),
data: vec![1, 2, 3],
data_size: 3,
..Needle::default()
};
let result = queue.submit(VolumeId(999), needle, false).await;
assert!(result.is_err());
match result {
Err(VolumeError::NotFound) => {} // expected
other => panic!("expected NotFound, got {:?}", other),
}
}
#[tokio::test]
async fn test_write_queue_concurrent_submissions() {
// Submit multiple concurrent writes -- all should complete (with errors since no volume).
let state = make_test_state();
let queue = WriteQueue::new(state, MAX_BATCH_SIZE);
let mut handles = Vec::new();
for i in 0..10u64 {
let q = queue.clone();
handles.push(tokio::spawn(async move {
let needle = Needle {
id: i.into(),
cookie: 0xABCD.into(),
data: vec![i as u8; 10],
data_size: 10,
..Needle::default()
};
q.submit(VolumeId(1), needle, false).await
}));
}
for handle in handles {
let result = handle.await.unwrap();
// All should fail with NotFound since there's no volume 1
assert!(matches!(result, Err(VolumeError::NotFound)));
}
}
#[tokio::test]
async fn test_write_queue_batching() {
// Verify that many concurrent writes get processed (testing the batching path).
let state = make_test_state();
let queue = WriteQueue::new(state, MAX_BATCH_SIZE);
// Submit MAX_BATCH_SIZE requests concurrently
let mut handles = Vec::new();
for i in 0..MAX_BATCH_SIZE as u64 {
let q = queue.clone();
handles.push(tokio::spawn(async move {
let needle = Needle {
id: i.into(),
cookie: 0x1111.into(),
data: vec![0u8; 4],
data_size: 4,
..Needle::default()
};
q.submit(VolumeId(42), needle, false).await
}));
}
let mut results = Vec::new();
for handle in handles {
results.push(handle.await.unwrap());
}
// All should complete (with NotFound errors since no volume exists)
assert_eq!(results.len(), MAX_BATCH_SIZE);
for r in results {
assert!(matches!(r, Err(VolumeError::NotFound)));
}
}
#[tokio::test]
async fn test_write_queue_dropped_sender() {
// When the queue is dropped, subsequent submits should fail gracefully.
let state = make_test_state();
let queue = WriteQueue::new(state, 1);
// Clone then drop the original -- the worker keeps running via its rx handle.
let queue2 = queue.clone();
drop(queue);
// This should still work since the worker is alive.
let needle = Needle {
id: 1.into(),
cookie: 0.into(),
data: vec![],
data_size: 0,
..Needle::default()
};
let result = queue2.submit(VolumeId(1), needle, false).await;
assert!(result.is_err()); // NotFound is fine -- the point is it doesn't panic
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,961 @@
//! EC bitrot detection — checksum sidecar.
//!
//! A per-volume sidecar file stores a CRC32C (Castagnoli) checksum for every
//! fixed-size block of every EC shard, so a scrub (and the reconstruction path)
//! can detect silent disk corruption in any shard — including cold parity shards
//! that are never read during normal serving.
//!
//! The sidecar is OPTIONAL: an absent or generation-mismatched sidecar simply
//! means "feature off" for that generation, so old binaries, JSON-only nodes,
//! and rollback deployments ignore it and degrade gracefully.
//!
//! On-disk layout of `<base>.ecsum` (legacy/generation 0) and `<base>.ecsum.v<N>`:
//!
//! ```text
//! [ magic(4) | format_version(2) | payload_len(4) | payload_crc32c(4) ] [ proto payload ]
//! ```
//!
//! All header fields are BIG-ENDIAN. The header's `payload_crc32c` lets a loader
//! detect corruption of the sidecar itself BEFORE trusting any contents, so a
//! rotted sidecar can never be mistaken for shard corruption. This format is
//! byte-identical to the Go implementation in
//! `weed/storage/erasure_coding/ec_bitrot.go`.
use std::fs::File;
use std::io::{self, Read, Write};
use prost::Message;
use crate::pb::volume_server_pb::{
ChecksumAlgorithm, EcBitrotProtection, EcShardChecksums, EcShardConfig,
};
use crate::storage::erasure_coding::ec_shard::MAX_SHARD_COUNT;
use crate::storage::needle::crc::CRC;
/// Canonical extension for the checksum sidecar. Generation 0 (legacy/fresh
/// encode) uses `<base>.ecsum`; vacuum generation N uses `<base>.ecsum.v<N>`,
/// mirroring the `.vif`/`.ecx` versioned convention.
pub const BITROT_SIDECAR_EXT: &str = ".ecsum";
/// Default checksum granularity (16 MiB). It is a power-of-two multiple of
/// `ERASURE_CODING_SMALL_BLOCK_SIZE` (1 MiB) and keeps the sidecar tiny
/// (~11 KB for a 30 GB volume) while localizing corruption to a 16 MiB region.
pub const DEFAULT_BITROT_BLOCK_SIZE: usize = 16 * 1024 * 1024;
/// Caps the block granularity so a loaded sidecar cannot force a huge
/// scrub/verify scratch buffer. Power-of-two multiple of 1 MiB.
pub const MAX_BITROT_BLOCK_SIZE: u32 = 64 * 1024 * 1024;
/// Magic "ECSU".
pub const BITROT_MAGIC: u32 = 0x4543_5355;
/// On-disk format version.
pub const BITROT_FORMAT_VERSION: u16 = 1;
/// Header size: magic(4) + version(2) + payload_len(4) + payload_crc32c(4).
pub const BITROT_HEADER_SIZE: usize = 14;
/// Resolved protection state of an EC volume's active generation after loading
/// and validating its sidecar.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BitrotStatus {
/// No sidecar, or a sidecar that does not describe the active generation.
/// The generation is unprotected; this is NOT corruption.
Off,
/// A complete, well-formed, generation-matching sidecar is loaded.
On,
/// A generation-matching sidecar that is malformed, incomplete, or
/// self-integrity-failed. The generation is unprotected pending repair of
/// the sidecar, and an integrity alarm should fire. The rebuild path treats
/// this as fail-closed.
Invalid,
}
impl std::fmt::Display for BitrotStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let s = match self {
BitrotStatus::On => "on",
BitrotStatus::Invalid => "invalid",
BitrotStatus::Off => "off",
};
f.write_str(s)
}
}
/// Error returned by [`load_bitrot_sidecar`]. A self-integrity failure is a
/// sidecar-integrity problem (the caller maps it to [`BitrotStatus::Invalid`]),
/// never a shard-corruption signal.
#[derive(Debug)]
pub enum BitrotLoadError {
/// File missing or other underlying I/O error.
Io(io::Error),
/// File shorter than the fixed header.
TooShort(usize),
/// Header magic did not match.
BadMagic(u32),
/// Header format version is unsupported.
UnsupportedVersion(u16),
/// Header payload_len disagrees with the actual payload length.
LengthMismatch { header: u32, actual: usize },
/// Header payload_crc32c disagrees with the computed CRC32C of the payload.
CrcMismatch { header: u32, computed: u32 },
/// Protobuf payload failed to decode.
Decode(prost::DecodeError),
}
impl std::fmt::Display for BitrotLoadError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BitrotLoadError::Io(e) => write!(f, "bitrot sidecar io error: {}", e),
BitrotLoadError::TooShort(n) => {
write!(f, "bitrot sidecar too short ({} bytes)", n)
}
BitrotLoadError::BadMagic(m) => write!(f, "bitrot sidecar bad magic {:#x}", m),
BitrotLoadError::UnsupportedVersion(v) => {
write!(f, "bitrot sidecar unsupported format version {}", v)
}
BitrotLoadError::LengthMismatch { header, actual } => write!(
f,
"bitrot sidecar length mismatch: header {}, actual {}",
header, actual
),
BitrotLoadError::CrcMismatch { header, computed } => write!(
f,
"bitrot sidecar self-integrity CRC mismatch: header {:#x}, computed {:#x}",
header, computed
),
BitrotLoadError::Decode(e) => write!(f, "unmarshal bitrot sidecar: {}", e),
}
}
}
impl std::error::Error for BitrotLoadError {}
impl From<io::Error> for BitrotLoadError {
fn from(e: io::Error) -> Self {
BitrotLoadError::Io(e)
}
}
/// Returns the sidecar path for a base file name and EC generation. Generation
/// 0 is the un-suffixed legacy path; generation N>0 is the versioned path,
/// consistent with how `.vif`/`.ecx` are versioned by the 2PC switch.
pub fn bitrot_sidecar_path(base: &str, generation: u32) -> String {
if generation == 0 {
format!("{}{}", base, BITROT_SIDECAR_EXT)
} else {
format!("{}{}.v{}", base, BITROT_SIDECAR_EXT, generation)
}
}
/// Remove the bitrot checksum sidecars for a base file name: the legacy
/// `<base>.ecsum` (generation 0) and any versioned `<base>.ecsum.v<N>`.
///
/// Already-gone is success; returns the first real removal failure (and surfaces
/// a directory-scan error) so a stale sidecar left behind is not reported as
/// cleaned. Mirrors Go's `RemoveBitrotSidecars` / `removeBitrotSidecars`.
pub fn remove_bitrot_sidecars(base: &str) -> io::Result<()> {
use std::fs;
use std::path::Path;
let rm = |path: std::path::PathBuf| -> io::Result<()> {
match fs::remove_file(&path) {
Err(e) if e.kind() != io::ErrorKind::NotFound => Err(e),
_ => Ok(()),
}
};
let mut first_err: Option<io::Error> = None;
let mut record = |res: io::Result<()>| {
if let Err(e) = res
&& first_err.is_none()
{
first_err = Some(e);
}
};
record(rm(format!("{}{}", base, BITROT_SIDECAR_EXT).into()));
let path = Path::new(base);
if let (Some(parent), Some(fname)) = (path.parent(), path.file_name()) {
let prefix = format!("{}{}.v", fname.display(), BITROT_SIDECAR_EXT);
match fs::read_dir(parent) {
Ok(entries) => {
for entry in entries.flatten() {
if entry.file_name().to_string_lossy().starts_with(&prefix) {
record(rm(entry.path()));
}
}
}
Err(e) if e.kind() != io::ErrorKind::NotFound => record(Err(e)),
Err(_) => {}
}
}
match first_err {
Some(e) => Err(e),
None => Ok(()),
}
}
/// Returns a fresh random per-encode identity used to detect a stale sidecar
/// left behind by an in-place re-encode.
pub fn new_encode_uuid() -> Vec<u8> {
use rand::Rng;
let mut b = vec![0u8; 16];
rand::rng().fill_bytes(&mut b);
b
}
/// Reports whether `block_size` is a power of two in [1 MiB, MAX_BITROT_BLOCK_SIZE].
pub fn is_pow2_multiple_of_1mib(block_size: u32) -> bool {
((1 << 20)..=MAX_BITROT_BLOCK_SIZE).contains(&block_size) && block_size.count_ones() == 1
}
/// Returns ceil(covered_size / block_size).
fn expected_block_count(covered_size: i64, block_size: i64) -> usize {
if block_size <= 0 {
return 0;
}
((covered_size + block_size - 1) / block_size) as usize
}
/// Packs a slice of u32 into little-endian bytes.
fn pack_u32_le(vals: &[u32]) -> Vec<u8> {
let mut out = Vec::with_capacity(vals.len() * 4);
for v in vals {
out.extend_from_slice(&v.to_le_bytes());
}
out
}
/// Unpacks little-endian bytes into a Vec<u32>. Trailing bytes (len % 4) are
/// ignored, matching the Go implementation's `len(b)/4` truncation.
fn unpack_u32_le(b: &[u8]) -> Vec<u32> {
let n = b.len() / 4;
let mut out = Vec::with_capacity(n);
for i in 0..n {
out.push(u32::from_le_bytes([
b[i * 4],
b[i * 4 + 1],
b[i * 4 + 2],
b[i * 4 + 3],
]));
}
out
}
/// Accumulates the per-block CRC32C of a single shard's byte stream as it is
/// written. Tolerates arbitrary chunk sizes that cross block boundaries, so it
/// works for both small encode buffers and the larger rebuild buffers.
pub struct ShardChecksumBuilder {
block_size: i64,
cur: CRC,
cur_len: i64,
total: i64,
blocks: Vec<u32>,
}
impl ShardChecksumBuilder {
/// Creates a builder over a given block size in bytes.
pub fn new(block_size: i64) -> Self {
ShardChecksumBuilder {
block_size,
cur: CRC(0),
cur_len: 0,
total: 0,
blocks: Vec::new(),
}
}
/// Feeds a chunk of arbitrary size, splitting it across block boundaries.
pub fn write(&mut self, mut p: &[u8]) {
while !p.is_empty() {
let room = self.block_size - self.cur_len;
let n = (p.len() as i64).min(room) as usize;
self.cur = self.cur.update(&p[..n]);
self.cur_len += n as i64;
self.total += n as i64;
p = &p[n..];
if self.cur_len == self.block_size {
self.blocks.push(self.cur.0);
self.cur = CRC(0);
self.cur_len = 0;
}
}
}
/// Flushes any partial last block and returns the covered size and the
/// packed little-endian u32 CRC array.
pub fn finalize(mut self) -> (i64, Vec<u8>) {
if self.cur_len > 0 {
self.blocks.push(self.cur.0);
self.cur = CRC(0);
self.cur_len = 0;
}
(self.total, pack_u32_le(&self.blocks))
}
}
/// Atomically writes `prot` to `path`, wrapped in the on-disk header with a
/// CRC32C over the serialized payload (temp file + rename).
pub fn save_bitrot_sidecar(path: &str, prot: &EcBitrotProtection) -> io::Result<()> {
// The header records payload_len as a uint32 and the buffer allocation below
// adds it to a constant. Bound the payload well under any overflow (a real
// manifest is a few KB) so neither the length field nor the buffer can wrap.
// Check the encoded length BEFORE serializing so an oversized manifest never
// allocates a huge buffer. Mirrors Go's SaveBitrotSidecar maxBitrotPayloadSize.
const MAX_BITROT_PAYLOAD_SIZE: usize = 1 << 30; // 1 GiB, vastly above any real sidecar
let payload_len = prot.encoded_len();
if payload_len > MAX_BITROT_PAYLOAD_SIZE {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("bitrot sidecar payload too large: {} bytes", payload_len),
));
}
let payload = prot.encode_to_vec();
let mut buf = Vec::with_capacity(BITROT_HEADER_SIZE + payload.len());
buf.extend_from_slice(&BITROT_MAGIC.to_be_bytes());
buf.extend_from_slice(&BITROT_FORMAT_VERSION.to_be_bytes());
buf.extend_from_slice(&(payload.len() as u32).to_be_bytes());
buf.extend_from_slice(&CRC::new(&payload).0.to_be_bytes());
buf.extend_from_slice(&payload);
let tmp = format!("{}.tmp", path);
{
let mut f = File::create(&tmp)?;
f.write_all(&buf)?;
f.sync_all()?;
}
if let Err(e) = std::fs::rename(&tmp, path) {
let _ = std::fs::remove_file(&tmp);
return Err(e);
}
Ok(())
}
/// Reads and self-integrity-checks a sidecar file. Returns the parsed message,
/// or an error if the file is missing, truncated, has a bad magic/version, or
/// fails the payload CRC. A self-integrity failure is a sidecar-integrity
/// problem (caller maps it to [`BitrotStatus::Invalid`]), never a shard
/// corruption signal.
pub fn load_bitrot_sidecar(path: &str) -> Result<EcBitrotProtection, BitrotLoadError> {
let mut data = Vec::new();
File::open(path)?.read_to_end(&mut data)?;
if data.len() < BITROT_HEADER_SIZE {
return Err(BitrotLoadError::TooShort(data.len()));
}
let magic = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
if magic != BITROT_MAGIC {
return Err(BitrotLoadError::BadMagic(magic));
}
let ver = u16::from_be_bytes([data[4], data[5]]);
if ver != BITROT_FORMAT_VERSION {
return Err(BitrotLoadError::UnsupportedVersion(ver));
}
let payload_len = u32::from_be_bytes([data[6], data[7], data[8], data[9]]);
let want_crc = u32::from_be_bytes([data[10], data[11], data[12], data[13]]);
let payload = &data[BITROT_HEADER_SIZE..];
if payload_len as usize != payload.len() {
return Err(BitrotLoadError::LengthMismatch {
header: payload_len,
actual: payload.len(),
});
}
let got = CRC::new(payload).0;
if got != want_crc {
return Err(BitrotLoadError::CrcMismatch {
header: want_crc,
computed: got,
});
}
EcBitrotProtection::decode(payload).map_err(BitrotLoadError::Decode)
}
/// Performs the disk-free manifest/syntax checks that every loader runs:
/// supported algorithm, valid block size, exactly one entry per shard id in the
/// active layout (no duplicates, no out-of-range ids), positive covered_size,
/// and a packed-CRC count consistent with covered_size.
///
/// It does NOT compare covered_size against on-disk shard lengths — that is a
/// per-node physical check done only for locally-held shards.
pub fn validate_manifest(
prot: &EcBitrotProtection,
data_shards: usize,
parity_shards: usize,
) -> Result<(), String> {
if prot.algorithm != ChecksumAlgorithm::ChecksumCrc32c as i32 {
return Err(format!("unsupported checksum algorithm {}", prot.algorithm));
}
if !is_pow2_multiple_of_1mib(prot.block_size) {
return Err(format!(
"invalid block_size {} (must be a power-of-two multiple of 1 MiB, at most {})",
prot.block_size, MAX_BITROT_BLOCK_SIZE
));
}
let bs = prot.block_size as i64;
let total = data_shards + parity_shards;
if total == 0 || total > MAX_SHARD_COUNT {
return Err(format!(
"invalid active layout: data={} parity={}",
data_shards, parity_shards
));
}
if prot.shards.len() != total {
return Err(format!(
"incomplete manifest: {} shard entries, expected {}",
prot.shards.len(),
total
));
}
let mut seen = [false; MAX_SHARD_COUNT];
for s in &prot.shards {
if s.shard_id >= total as u32 {
return Err(format!(
"shard id {} out of range [0,{})",
s.shard_id, total
));
}
if seen[s.shard_id as usize] {
return Err(format!("duplicate shard id {}", s.shard_id));
}
seen[s.shard_id as usize] = true;
if s.covered_size <= 0 {
return Err(format!(
"shard {} has non-positive covered_size {}",
s.shard_id, s.covered_size
));
}
let want_count = expected_block_count(s.covered_size, bs);
if s.block_crc32c.len() != want_count * 4 {
return Err(format!(
"shard {} crc count mismatch: {} bytes, expected {} (covered_size={} block_size={})",
s.shard_id,
s.block_crc32c.len(),
want_count * 4,
s.covered_size,
prot.block_size
));
}
}
Ok(())
}
/// Resolves the protection status of a loaded-or-missing sidecar against the
/// active generation and layout. `loaded` is the result of attempting to load
/// the sidecar at the active generation's path.
///
/// - missing sidecar (NotFound) => [`BitrotStatus::Off`]
/// - load/self-integrity failure => [`BitrotStatus::Invalid`]
/// - generation mismatch => [`BitrotStatus::Off`]
/// - manifest validation failure => [`BitrotStatus::Invalid`]
/// - otherwise => [`BitrotStatus::On`]
pub fn resolve_status(
loaded: &Result<EcBitrotProtection, BitrotLoadError>,
active_generation: u32,
data_shards: usize,
parity_shards: usize,
) -> BitrotStatus {
match loaded {
Err(BitrotLoadError::Io(e)) if e.kind() == io::ErrorKind::NotFound => BitrotStatus::Off,
Err(_) => BitrotStatus::Invalid,
Ok(prot) => {
if prot.generation != active_generation {
return BitrotStatus::Off;
}
if validate_manifest(prot, data_shards, parity_shards).is_err() {
return BitrotStatus::Invalid;
}
BitrotStatus::On
}
}
}
/// Whether a generation-matching sidecar agrees with the geometry the volume is
/// mounted with. Both files record the layout the generation was encoded with,
/// so a disagreement means one of them is wrong and reads through the other
/// would land at the wrong shard offsets — the caller fails the mount rather
/// than merely dropping protection. A sidecar that records no EC config has
/// nothing to contradict.
pub fn geometry_matches(
prot: &EcBitrotProtection,
data_shards: usize,
parity_shards: usize,
block_size: i64,
) -> bool {
match &prot.ec_shard_config {
None => true,
Some(cfg) => {
cfg.data_shards as usize == data_shards
&& cfg.parity_shards as usize == parity_shards
&& cfg.block_size == block_size
}
}
}
/// Returns the [`EcShardChecksums`] entry for a shard id, or `None`.
pub fn shard_checksums(prot: &EcBitrotProtection, shard_id: u32) -> Option<&EcShardChecksums> {
prot.shards.iter().find(|s| s.shard_id == shard_id)
}
/// Reads a shard file at `path` in `block_size` chunks and compares each block's
/// CRC32C against the manifest entry. Returns the list of mismatching block
/// indices (empty == clean), or a fatal `io::Error` for genuine I/O problems.
///
/// A length mismatch (truncation or unexpected trailing bytes) is itself shard
/// corruption: every block index is reported as mismatched so the caller treats
/// the shard as bad. It does not interpret the result — the caller (scrub /
/// rebuild) arbitrates shard-vs-sidecar via Reed-Solomon before acting.
pub fn verify_shard_file_blocks(
path: &str,
entry: &EcShardChecksums,
block_size: i64,
) -> io::Result<Vec<usize>> {
verify_shard_blocks(&File::open(path)?, entry, block_size)
}
/// Same verification against an ALREADY-OPEN shard handle.
///
/// Go's `ChecksumScrub` reads through `shard.ReadAt`, i.e. the handle the
/// EcVolumeShard already holds, so a concurrent teardown that unlinks the shard
/// cannot turn an intentional removal into a scrub read error. A scrub that
/// runs with the store lock released has to read the same way — see
/// `EcChecksumScrubPlan`.
pub fn verify_shard_blocks(
f: &File,
entry: &EcShardChecksums,
block_size: i64,
) -> io::Result<Vec<usize>> {
let file_size = f.metadata()?.len() as i64;
let want = unpack_u32_le(&entry.block_crc32c);
if file_size != entry.covered_size {
// Length drift is shard corruption: report every block as mismatched.
return Ok((0..want.len()).collect());
}
let mut mismatched = Vec::new();
let mut buf = vec![0u8; block_size.max(1) as usize];
let mut offset: i64 = 0;
for (i, want_crc) in want.iter().enumerate() {
let to_read = (entry.covered_size - offset).min(block_size);
if to_read <= 0 {
break;
}
let to_read = to_read as usize;
read_full_at(f, &mut buf[..to_read], offset as u64)?;
if CRC::new(&buf[..to_read]).0 != *want_crc {
mismatched.push(i);
}
offset += to_read as i64;
}
Ok(mismatched)
}
/// Reads exactly `buf.len()` bytes from `f` at `offset`, erroring on early EOF.
fn read_full_at(f: &File, buf: &mut [u8], offset: u64) -> io::Result<()> {
let mut total = 0usize;
while total < buf.len() {
#[cfg(unix)]
let n = {
use std::os::unix::fs::FileExt;
f.read_at(&mut buf[total..], offset + total as u64)?
};
#[cfg(not(unix))]
let n = {
use std::io::{Read, Seek, SeekFrom};
let mut fc = f.try_clone()?;
fc.seek(SeekFrom::Start(offset + total as u64))?;
fc.read(&mut buf[total..])?
};
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"short read on shard block",
));
}
total += n;
}
Ok(())
}
/// Builds the `EcShardConfig` proto for the given layout. The bitrot sidecar
/// carries its own top-level encode_uuid, so the nested config leaves it empty.
pub fn ec_shard_config(data_shards: u32, parity_shards: u32, block_size: i64) -> EcShardConfig {
EcShardConfig {
data_shards,
parity_shards,
encode_ts_ns: 0,
block_size,
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Cross-binary byte-exact proof: CANONICAL_HEX equals `canonicalInteropHex`
/// in weed/storage/erasure_coding/ec_bitrot_interop_test.go (the Go reference
/// asserts the same constant). Both binaries port the same format, so a
/// sidecar written by either must be byte-identical; this pins the Rust side
/// of that guarantee. If you change the format, regenerate and update BOTH.
#[test]
fn test_byte_exact_go_interop() {
const CANONICAL_HEX: &str = "45435355000100000039cc1b826a080110808080082204080a10042a0a108080401a04040302012a0c0801108080401a04080706053210000102030405060708090a0b0c0d0e0f";
let prot = EcBitrotProtection {
algorithm: ChecksumAlgorithm::ChecksumCrc32c as i32,
block_size: DEFAULT_BITROT_BLOCK_SIZE as u32,
generation: 0,
ec_shard_config: Some(EcShardConfig {
data_shards: 10,
parity_shards: 4,
encode_ts_ns: 0,
block_size: 0,
}),
shards: vec![
EcShardChecksums {
shard_id: 0,
covered_size: 1024 * 1024,
block_crc32c: pack_u32_le(&[0x0102_0304]),
},
EcShardChecksums {
shard_id: 1,
covered_size: 1024 * 1024,
block_crc32c: pack_u32_le(&[0x0506_0708]),
},
],
encode_uuid: (0..16u8).collect(),
};
let dir = std::env::temp_dir().join(format!("ecsum_interop_{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("v1.ecsum");
let path = path.to_str().unwrap();
save_bitrot_sidecar(path, &prot).unwrap();
let bytes = std::fs::read(path).unwrap();
let hex: String = bytes.iter().map(|b| format!("{:02x}", b)).collect();
assert_eq!(
hex, CANONICAL_HEX,
"Rust .ecsum bytes drifted from the Go canonical form"
);
let _ = std::fs::remove_file(path);
}
#[test]
fn test_sidecar_path_generations() {
assert_eq!(bitrot_sidecar_path("/d/1", 0), "/d/1.ecsum");
assert_eq!(bitrot_sidecar_path("/d/1", 3), "/d/1.ecsum.v3");
}
/// Mirrors Go's TestRemoveBitrotSidecars: legacy + versioned sidecars go,
/// unrelated files (a shard, a longer-vid sidecar) survive, absent is success.
#[test]
fn test_remove_bitrot_sidecars() {
let tmp = tempfile::TempDir::new().unwrap();
let base = tmp.path().join("5").to_str().unwrap().to_string();
for p in [
format!("{}.ecsum", base),
format!("{}.ecsum.v1", base),
format!("{}.ecsum.v7", base),
] {
std::fs::write(&p, b"x").unwrap();
}
let keep_shard = format!("{}.ec00", base);
let keep_other_vid = format!("{}0.ecsum", base);
std::fs::write(&keep_shard, b"x").unwrap();
std::fs::write(&keep_other_vid, b"x").unwrap();
remove_bitrot_sidecars(&base).unwrap();
for p in [
format!("{}.ecsum", base),
format!("{}.ecsum.v1", base),
format!("{}.ecsum.v7", base),
] {
assert!(
!std::path::Path::new(&p).exists(),
"{} should be removed",
p
);
}
assert!(std::path::Path::new(&keep_shard).exists());
assert!(std::path::Path::new(&keep_other_vid).exists());
// Already-gone is success.
remove_bitrot_sidecars(&base).unwrap();
}
#[test]
fn test_is_pow2_multiple_of_1mib() {
assert!(is_pow2_multiple_of_1mib(1 << 20)); // 1 MiB
assert!(is_pow2_multiple_of_1mib(16 * 1024 * 1024)); // 16 MiB default
assert!(is_pow2_multiple_of_1mib(1 << 25));
assert!(is_pow2_multiple_of_1mib(MAX_BITROT_BLOCK_SIZE)); // 64 MiB boundary
assert!(!is_pow2_multiple_of_1mib(0));
assert!(!is_pow2_multiple_of_1mib(1 << 19)); // 512 KiB, too small
assert!(!is_pow2_multiple_of_1mib(3 << 20)); // 3 MiB, not pow2
assert!(!is_pow2_multiple_of_1mib(128 * 1024 * 1024)); // pow2 but > MAX_BITROT_BLOCK_SIZE
assert!(!is_pow2_multiple_of_1mib(
DEFAULT_BITROT_BLOCK_SIZE as u32 + 1
));
}
#[test]
fn test_expected_block_count() {
assert_eq!(expected_block_count(0, 16), 0);
assert_eq!(expected_block_count(1, 16), 1);
assert_eq!(expected_block_count(16, 16), 1);
assert_eq!(expected_block_count(17, 16), 2);
assert_eq!(expected_block_count(32, 16), 2);
assert_eq!(expected_block_count(100, 0), 0);
}
#[test]
fn test_pack_unpack_roundtrip() {
let vals = vec![0x0102_0304u32, 0xdead_beef, 0, u32::MAX];
let packed = pack_u32_le(&vals);
assert_eq!(packed.len(), 16);
// Verify little-endian byte order of the first entry.
assert_eq!(&packed[0..4], &[0x04, 0x03, 0x02, 0x01]);
assert_eq!(unpack_u32_le(&packed), vals);
}
#[test]
fn test_builder_block_boundaries() {
// block_size = 4; feed 10 bytes in chunks that cross boundaries.
let mut b = ShardChecksumBuilder::new(4);
let data = b"0123456789";
b.write(&data[0..3]); // partial block 0
b.write(&data[3..7]); // completes block 0 (idx 3), fills block 1
b.write(&data[7..10]); // partial block 2
let (covered, packed) = b.finalize();
assert_eq!(covered, 10);
let crcs = unpack_u32_le(&packed);
// ceil(10/4) = 3 blocks
assert_eq!(crcs.len(), 3);
// Compare against direct per-block CRCs.
assert_eq!(crcs[0], CRC::new(&data[0..4]).0);
assert_eq!(crcs[1], CRC::new(&data[4..8]).0);
assert_eq!(crcs[2], CRC::new(&data[8..10]).0);
}
#[test]
fn test_builder_exact_block_multiple() {
let mut b = ShardChecksumBuilder::new(4);
b.write(b"01234567"); // exactly 2 blocks, no partial
let (covered, packed) = b.finalize();
assert_eq!(covered, 8);
assert_eq!(unpack_u32_le(&packed).len(), 2);
}
#[test]
fn test_save_load_roundtrip() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("vol.ecsum").to_str().unwrap().to_string();
let mut builder = ShardChecksumBuilder::new(DEFAULT_BITROT_BLOCK_SIZE as i64);
builder.write(b"hello world");
let (covered, packed) = builder.finalize();
let prot = EcBitrotProtection {
algorithm: ChecksumAlgorithm::ChecksumCrc32c as i32,
block_size: DEFAULT_BITROT_BLOCK_SIZE as u32,
generation: 0,
ec_shard_config: Some(ec_shard_config(10, 4, 0)),
shards: vec![EcShardChecksums {
shard_id: 0,
covered_size: covered,
block_crc32c: packed,
}],
encode_uuid: new_encode_uuid(),
};
save_bitrot_sidecar(&path, &prot).unwrap();
let loaded = load_bitrot_sidecar(&path).unwrap();
assert_eq!(loaded, prot);
}
#[test]
fn test_load_rejects_bad_magic() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("bad.ecsum").to_str().unwrap().to_string();
std::fs::write(&path, vec![0u8; BITROT_HEADER_SIZE + 4]).unwrap();
match load_bitrot_sidecar(&path) {
Err(BitrotLoadError::BadMagic(_)) => {}
other => panic!("expected BadMagic, got {:?}", other),
}
}
#[test]
fn test_load_rejects_corrupted_payload() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp
.path()
.join("corrupt.ecsum")
.to_str()
.unwrap()
.to_string();
let prot = EcBitrotProtection {
algorithm: ChecksumAlgorithm::ChecksumCrc32c as i32,
block_size: DEFAULT_BITROT_BLOCK_SIZE as u32,
generation: 0,
ec_shard_config: Some(ec_shard_config(10, 4, 0)),
shards: vec![EcShardChecksums {
shard_id: 0,
covered_size: 5,
block_crc32c: pack_u32_le(&[CRC::new(b"hello").0]),
}],
encode_uuid: vec![0u8; 16],
};
save_bitrot_sidecar(&path, &prot).unwrap();
// Flip a byte in the payload (after the 14-byte header).
let mut data = std::fs::read(&path).unwrap();
let last = data.len() - 1;
data[last] ^= 0xff;
std::fs::write(&path, &data).unwrap();
match load_bitrot_sidecar(&path) {
Err(BitrotLoadError::CrcMismatch { .. }) => {}
other => panic!("expected CrcMismatch, got {:?}", other),
}
}
#[test]
fn test_load_missing_is_notfound() {
let res = load_bitrot_sidecar("/nonexistent/path/x.ecsum");
match res {
Err(BitrotLoadError::Io(e)) => assert_eq!(e.kind(), io::ErrorKind::NotFound),
other => panic!("expected Io(NotFound), got {:?}", other),
}
}
fn good_manifest() -> EcBitrotProtection {
let mut shards = Vec::new();
for id in 0..14u32 {
shards.push(EcShardChecksums {
shard_id: id,
covered_size: 5,
block_crc32c: pack_u32_le(&[CRC::new(b"hello").0]),
});
}
EcBitrotProtection {
algorithm: ChecksumAlgorithm::ChecksumCrc32c as i32,
block_size: DEFAULT_BITROT_BLOCK_SIZE as u32,
generation: 0,
ec_shard_config: Some(ec_shard_config(10, 4, 0)),
shards,
encode_uuid: vec![0u8; 16],
}
}
#[test]
fn test_validate_manifest_ok() {
assert!(validate_manifest(&good_manifest(), 10, 4).is_ok());
}
#[test]
fn test_validate_rejects_wrong_algorithm() {
let mut m = good_manifest();
m.algorithm = ChecksumAlgorithm::ChecksumNone as i32;
assert!(validate_manifest(&m, 10, 4).is_err());
}
#[test]
fn test_validate_rejects_bad_block_size() {
let mut m = good_manifest();
m.block_size = 3 << 20;
assert!(validate_manifest(&m, 10, 4).is_err());
}
#[test]
fn test_validate_rejects_incomplete() {
let mut m = good_manifest();
m.shards.pop();
assert!(validate_manifest(&m, 10, 4).is_err());
}
#[test]
fn test_validate_rejects_duplicate_shard_id() {
let mut m = good_manifest();
m.shards[1].shard_id = 0;
assert!(validate_manifest(&m, 10, 4).is_err());
}
#[test]
fn test_validate_rejects_out_of_range_id() {
let mut m = good_manifest();
m.shards[13].shard_id = 14;
assert!(validate_manifest(&m, 10, 4).is_err());
}
#[test]
fn test_validate_rejects_nonpositive_covered_size() {
let mut m = good_manifest();
m.shards[0].covered_size = 0;
assert!(validate_manifest(&m, 10, 4).is_err());
}
#[test]
fn test_validate_rejects_crc_count_mismatch() {
let mut m = good_manifest();
m.shards[0].block_crc32c = vec![0u8; 8]; // 2 entries but covered_size=5 => want 1
assert!(validate_manifest(&m, 10, 4).is_err());
}
#[test]
fn test_resolve_status() {
// Missing => Off.
let notfound: Result<EcBitrotProtection, BitrotLoadError> = Err(BitrotLoadError::Io(
io::Error::new(io::ErrorKind::NotFound, "x"),
));
assert_eq!(resolve_status(&notfound, 0, 10, 4), BitrotStatus::Off);
// Integrity failure => Invalid.
let bad: Result<EcBitrotProtection, BitrotLoadError> = Err(BitrotLoadError::BadMagic(0));
assert_eq!(resolve_status(&bad, 0, 10, 4), BitrotStatus::Invalid);
// Generation mismatch => Off.
let mut m = good_manifest();
m.generation = 7;
let ok: Result<EcBitrotProtection, BitrotLoadError> = Ok(m);
assert_eq!(resolve_status(&ok, 0, 10, 4), BitrotStatus::Off);
// Matching + valid => On.
let ok2: Result<EcBitrotProtection, BitrotLoadError> = Ok(good_manifest());
assert_eq!(resolve_status(&ok2, 0, 10, 4), BitrotStatus::On);
// Matching generation but invalid manifest => Invalid.
let mut bad_m = good_manifest();
bad_m.shards.pop();
let ok3: Result<EcBitrotProtection, BitrotLoadError> = Ok(bad_m);
assert_eq!(resolve_status(&ok3, 0, 10, 4), BitrotStatus::Invalid);
}
#[test]
fn test_verify_shard_file_blocks() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("s.ec00").to_str().unwrap().to_string();
let block_size: i64 = 4;
let data = b"0123456789"; // 10 bytes, 3 blocks
std::fs::write(&path, data).unwrap();
let mut b = ShardChecksumBuilder::new(block_size);
b.write(data);
let (covered, packed) = b.finalize();
let entry = EcShardChecksums {
shard_id: 0,
covered_size: covered,
block_crc32c: packed,
};
// Clean file => no mismatches.
let mm = verify_shard_file_blocks(&path, &entry, block_size).unwrap();
assert!(mm.is_empty());
// Corrupt block index 1 (bytes 4..8).
let mut corrupt = data.to_vec();
corrupt[5] ^= 0xff;
std::fs::write(&path, &corrupt).unwrap();
let mm = verify_shard_file_blocks(&path, &entry, block_size).unwrap();
assert_eq!(mm, vec![1]);
// Truncation => all blocks mismatched.
std::fs::write(&path, b"012").unwrap();
let mm = verify_shard_file_blocks(&path, &entry, block_size).unwrap();
assert_eq!(mm, vec![0, 1, 2]);
}
}

View File

@@ -0,0 +1,697 @@
//! EC decoding: reconstruct a .dat file from EC shards.
//!
//! Rebuilds the original .dat + .idx files from data shards (.ec00-.ec09)
//! and the sorted index (.ecx) + deletion journal (.ecj).
use std::fs::File;
use std::io::{self, Read, Write};
use crate::storage::erasure_coding::ec_shard::*;
use crate::storage::idx;
use crate::storage::needle::needle::get_actual_size;
use crate::storage::super_block::SUPER_BLOCK_SIZE;
use crate::storage::types::*;
use crate::storage::volume::{fsync_dir, volume_file_name};
/// Calculate .dat file size from the max offset entry in .ecx.
/// Reads the volume version from the first EC shard (.ec00) superblock,
/// then scans .ecx entries to find the largest (offset + needle_actual_size).
///
/// `dir` is used both for reading `.ec00` and `.ecx`. For split-disk
/// reconciled volumes call [`find_dat_file_size_with_dirs`] instead.
pub fn find_dat_file_size(dir: &str, collection: &str, volume_id: VolumeId) -> io::Result<i64> {
find_dat_file_size_with_dirs(dir, dir, collection, volume_id)
}
/// Like [`find_dat_file_size`] but lets the caller pass separate dirs
/// for `.ec00` (the data shard) and `.ecx` (the sealed index). This
/// is the form needed when shards are split across data dirs and the
/// `.ecx` lives on a sibling disk's idx dir (#9252).
pub fn find_dat_file_size_with_dirs(
ec00_dir: &str,
ecx_dir: &str,
collection: &str,
volume_id: VolumeId,
) -> io::Result<i64> {
let ec00_base = volume_file_name(ec00_dir, collection, volume_id);
let ecx_base = volume_file_name(ecx_dir, collection, volume_id);
// Read volume version from .ec00 superblock
let ec00_path = format!("{}.ec00", ec00_base);
let mut ec00 = File::open(&ec00_path)?;
let mut sb_buf = [0u8; SUPER_BLOCK_SIZE];
ec00.read_exact(&mut sb_buf)?;
let version = Version(sb_buf[0]);
// Start with at least the superblock size
let mut dat_size: i64 = SUPER_BLOCK_SIZE as i64;
// Scan .ecx entries
let ecx_path = format!("{}.ecx", ecx_base);
let ecx_data = std::fs::read(&ecx_path)?;
let entry_count = ecx_data.len() / NEEDLE_MAP_ENTRY_SIZE;
for i in 0..entry_count {
let start = i * NEEDLE_MAP_ENTRY_SIZE;
let (_, offset, size) =
idx_entry_from_bytes(&ecx_data[start..start + NEEDLE_MAP_ENTRY_SIZE]);
if size.is_deleted() {
continue;
}
let entry_stop = offset.to_actual_offset() + get_actual_size(size, version);
if entry_stop > dat_size {
dat_size = entry_stop;
}
}
Ok(dat_size)
}
/// What it takes to rebuild a volume's .dat from its EC data shards.
///
/// Mirrors Go's `WriteDatFile(baseFileName, datFileSize,
/// encodedDatFileSize, shardFileNames)` shape — Go passes per-shard
/// paths so a reconciled volume with shards split across disks of the
/// same volume server can still be decoded back to a regular .dat
/// (seaweedfs/seaweedfs#9252).
#[derive(Clone, Copy, Debug)]
pub struct DatRebuild<'a> {
/// Where the produced `.dat` is written.
pub dat_dir: &'a str,
pub collection: &'a str,
pub volume_id: VolumeId,
/// The number of bytes to write, i.e. the live data extent from
/// [`find_dat_file_size`].
pub dat_file_size: i64,
/// The .dat size at encode time, which fixed the shard block layout:
/// deletions can move the live extent below the large-block row
/// boundary, and deriving the layout from the shrunk extent would read
/// the shards in the wrong block order. Zero when the .vif does not
/// record the encode-time size; the layout is then inferred from the
/// shard size.
pub encoded_dat_file_size: i64,
pub data_shards: usize,
/// `shard_dirs[i]` is the directory holding shard `i`. `None` means every
/// data shard sits in `dat_dir`.
pub shard_dirs: Option<&'a [String]>,
/// The volume's shard block layout, e.g. `EcVolume::large_block_size()`
/// / `small_block_size()` from its .vif EC config.
pub large_block_size: usize,
pub small_block_size: usize,
}
/// Reconstruct a .dat file from EC data shards.
///
/// Reads from .ec00-.ec09 and writes a new .dat file, from one directory or
/// from the per-shard directories of a cross-disk reconciled volume.
pub fn write_dat_file_from_shards(spec: &DatRebuild<'_>) -> io::Result<()> {
let DatRebuild {
dat_dir,
collection,
volume_id,
dat_file_size,
encoded_dat_file_size,
data_shards,
shard_dirs,
large_block_size,
small_block_size,
} = *spec;
let same_dir: Vec<String>;
let shard_dirs: &[String] = match shard_dirs {
Some(dirs) => dirs,
None => {
same_dir = vec![dat_dir.to_string(); data_shards];
&same_dir
}
};
if data_shards == 0 {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"no data shards",
));
}
if shard_dirs.len() < data_shards {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"shard_dirs len {} < data_shards {}",
shard_dirs.len(),
data_shards
),
));
}
let base = volume_file_name(dat_dir, collection, volume_id);
let dat_path = format!("{}.dat", base);
// Write to a temp file and atomically rename into place, so a crash
// mid-write never leaves a partial .dat at the final name beside the
// source shards.
let tmp_path = format!("{}.tmp", dat_path);
let write_result = (|| -> io::Result<()> {
// Open data shards from their individual home dirs.
let mut shards: Vec<EcVolumeShard> = (0..data_shards as u8)
.map(|i| EcVolumeShard::new(&shard_dirs[i as usize], collection, volume_id, i))
.collect();
for shard in &mut shards {
shard.open()?;
}
let mut encoded_dat_file_size = encoded_dat_file_size;
if encoded_dat_file_size <= 0 {
// .vif without the encode-time size: infer the padded layout from
// the physical shard size, which reads the shards in the same
// block order.
let shard_size = std::fs::metadata(shards[0].file_name())?.len() as i64;
// A shard size that is an exact multiple of the large block size
// is ambiguous: N large rows, or N-1 large rows plus a full
// small-block region. The two layouts only agree below the last
// large row.
let large = large_block_size as i64;
if shard_size % large == 0
&& dat_file_size > (shard_size / large - 1) * large * data_shards as i64
{
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"shard size {} does not identify the block layout; re-encode to record the dat size in .vif",
shard_size
),
));
}
encoded_dat_file_size = data_shards as i64 * shard_size;
}
if dat_file_size > encoded_dat_file_size {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"dat file size {} exceeds encoded dat file size {}",
dat_file_size, encoded_dat_file_size
),
));
}
let mut dat_file = File::create(&tmp_path)?;
let mut remaining = dat_file_size;
let mut encoded_remaining = encoded_dat_file_size;
let large_row_size = (large_block_size * data_shards) as i64;
let mut shard_offset: u64 = 0;
// Read large blocks
while encoded_remaining >= large_row_size && remaining > 0 {
for (i, shard) in shards[..data_shards].iter().enumerate() {
let to_write = large_block_size.min(remaining as usize);
let mut buf = vec![0u8; to_write];
let n = shard.read_at(&mut buf, shard_offset)?;
if n != to_write {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
format!("short read of large block on shard {}", i),
));
}
dat_file.write_all(&buf)?;
remaining -= to_write as i64;
if remaining <= 0 {
break;
}
}
encoded_remaining -= large_row_size;
shard_offset += large_block_size as u64;
}
// Read small blocks
while remaining > 0 {
for (i, shard) in shards[..data_shards].iter().enumerate() {
let to_write = small_block_size.min(remaining as usize);
let mut buf = vec![0u8; to_write];
let n = shard.read_at(&mut buf, shard_offset)?;
if n != to_write {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
format!("short read of small block on shard {}", i),
));
}
dat_file.write_all(&buf)?;
remaining -= to_write as i64;
if remaining <= 0 {
break;
}
}
shard_offset += small_block_size as u64;
}
for shard in &mut shards {
shard.close();
}
// fsync, rename, then fsync the dir so the decoded .dat is durable and
// atomically published before the caller deletes the source shards.
dat_file.sync_all()?;
drop(dat_file);
// Windows rename does not replace an existing file on every version;
// remove the destination first, matching the compaction commit path.
#[cfg(windows)]
{
let _ = std::fs::remove_file(&dat_path);
}
std::fs::rename(&tmp_path, &dat_path)?;
fsync_dir(&dat_path)?;
Ok(())
})();
if write_result.is_err() {
let _ = std::fs::remove_file(&tmp_path);
}
write_result
}
/// Write .idx file from .ecx index + .ecj deletion journal.
///
/// Copies sorted .ecx entries to .idx, then appends tombstones for
/// deleted needles from .ecj.
pub fn write_idx_file_from_ec_index(
dir: &str,
collection: &str,
volume_id: VolumeId,
) -> io::Result<()> {
let base = volume_file_name(dir, collection, volume_id);
let ecx_path = format!("{}.ecx", base);
let ecj_path = format!("{}.ecj", base);
let idx_path = format!("{}.idx", base);
// Write to a temp file and atomically rename into place, so a crash
// mid-write never leaves a partial .idx at the final name beside the
// source shards.
let tmp_path = format!("{}.tmp", idx_path);
let write_result = (|| -> io::Result<()> {
// Copy .ecx to the temp .idx
std::fs::copy(&ecx_path, &tmp_path)?;
// Append deletions from .ecj as tombstones. Read the journal directly
// and treat only NotFound as "no journal": Path::exists would also
// swallow a permission/IO error and silently skip deletions, which
// would resurrect deleted needles as live.
let mut idx_file = std::fs::OpenOptions::new().append(true).open(&tmp_path)?;
match std::fs::read(&ecj_path) {
Ok(ecj_data) => {
let count = ecj_data.len() / NEEDLE_ID_SIZE;
for i in 0..count {
let start = i * NEEDLE_ID_SIZE;
let needle_id = NeedleId::from_bytes(&ecj_data[start..start + NEEDLE_ID_SIZE]);
idx::write_index_entry(
&mut idx_file,
needle_id,
Offset::default(),
TOMBSTONE_FILE_SIZE,
)?;
}
}
Err(e) if e.kind() == io::ErrorKind::NotFound => {}
Err(e) => return Err(e),
}
// fsync, rename, then fsync the dir so the decoded .idx is durable and
// atomically published before the caller deletes the source shards.
idx_file.sync_all()?;
drop(idx_file);
// Windows rename does not replace an existing file on every version;
// remove the destination first, matching the compaction commit path.
#[cfg(windows)]
{
let _ = std::fs::remove_file(&idx_path);
}
std::fs::rename(&tmp_path, &idx_path)?;
fsync_dir(&idx_path)?;
Ok(())
})();
if write_result.is_err() {
let _ = std::fs::remove_file(&tmp_path);
}
write_result
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::erasure_coding::ec_encoder;
use crate::storage::needle::needle::Needle;
use crate::storage::needle_map::NeedleMapKind;
use crate::storage::volume::{Volume, VolumeSpec};
use tempfile::TempDir;
#[test]
fn test_ec_full_round_trip() {
let tmp = TempDir::new().unwrap();
let dir = tmp.path().to_str().unwrap();
// Create volume with data
let mut v = Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
let test_data: Vec<(NeedleId, Vec<u8>)> = (1..=3)
.map(|i| {
let data = format!("EC round trip data for needle {}", i);
(NeedleId(i), data.into_bytes())
})
.collect();
for (id, data) in &test_data {
let mut n = Needle {
id: *id,
cookie: Cookie(id.0 as u32),
data: data.clone(),
data_size: data.len() as u32,
..Needle::default()
};
v.write_needle(&mut n, true, false).unwrap();
}
v.sync_to_disk().unwrap();
let original_dat_size = v.dat_file_size().unwrap();
v.close();
// Read original .dat for comparison
let original_dat = std::fs::read(format!("{}/1.dat", dir)).unwrap();
// Encode to EC
let data_shards = 10;
let parity_shards = 4;
let block_size =
ec_encoder::write_ec_files(dir, dir, "", VolumeId(1), data_shards, parity_shards)
.unwrap();
// Delete original .dat and .idx
std::fs::remove_file(format!("{}/1.dat", dir)).unwrap();
std::fs::remove_file(format!("{}/1.idx", dir)).unwrap();
// Reconstruct from EC shards
write_dat_file_from_shards(&DatRebuild {
dat_dir: dir,
collection: "",
volume_id: VolumeId(1),
dat_file_size: original_dat_size as i64,
encoded_dat_file_size: original_dat_size as i64,
data_shards,
shard_dirs: None,
large_block_size: block_size as usize,
small_block_size: block_size as usize,
})
.unwrap();
write_idx_file_from_ec_index(dir, "", VolumeId(1)).unwrap();
// Atomic publish must rename the temp files away, never leaving them behind.
assert!(!std::path::Path::new(&format!("{}/1.dat.tmp", dir)).exists());
assert!(!std::path::Path::new(&format!("{}/1.idx.tmp", dir)).exists());
// Verify reconstructed .dat matches original
let reconstructed_dat = std::fs::read(format!("{}/1.dat", dir)).unwrap();
assert_eq!(
original_dat[..original_dat_size as usize],
reconstructed_dat[..original_dat_size as usize],
"reconstructed .dat should match original"
);
// Verify we can load and read from reconstructed volume
let v2 = Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
for (id, expected_data) in &test_data {
let mut n = Needle {
id: *id,
..Needle::default()
};
v2.read_needle(&mut n).unwrap();
assert_eq!(&n.data, expected_data, "needle {} data should match", id);
}
}
#[test]
fn test_decode_missing_shard_leaves_no_dat() {
let tmp = TempDir::new().unwrap();
let dir = tmp.path().to_str().unwrap();
// No shard files exist, so de-striping must fail and publish nothing:
// neither the final .dat nor a partial .dat.tmp may remain.
let res = write_dat_file_from_shards(&DatRebuild {
dat_dir: dir,
collection: "",
volume_id: VolumeId(7),
dat_file_size: 100,
encoded_dat_file_size: 100,
data_shards: 10,
shard_dirs: None,
large_block_size: ERASURE_CODING_LARGE_BLOCK_SIZE,
small_block_size: ERASURE_CODING_SMALL_BLOCK_SIZE,
});
assert!(res.is_err());
assert!(!std::path::Path::new(&format!("{}/7.dat", dir)).exists());
assert!(!std::path::Path::new(&format!("{}/7.dat.tmp", dir)).exists());
}
// Decoding when .vif does not record the encode-time size: the layout is
// inferred from the shard size, except when that is an exact large-block
// multiple and the live extent reaches the ambiguous region.
#[test]
fn test_write_dat_file_fallback_layout() {
use crate::storage::erasure_coding::ec_bitrot::{
DEFAULT_BITROT_BLOCK_SIZE, ShardChecksumBuilder,
};
use reed_solomon_erasure::galois_8::ReedSolomon;
const LARGE: usize = 10000;
const SMALL: usize = 100;
let data_shards = 10usize;
let parity_shards = 4usize;
let large_row_size = (LARGE * data_shards) as i64;
let tmp = TempDir::new().unwrap();
let encode = |name: &str, dat_size: i64| -> (String, Vec<String>, Vec<u8>) {
let dir = tmp.path().join(name);
std::fs::create_dir(&dir).unwrap();
let dir = dir.to_str().unwrap().to_string();
let original: Vec<u8> = (0..dat_size as usize)
.map(|i| ((i as u64).wrapping_mul(2654435761) >> 8) as u8)
.collect();
std::fs::write(format!("{}/1.dat", dir), &original).unwrap();
let dat_file = File::open(format!("{}/1.dat", dir)).unwrap();
let rs = ReedSolomon::new(data_shards, parity_shards).unwrap();
let total = data_shards + parity_shards;
let mut shards: Vec<EcVolumeShard> = (0..total as u8)
.map(|i| EcVolumeShard::new(&dir, "", VolumeId(1), i))
.collect();
for shard in &mut shards {
shard.create().unwrap();
}
let mut builders: Vec<ShardChecksumBuilder> = (0..total)
.map(|_| ShardChecksumBuilder::new(DEFAULT_BITROT_BLOCK_SIZE as i64))
.collect();
ec_encoder::encode_dat_file(
&dat_file,
dat_size,
&rs,
&mut shards,
&mut builders,
ec_encoder::EcEncodeLayout {
data_shards,
parity_shards,
buffer_size: SMALL,
large_block_size: LARGE,
small_block_size: SMALL,
},
)
.unwrap();
for shard in &mut shards {
shard.close();
}
let shard_dirs: Vec<String> = (0..data_shards).map(|_| dir.clone()).collect();
(dir, shard_dirs, original)
};
let decode_to = |dir: &str,
sub: &str,
live: i64,
encoded: i64,
shard_dirs: &[String]|
-> io::Result<Vec<u8>> {
let out = format!("{}/{}", dir, sub);
std::fs::create_dir_all(&out).unwrap();
write_dat_file_from_shards(&DatRebuild {
dat_dir: &out,
collection: "",
volume_id: VolumeId(1),
dat_file_size: live,
encoded_dat_file_size: encoded,
data_shards: 10,
shard_dirs: Some(shard_dirs),
large_block_size: LARGE,
small_block_size: SMALL,
})?;
Ok(std::fs::read(format!("{}/1.dat", out)).unwrap())
};
// a small-block tail that is not a large-block multiple is unambiguous
let (dir, shard_dirs, original) = encode("plain", large_row_size + 2530);
let live = large_row_size / 2;
let decoded = decode_to(&dir, "out", live, 0, &shard_dirs).unwrap();
assert_eq!(&original[..live as usize], &decoded[..]);
// datSize just under one large row: the full small-block region makes
// each shard exactly one large block, indistinguishable from one large row
let (dir, shard_dirs, _) = encode("ambig1", large_row_size - 1);
let err = decode_to(&dir, "out", large_row_size / 2, 0, &shard_dirs).unwrap_err();
assert!(
err.to_string()
.contains("does not identify the block layout")
);
// two-row equivalent: decoding within the agreed prefix still works
let (dir, shard_dirs, original) = encode("ambig2", 2 * large_row_size - 1);
let decoded = decode_to(&dir, "outa", large_row_size, 0, &shard_dirs).unwrap();
assert_eq!(&original[..large_row_size as usize], &decoded[..]);
let err = decode_to(&dir, "outb", large_row_size + 1, 0, &shard_dirs).unwrap_err();
assert!(
err.to_string()
.contains("does not identify the block layout")
);
}
// Decoding after deletions moved the live extent below the large-block row
// boundary: the shard block layout is fixed by the encode-time .dat size,
// so de-striping must not derive it from the shrunk live extent.
#[test]
fn test_write_dat_file_after_tail_deletion() {
use crate::storage::erasure_coding::ec_bitrot::{
DEFAULT_BITROT_BLOCK_SIZE, ShardChecksumBuilder,
};
use reed_solomon_erasure::galois_8::ReedSolomon;
const LARGE: usize = 10000;
const SMALL: usize = 100;
let data_shards = 10usize;
let parity_shards = 4usize;
let large_row_size = (LARGE * data_shards) as i64;
let small_row_size = (SMALL * data_shards) as i64;
// one full large-block row plus a small-block tail
let dat_size = large_row_size + 2 * small_row_size + 530;
let tmp = TempDir::new().unwrap();
let dir = tmp.path().to_str().unwrap();
let original: Vec<u8> = (0..dat_size as usize)
.map(|i| ((i as u64).wrapping_mul(2654435761) >> 8) as u8)
.collect();
std::fs::write(format!("{}/1.dat", dir), &original).unwrap();
let dat_file = File::open(format!("{}/1.dat", dir)).unwrap();
let rs = ReedSolomon::new(data_shards, parity_shards).unwrap();
let total = data_shards + parity_shards;
let mut shards: Vec<EcVolumeShard> = (0..total as u8)
.map(|i| EcVolumeShard::new(dir, "", VolumeId(1), i))
.collect();
for shard in &mut shards {
shard.create().unwrap();
}
let mut builders: Vec<ShardChecksumBuilder> = (0..total)
.map(|_| ShardChecksumBuilder::new(DEFAULT_BITROT_BLOCK_SIZE as i64))
.collect();
ec_encoder::encode_dat_file(
&dat_file,
dat_size,
&rs,
&mut shards,
&mut builders,
ec_encoder::EcEncodeLayout {
data_shards,
parity_shards,
buffer_size: SMALL,
large_block_size: LARGE,
small_block_size: SMALL,
},
)
.unwrap();
for shard in &mut shards {
shard.close();
}
let shard_size = std::fs::metadata(format!("{}/1.ec00", dir)).unwrap().len() as i64;
let padded_size = data_shards as i64 * shard_size;
let shard_dirs: Vec<String> = (0..data_shards).map(|_| dir.to_string()).collect();
// decode into a separate dir so the output does not collide with the source .dat
let out_dir = tmp.path().join("out");
std::fs::create_dir(&out_dir).unwrap();
let out = out_dir.to_str().unwrap();
let decode = |live_size: i64, encoded_size: i64| -> Vec<u8> {
write_dat_file_from_shards(&DatRebuild {
dat_dir: out,
collection: "",
volume_id: VolumeId(1),
dat_file_size: live_size,
encoded_dat_file_size: encoded_size,
data_shards,
shard_dirs: Some(&shard_dirs),
large_block_size: LARGE,
small_block_size: SMALL,
})
.unwrap();
let path = format!("{}/1.dat", out);
let decoded = std::fs::read(&path).unwrap();
std::fs::remove_file(&path).unwrap();
decoded
};
for live_size in [
LARGE as i64 - 1, // within the first large block
LARGE as i64 + 42, // partial second large block
large_row_size / 2, // mid large row
large_row_size, // exactly the large row
large_row_size + 5 * SMALL as i64, // into the small-block tail
dat_size, // nothing deleted
] {
assert_eq!(
&original[..live_size as usize],
&decode(live_size, dat_size)[..],
"live size {} with encode-time layout",
live_size
);
assert_eq!(
&original[..live_size as usize],
&decode(live_size, padded_size)[..],
"live size {} with padded layout",
live_size
);
}
// deriving the layout from the shrunk live extent reorders the data
let control = decode(large_row_size / 2, large_row_size / 2);
assert_ne!(&original[..(large_row_size / 2) as usize], &control[..]);
// the live extent can never exceed the encode-time size
assert!(
write_dat_file_from_shards(&DatRebuild {
dat_dir: out,
collection: "",
volume_id: VolumeId(1),
dat_file_size: dat_size + 1,
encoded_dat_file_size: dat_size,
data_shards,
shard_dirs: Some(&shard_dirs),
large_block_size: LARGE,
small_block_size: SMALL,
})
.is_err()
);
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,266 @@
//! EC data location: maps needle offset/size to shard intervals.
//!
//! Determines which shard(s) contain data for a given needle and at what
//! offsets within those shards. Handles both large (1GB) and small (1MB)
//! block sections.
use crate::storage::erasure_coding::ec_shard::*;
use crate::storage::types::*;
/// An interval to read from EC shards.
#[derive(Debug, Clone)]
pub struct Interval {
pub block_index: usize,
pub inner_block_offset: i64,
pub size: i64,
pub is_large_block: bool,
pub large_block_rows_count: usize,
}
impl Interval {
pub fn to_shard_id_and_offset(
&self,
data_shards: u32,
large_block_size: i64,
small_block_size: i64,
) -> (ShardId, i64) {
let data_shards_usize = data_shards as usize;
let shard_id = (self.block_index % data_shards_usize) as ShardId;
let row_index = self.block_index / data_shards_usize;
let block_size = if self.is_large_block {
large_block_size
} else {
small_block_size
};
let mut offset = row_index as i64 * block_size + self.inner_block_offset;
if !self.is_large_block {
// Small blocks come after large blocks in the shard file
offset += self.large_block_rows_count as i64 * large_block_size;
}
(shard_id, offset)
}
}
/// Locate the EC shard intervals needed to read data at the given offset and size.
///
/// `shard_size` is the size of a single shard file.
pub fn locate_data(
offset: i64,
size: Size,
shard_size: i64,
data_shards: u32,
large_block_size: i64,
small_block_size: i64,
) -> Vec<Interval> {
let mut intervals = Vec::new();
let data_size = size.0 as i64;
if data_size <= 0 || shard_size <= 0 {
return intervals;
}
let large_row_size = large_block_size * data_shards as i64;
let small_row_size = small_block_size * data_shards as i64;
// Number of large block rows. Mirrors Go's shardDatSize/largeBlockLength:
// the caller's ecd-size fallback already subtracts 1 to disambiguate the
// exact-multiple case, so no further -1 here — a shard size that IS an
// exact multiple (dat_file_size path) means real full large rows.
let n_large_block_rows = (shard_size / large_block_size) as usize;
let large_section_size = n_large_block_rows as i64 * large_row_size;
let mut remaining_offset = offset;
let mut remaining_size = data_size;
// In large block section?
if remaining_offset < large_section_size {
let available_in_large = large_section_size - remaining_offset;
let to_read = remaining_size.min(available_in_large);
add_intervals(
&mut intervals,
remaining_offset,
to_read,
large_block_size,
large_row_size,
true,
n_large_block_rows,
);
remaining_offset += to_read;
remaining_size -= to_read;
}
// In small block section?
if remaining_size > 0 {
let small_offset = remaining_offset - large_section_size;
add_intervals(
&mut intervals,
small_offset,
remaining_size,
small_block_size,
small_row_size,
false,
n_large_block_rows,
);
}
intervals
}
fn add_intervals(
intervals: &mut Vec<Interval>,
offset: i64,
size: i64,
block_size: i64,
_row_size: i64,
is_large_block: bool,
large_block_rows_count: usize,
) {
let mut pos = offset;
let end = offset + size;
while pos < end {
let block_index = (pos / block_size) as usize;
let inner_offset = pos % block_size;
let remaining_in_block = block_size - inner_offset;
let interval_size = remaining_in_block.min(end - pos);
intervals.push(Interval {
block_index,
inner_block_offset: inner_offset,
size: interval_size,
is_large_block,
large_block_rows_count,
});
pos += interval_size;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_interval_to_shard_id() {
let data_shards = 10;
let large_block_size = ERASURE_CODING_LARGE_BLOCK_SIZE as i64;
let _shard_size = 1024 * 1024; // Example shard size
// Block index 0 → shard 0
let interval = Interval {
block_index: 0,
inner_block_offset: 100,
size: 50,
is_large_block: true,
large_block_rows_count: 1,
};
let (shard_id, offset) = interval.to_shard_id_and_offset(
data_shards,
ERASURE_CODING_LARGE_BLOCK_SIZE as i64,
ERASURE_CODING_SMALL_BLOCK_SIZE as i64,
);
assert_eq!(shard_id, 0);
assert_eq!(offset, 100);
// Block index 5 → shard 5
let interval = Interval {
block_index: 5,
inner_block_offset: 0,
size: 1024,
is_large_block: true,
large_block_rows_count: 1,
};
let (shard_id, _offset) = interval.to_shard_id_and_offset(
data_shards,
ERASURE_CODING_LARGE_BLOCK_SIZE as i64,
ERASURE_CODING_SMALL_BLOCK_SIZE as i64,
);
assert_eq!(shard_id, 5);
// Block index 12 (data_shards=10) → row_index 1, shard_id 2
let interval = Interval {
block_index: 12,
inner_block_offset: 200,
size: 50,
is_large_block: true,
large_block_rows_count: 5,
};
let (shard_id, offset) = interval.to_shard_id_and_offset(
data_shards,
ERASURE_CODING_LARGE_BLOCK_SIZE as i64,
ERASURE_CODING_SMALL_BLOCK_SIZE as i64,
);
assert_eq!(shard_id, 2); // 12 % 10 = 2
assert_eq!(offset, large_block_size + 200); // row 1 offset + inner_block_offset
// Block index 10 → shard 0 (second row)
let interval = Interval {
block_index: 10,
inner_block_offset: 0,
size: 100,
is_large_block: true,
large_block_rows_count: 2,
};
let (shard_id, offset) = interval.to_shard_id_and_offset(
data_shards,
ERASURE_CODING_LARGE_BLOCK_SIZE as i64,
ERASURE_CODING_SMALL_BLOCK_SIZE as i64,
);
assert_eq!(shard_id, 0);
assert_eq!(offset, ERASURE_CODING_LARGE_BLOCK_SIZE as i64); // row 1 offset
}
#[test]
fn test_locate_data_small_file() {
// Small file: 100 bytes at offset 50, shard size = 1MB
let intervals = locate_data(
50,
Size(100),
1024 * 1024,
10,
ERASURE_CODING_LARGE_BLOCK_SIZE as i64,
ERASURE_CODING_SMALL_BLOCK_SIZE as i64,
);
assert!(!intervals.is_empty());
// Should be a single small block interval (no large block rows for 1MB shard)
assert_eq!(intervals.len(), 1);
assert!(!intervals[0].is_large_block);
}
#[test]
fn test_locate_data_empty() {
let intervals = locate_data(
0,
Size(0),
1024 * 1024,
10,
ERASURE_CODING_LARGE_BLOCK_SIZE as i64,
ERASURE_CODING_SMALL_BLOCK_SIZE as i64,
);
assert!(intervals.is_empty());
}
#[test]
fn test_small_block_after_large() {
let interval = Interval {
block_index: 0,
inner_block_offset: 0,
size: 100,
is_large_block: false,
large_block_rows_count: 2,
};
let (_shard_id, offset) = interval.to_shard_id_and_offset(
10,
ERASURE_CODING_LARGE_BLOCK_SIZE as i64,
ERASURE_CODING_SMALL_BLOCK_SIZE as i64,
);
// Should be after 2 large block rows
assert_eq!(offset, 2 * ERASURE_CODING_LARGE_BLOCK_SIZE as i64);
}
}

View File

@@ -0,0 +1,254 @@
//! EcVolumeShard: a single shard file (.ec00-.ec13) of an erasure-coded volume.
use std::fs::{self, File, OpenOptions};
use std::io::{self, Write};
use crate::storage::types::*;
use crate::storage::volume_open::open_volume_file;
pub const DATA_SHARDS_COUNT: usize = 10;
pub const PARITY_SHARDS_COUNT: usize = 4;
pub const TOTAL_SHARDS_COUNT: usize = DATA_SHARDS_COUNT + PARITY_SHARDS_COUNT;
pub const MAX_SHARD_COUNT: usize = 32;
pub const MIN_TOTAL_DISKS: usize = TOTAL_SHARDS_COUNT / PARITY_SHARDS_COUNT + 1;
pub const ERASURE_CODING_LARGE_BLOCK_SIZE: usize = 1024 * 1024 * 1024; // 1GB
pub const ERASURE_CODING_SMALL_BLOCK_SIZE: usize = 1024 * 1024; // 1MB
pub type ShardId = u8;
/// A single erasure-coded shard file.
pub struct EcVolumeShard {
pub volume_id: VolumeId,
pub shard_id: ShardId,
pub collection: String,
pub dir: String,
pub disk_type: DiskType,
ecd_file: Option<File>,
ecd_file_size: i64,
}
impl EcVolumeShard {
/// Create a new shard reference (does not open the file).
pub fn new(dir: &str, collection: &str, volume_id: VolumeId, shard_id: ShardId) -> Self {
EcVolumeShard {
volume_id,
shard_id,
collection: collection.to_string(),
dir: dir.to_string(),
disk_type: DiskType::default(),
ecd_file: None,
ecd_file_size: 0,
}
}
/// Shard file name, e.g. "dir/collection_42.ec03"
pub fn file_name(&self) -> String {
let base =
crate::storage::volume::volume_file_name(&self.dir, &self.collection, self.volume_id);
format!("{}.ec{:02}", base, self.shard_id)
}
/// Open the shard file for reading.
pub fn open(&mut self) -> io::Result<()> {
let path = self.file_name();
let file = open_volume_file(OpenOptions::new().read(true), &path)?;
self.ecd_file_size = file.metadata()?.len() as i64;
self.ecd_file = Some(file);
Ok(())
}
/// Create the shard file for writing.
pub fn create(&mut self) -> io::Result<()> {
let path = self.file_name();
let file = open_volume_file(
OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(true),
&path,
)?;
self.ecd_file = Some(file);
self.ecd_file_size = 0;
Ok(())
}
/// Read data at a specific offset.
pub fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<usize> {
let file = self
.ecd_file
.as_ref()
.ok_or_else(|| io::Error::other("shard file not open"))?;
#[cfg(unix)]
{
use std::os::unix::fs::FileExt;
file.read_at(buf, offset)
}
#[cfg(not(unix))]
{
use std::io::{Read, Seek, SeekFrom};
// File::read_at is unix-only; fall back to seek + read.
// We need a mutable reference for seek/read, so clone the handle.
let mut f = file.try_clone()?;
f.seek(SeekFrom::Start(offset))?;
f.read(buf)
}
}
/// Write data to the shard file (appends).
pub fn write_all(&mut self, data: &[u8]) -> io::Result<()> {
let file = self
.ecd_file
.as_mut()
.ok_or_else(|| io::Error::other("shard file not open"))?;
file.write_all(data)?;
self.ecd_file_size += data.len() as i64;
Ok(())
}
pub fn file_size(&self) -> i64 {
self.ecd_file_size
}
/// A duplicate of the mounted shard handle, for a reader that has to
/// outlive the store guard.
///
/// This is the same descriptor `read_at` serves from, so it carries the
/// `O_NOATIME` from `open_volume_file` and keeps pointing at the shard
/// that was mounted, whatever later happens to the path. `dup` shares the
/// kernel file offset, which is why every read through it must be
/// positional (`read_at`), never seek-based.
pub fn try_clone_file(&self) -> io::Result<File> {
self.ecd_file
.as_ref()
.ok_or_else(|| io::Error::other("shard file not open"))?
.try_clone()
}
/// Protobuf descriptor for this shard. Mirrors Go's ToEcShardInfo.
pub fn to_ec_shard_info(&self) -> crate::pb::volume_server_pb::EcShardInfo {
crate::pb::volume_server_pb::EcShardInfo {
shard_id: self.shard_id as u32,
size: self.file_size(),
collection: self.collection.clone(),
volume_id: self.volume_id.0,
..Default::default()
}
}
/// Close the shard file.
pub fn close(&mut self) {
if let Some(ref file) = self.ecd_file {
let _ = file.sync_all();
}
self.ecd_file = None;
}
/// Delete the shard file from disk.
pub fn destroy(&mut self) {
self.close();
let _ = fs::remove_file(self.file_name());
}
}
/// ShardBits: bitmap tracking which shards are present.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ShardBits(pub u32);
impl ShardBits {
pub fn add_shard_id(&mut self, id: ShardId) {
assert!((id as usize) < 32, "shard id {} out of bounds (max 31)", id,);
self.0 |= 1 << id;
}
pub fn remove_shard_id(&mut self, id: ShardId) {
assert!((id as usize) < 32, "shard id {} out of bounds (max 31)", id,);
self.0 &= !(1 << id);
}
pub fn has_shard_id(&self, id: ShardId) -> bool {
if (id as usize) >= 32 {
return false;
}
self.0 & (1 << id) != 0
}
pub fn shard_id_count(&self) -> usize {
self.0.count_ones() as usize
}
/// Iterator over present shard IDs.
pub fn shard_ids(&self) -> Vec<ShardId> {
let mut ids = Vec::with_capacity(self.shard_id_count());
for i in 0..32 {
if self.has_shard_id(i) {
ids.push(i);
}
}
ids
}
pub fn minus(&self, other: ShardBits) -> ShardBits {
ShardBits(self.0 & !other.0)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_shard_bits() {
let mut bits = ShardBits::default();
assert_eq!(bits.shard_id_count(), 0);
bits.add_shard_id(0);
bits.add_shard_id(3);
bits.add_shard_id(13);
assert_eq!(bits.shard_id_count(), 3);
assert!(bits.has_shard_id(0));
assert!(bits.has_shard_id(3));
assert!(!bits.has_shard_id(1));
bits.remove_shard_id(3);
assert!(!bits.has_shard_id(3));
assert_eq!(bits.shard_id_count(), 2);
}
#[test]
fn test_shard_bits_ids() {
let mut bits = ShardBits::default();
bits.add_shard_id(1);
bits.add_shard_id(5);
bits.add_shard_id(9);
assert_eq!(bits.shard_ids(), vec![1, 5, 9]);
}
#[test]
fn test_shard_bits_minus() {
let mut a = ShardBits::default();
a.add_shard_id(0);
a.add_shard_id(1);
a.add_shard_id(2);
let mut b = ShardBits::default();
b.add_shard_id(1);
let c = a.minus(b);
assert_eq!(c.shard_ids(), vec![0, 2]);
}
#[test]
fn test_shard_file_name() {
let shard = EcVolumeShard::new("/data", "pics", VolumeId(42), 3);
assert_eq!(shard.file_name(), "/data/pics_42.ec03");
}
#[test]
fn test_shard_file_name_no_collection() {
let shard = EcVolumeShard::new("/data", "", VolumeId(7), 13);
assert_eq!(shard.file_name(), "/data/7.ec13");
}
}

View File

@@ -0,0 +1,17 @@
//! Erasure coding module for volume data protection.
//!
//! Encodes a volume's .dat file into 10 data + 4 parity shards using
//! Reed-Solomon erasure coding. Can reconstruct from any 10 of 14 shards.
pub mod ec_bitrot;
pub mod ec_decoder;
pub mod ec_encoder;
pub mod ec_locate;
pub mod ec_shard;
pub mod ec_volume;
pub use ec_shard::{
DATA_SHARDS_COUNT, EcVolumeShard, MAX_SHARD_COUNT, MIN_TOTAL_DISKS, PARITY_SHARDS_COUNT,
ShardId, TOTAL_SHARDS_COUNT,
};
pub use ec_volume::EcVolume;

View File

@@ -0,0 +1,252 @@
//! Index file (.idx) format: sequential 17-byte entries.
//!
//! Each entry: NeedleId(8) + Offset(5) + Size(4) = 17 bytes.
use crate::storage::needle::needle::get_actual_size;
use crate::storage::types::*;
use std::io::{self, Read, Seek, SeekFrom};
pub(crate) const ROWS_TO_READ: usize = 1024;
/// Walk all entries in an .idx file, calling `f` for each.
/// Mirrors Go's `WalkIndexFile()`.
pub fn walk_index_file<R, F>(reader: &mut R, start_from: u64, mut f: F) -> io::Result<()>
where
R: Read + Seek,
F: FnMut(NeedleId, Offset, Size) -> io::Result<()>,
{
let reader_offset = start_from * NEEDLE_MAP_ENTRY_SIZE as u64;
reader.seek(SeekFrom::Start(reader_offset))?;
let mut buf = vec![0u8; NEEDLE_MAP_ENTRY_SIZE * ROWS_TO_READ];
loop {
let count = match reader.read(&mut buf) {
Ok(0) => return Ok(()),
Ok(n) => n,
Err(ref e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(()),
Err(e) => return Err(e),
};
let mut i = 0;
while i + NEEDLE_MAP_ENTRY_SIZE <= count {
let (key, offset, size) = idx_entry_from_bytes(&buf[i..i + NEEDLE_MAP_ENTRY_SIZE]);
f(key, offset, size)?;
i += NEEDLE_MAP_ENTRY_SIZE;
}
}
}
/// Verify the integrity of an .idx/.ecx index file: walk every entry, sort by
/// (offset, size), flag overlapping needles, and check the file is a whole
/// number of entries. Mirrors Go's `idx.CheckIndexFile`. No data-file reads.
pub fn check_index_file<R: Read + Seek>(
reader: &mut R,
idx_file_size: i64,
version: Version,
) -> (u64, Vec<String>) {
let mut errs = Vec::new();
// (walk index, id, actual offset, size)
let mut entries: Vec<(usize, NeedleId, i64, Size)> = Vec::new();
let mut i = 0usize;
if let Err(e) = walk_index_file(reader, 0, |id, offset, size| {
entries.push((i, id, offset.to_actual_offset(), size));
i += 1;
Ok(())
}) {
errs.push(format!("walk index file: {}", e));
}
entries.sort_by(|a, b| a.2.cmp(&b.2).then(a.3.0.cmp(&b.3.0)));
// Offset-0 logical tombstones (remote-tier deletes) occupy no physical extent,
// so they cannot overlap anything — exclude them from the overlap check. They
// are still counted below for the index-size check.
let physical: Vec<(usize, NeedleId, i64, Size)> = entries
.iter()
.copied()
.filter(|e| !(e.2 == 0 && e.3.is_deleted()))
.collect();
for j in 1..physical.len() {
let (index, id, offset, size) = physical[j];
let (_, last_id, last_offset, last_size) = physical[j - 1];
let end = match get_actual_size(size, version) {
0 => offset,
s => offset + s - 1,
};
let last_end = match get_actual_size(last_size, version) {
0 => last_offset,
s => last_offset + s - 1,
};
if offset <= last_end {
errs.push(format!(
"needle {} (#{}) at [{}-{}] overlaps needle {} at [{}-{}]",
id.0,
index + 1,
offset,
end,
last_id.0,
last_offset,
last_end
));
}
}
let count = entries.len() as u64;
let expected = count as i64 * NEEDLE_MAP_ENTRY_SIZE as i64;
if expected != idx_file_size {
errs.push(format!(
"expected an index file of size {}, got {}",
idx_file_size, expected
));
}
(count, errs)
}
/// Write a single index entry to a writer.
pub fn write_index_entry<W: io::Write>(
writer: &mut W,
key: NeedleId,
offset: Offset,
size: Size,
) -> io::Result<()> {
let mut buf = [0u8; NEEDLE_MAP_ENTRY_SIZE];
idx_entry_to_bytes(&mut buf, key, offset, size);
writer.write_all(&buf)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn test_walk_index_file() {
// Create a small index with 3 entries
let mut data = Vec::new();
let entries = vec![
(NeedleId(1), Offset::from_actual_offset(0), Size(100)),
(NeedleId(2), Offset::from_actual_offset(128), Size(200)),
(NeedleId(3), Offset::from_actual_offset(384), Size(300)),
];
for (key, offset, size) in &entries {
let mut buf = [0u8; NEEDLE_MAP_ENTRY_SIZE];
idx_entry_to_bytes(&mut buf, *key, *offset, *size);
data.extend_from_slice(&buf);
}
let mut cursor = Cursor::new(data);
let mut collected = Vec::new();
walk_index_file(&mut cursor, 0, |key, offset, size| {
collected.push((key, offset.to_actual_offset(), size));
Ok(())
})
.unwrap();
assert_eq!(collected.len(), 3);
assert_eq!(collected[0].0, NeedleId(1));
assert_eq!(collected[0].1, 0);
assert_eq!(collected[0].2, Size(100));
assert_eq!(collected[1].0, NeedleId(2));
assert_eq!(collected[2].0, NeedleId(3));
}
#[test]
fn test_walk_empty() {
let mut cursor = Cursor::new(Vec::new());
let mut count = 0;
walk_index_file(&mut cursor, 0, |_, _, _| {
count += 1;
Ok(())
})
.unwrap();
assert_eq!(count, 0);
}
fn idx_bytes(entries: &[(NeedleId, Offset, Size)]) -> Vec<u8> {
let mut data = Vec::new();
for (key, offset, size) in entries {
let mut buf = [0u8; NEEDLE_MAP_ENTRY_SIZE];
idx_entry_to_bytes(&mut buf, *key, *offset, *size);
data.extend_from_slice(&buf);
}
data
}
#[test]
fn test_check_index_file_clean() {
let data = idx_bytes(&[
(NeedleId(1), Offset::from_actual_offset(0), Size(50)),
(NeedleId(2), Offset::from_actual_offset(100_000), Size(50)),
]);
let size = data.len() as i64;
let (count, errs) = check_index_file(&mut Cursor::new(data), size, Version(3));
assert_eq!(count, 2);
assert!(errs.is_empty(), "{:?}", errs);
}
#[test]
fn test_check_index_file_detects_overlap() {
let data = idx_bytes(&[
(NeedleId(1), Offset::from_actual_offset(0), Size(100)),
(NeedleId(2), Offset::from_actual_offset(8), Size(100)),
]);
let size = data.len() as i64;
let (count, errs) = check_index_file(&mut Cursor::new(data), size, Version(3));
assert_eq!(count, 2);
assert_eq!(errs.len(), 1, "{:?}", errs);
assert!(errs[0].contains("overlaps"), "{:?}", errs);
}
#[test]
fn test_check_index_file_ignores_offset0_tombstone() {
// A live needle near the start plus an offset-0 logical tombstone (remote-tier
// delete) must NOT be flagged as overlapping; the tombstone row still counts.
let data = idx_bytes(&[
(NeedleId(1), Offset::from_actual_offset(8), Size(100)),
(NeedleId(2), Offset::from_actual_offset(0), Size(-1)),
]);
let size = data.len() as i64;
let (count, errs) = check_index_file(&mut Cursor::new(data), size, Version(3));
assert_eq!(count, 2, "tombstone row is still counted: {:?}", errs);
assert!(
errs.is_empty(),
"offset-0 tombstone must not overlap: {:?}",
errs
);
}
#[test]
fn test_check_index_file_detects_size_mismatch() {
let data = idx_bytes(&[(NeedleId(1), Offset::from_actual_offset(0), Size(50))]);
let claimed = data.len() as i64 + 5;
let (_count, errs) = check_index_file(&mut Cursor::new(data), claimed, Version(3));
assert!(
errs.iter().any(|e| e.contains("expected an index file")),
"{:?}",
errs
);
}
#[test]
fn test_write_index_entry() {
let mut buf = Vec::new();
write_index_entry(
&mut buf,
NeedleId(42),
Offset::from_actual_offset(8 * 10),
Size(512),
)
.unwrap();
assert_eq!(buf.len(), NEEDLE_MAP_ENTRY_SIZE);
let (key, offset, size) = idx_entry_from_bytes(&buf);
assert_eq!(key, NeedleId(42));
assert_eq!(offset.to_actual_offset(), 80);
assert_eq!(size, Size(512));
}
}

View File

@@ -0,0 +1,16 @@
pub mod disk_location;
pub mod erasure_coding;
pub mod idx;
pub mod needle;
pub mod needle_map;
pub mod store;
pub mod store_ec_mirror;
pub mod store_ec_reconcile;
pub mod super_block;
pub mod types;
pub mod volume;
pub mod volume_idx_rebuild;
pub mod volume_idx_repair;
pub mod volume_open;
pub mod volume_report;
pub mod volume_report_hash;

View File

@@ -0,0 +1,75 @@
//! CRC32-Castagnoli checksum for needle data integrity.
//!
//! Matches Go's `crc32.MakeTable(crc32.Castagnoli)` exactly.
//! The CRC is stored as raw u32 (not the `.Value()` legacy transform).
/// CRC32-Castagnoli checksum wrapper.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct CRC(pub u32);
impl CRC {
/// Compute CRC from a byte slice (starting from 0).
pub fn new(data: &[u8]) -> Self {
CRC(0).update(data)
}
/// Update the CRC with additional bytes.
pub fn update(self, data: &[u8]) -> Self {
CRC(crc32c::crc32c_append(self.0, data))
}
/// Legacy `.Value()` function — deprecated in Go but needed for backward compat check.
/// Formula: (crc >> 15 | crc << 17) + 0xa282ead8
pub fn legacy_value(&self) -> u32 {
self.0.rotate_right(15).wrapping_add(0xa282ead8)
}
}
impl From<u32> for CRC {
fn from(v: u32) -> Self {
CRC(v)
}
}
impl From<CRC> for u32 {
fn from(c: CRC) -> Self {
c.0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_crc_empty() {
let crc = CRC::new(&[]);
assert_eq!(crc.0, 0);
}
#[test]
fn test_crc_known_value() {
// CRC32-C of "hello" — verify it produces a non-zero deterministic value
let crc = CRC::new(b"hello");
assert_ne!(crc.0, 0);
// Same input produces same output
assert_eq!(crc, CRC::new(b"hello"));
}
#[test]
fn test_crc_incremental() {
let crc1 = CRC::new(b"hello world");
let crc2 = CRC::new(b"hello").update(b" world");
assert_eq!(crc1, crc2);
}
#[test]
fn test_crc_legacy_value() {
let crc = CRC(0x12345678);
let v = crc.legacy_value();
// (0x12345678 >> 15 | 0x12345678 << 17) + 0xa282ead8, worked out by hand so
// the test checks the rotate rather than restating it.
let expected = 0x4f730f40_u32;
assert_eq!(v, expected);
}
}

View File

@@ -0,0 +1,11 @@
pub mod crc;
#[expect(
clippy::module_inception,
reason = "needle/needle.rs mirrors the Go package layout"
)]
pub mod needle;
pub mod ttl;
pub use crc::CRC;
pub use needle::Needle;
pub use ttl::TTL;

View File

@@ -0,0 +1,972 @@
//! Needle: the individual file object stored in a volume.
//!
//! Binary format (Version 2/3):
//! Header (16 bytes): Cookie(4) + NeedleId(8) + Size(4)
//! Body (Size bytes):
//! DataSize(4) + Data(DataSize) + Flags(1)
//! [if HasName]: NameSize(1) + Name(NameSize)
//! [if HasMime]: MimeSize(1) + Mime(MimeSize)
//! [if HasLastMod]: LastModified(5)
//! [if HasTtl]: TTL(2)
//! [if HasPairs]: PairsSize(2) + Pairs(PairsSize)
//! Tail:
//! Checksum(4) + [if V3: AppendAtNs(8)] + Padding(0-7)
use super::crc::CRC;
use super::ttl::TTL;
use crate::storage::types::*;
// Flag bits (matching Go constants)
pub const FLAG_IS_COMPRESSED: u8 = 0x01;
pub const FLAG_HAS_NAME: u8 = 0x02;
pub const FLAG_HAS_MIME: u8 = 0x04;
pub const FLAG_HAS_LAST_MODIFIED_DATE: u8 = 0x08;
pub const FLAG_HAS_TTL: u8 = 0x10;
pub const FLAG_HAS_PAIRS: u8 = 0x20;
pub const FLAG_IS_CHUNK_MANIFEST: u8 = 0x80;
pub const LAST_MODIFIED_BYTES_LENGTH: usize = 5;
pub const TTL_BYTES_LENGTH: usize = 2;
#[derive(Debug, Clone, Default)]
pub struct Needle {
pub cookie: Cookie,
pub id: NeedleId,
pub size: Size, // sum of body content fields
// Version 2+ fields
pub data_size: u32,
pub data: Vec<u8>,
pub flags: u8,
pub name_size: u8,
pub name: Vec<u8>, // max 255 bytes
pub mime_size: u8,
pub mime: Vec<u8>, // max 255 bytes
pub pairs_size: u16,
pub pairs: Vec<u8>, // max 64KB, JSON
pub last_modified: u64, // stored as 5 bytes on disk
pub ttl: Option<TTL>,
// Tail fields
pub checksum: CRC,
pub append_at_ns: u64, // Version 3 only
pub padding: Vec<u8>,
}
impl Needle {
// ---- Flag accessors (matching Go) ----
pub fn is_compressed(&self) -> bool {
self.flags & FLAG_IS_COMPRESSED != 0
}
pub fn set_is_compressed(&mut self) {
self.flags |= FLAG_IS_COMPRESSED;
}
pub fn has_name(&self) -> bool {
self.flags & FLAG_HAS_NAME != 0
}
pub fn set_has_name(&mut self) {
self.flags |= FLAG_HAS_NAME;
}
pub fn has_mime(&self) -> bool {
self.flags & FLAG_HAS_MIME != 0
}
pub fn set_has_mime(&mut self) {
self.flags |= FLAG_HAS_MIME;
}
pub fn has_last_modified_date(&self) -> bool {
self.flags & FLAG_HAS_LAST_MODIFIED_DATE != 0
}
pub fn set_has_last_modified_date(&mut self) {
self.flags |= FLAG_HAS_LAST_MODIFIED_DATE;
}
pub fn has_ttl(&self) -> bool {
self.flags & FLAG_HAS_TTL != 0
}
pub fn set_has_ttl(&mut self) {
self.flags |= FLAG_HAS_TTL;
}
pub fn has_pairs(&self) -> bool {
self.flags & FLAG_HAS_PAIRS != 0
}
pub fn set_has_pairs(&mut self) {
self.flags |= FLAG_HAS_PAIRS;
}
pub fn is_chunk_manifest(&self) -> bool {
self.flags & FLAG_IS_CHUNK_MANIFEST != 0
}
pub fn set_is_chunk_manifest(&mut self) {
self.flags |= FLAG_IS_CHUNK_MANIFEST;
}
// ---- Header parsing ----
/// Parse the 16-byte needle header.
pub fn parse_header(bytes: &[u8]) -> (Cookie, NeedleId, Size) {
assert!(bytes.len() >= NEEDLE_HEADER_SIZE);
let cookie = Cookie::from_bytes(&bytes[0..COOKIE_SIZE]);
let id = NeedleId::from_bytes(&bytes[COOKIE_SIZE..COOKIE_SIZE + NEEDLE_ID_SIZE]);
let size = Size::from_bytes(&bytes[COOKIE_SIZE + NEEDLE_ID_SIZE..NEEDLE_HEADER_SIZE]);
(cookie, id, size)
}
/// Parse needle header into self.
pub fn read_header(&mut self, bytes: &[u8]) {
let (cookie, id, size) = Self::parse_header(bytes);
self.cookie = cookie;
self.id = id;
self.size = size;
}
// ---- Body reading (Version 2/3) ----
/// Read version 2/3 body metadata only — skips copying the data payload.
/// Sets `data_size` and all metadata fields but leaves `data` empty.
pub fn read_body_v2_meta_only(&mut self, bytes: &[u8]) -> Result<(), NeedleError> {
let len_bytes = bytes.len();
let mut index = 0;
// DataSize (4 bytes)
if index + 4 > len_bytes {
return Err(NeedleError::IndexOutOfRange(1));
}
self.data_size = u32::from_be_bytes([
bytes[index],
bytes[index + 1],
bytes[index + 2],
bytes[index + 3],
]);
index += 4;
// Skip data bytes (do NOT copy them)
if index + self.data_size as usize > len_bytes {
return Err(NeedleError::IndexOutOfRange(1));
}
index += self.data_size as usize;
// Read non-data metadata
self.read_body_v2_non_data(&bytes[index..])?;
Ok(())
}
/// Read full needle from bytes but skip copying the data payload.
/// Sets all metadata fields, checksum, etc. but leaves `data` empty.
pub fn read_bytes_meta_only(
&mut self,
bytes: &[u8],
offset: i64,
expected_size: Size,
version: Version,
) -> Result<(), NeedleError> {
self.read_header(bytes);
if self.size != expected_size {
return Err(NeedleError::SizeMismatch {
offset,
id: self.id,
found: self.size,
expected: expected_size,
});
}
let body_start = NEEDLE_HEADER_SIZE;
let body_end = body_start + self.size.0 as usize;
if version == VERSION_1 {
// V1 has no metadata — data is the entire body
self.data_size = self.size.0 as u32;
} else if self.size.0 == 0 {
// Tombstones have no DataSize/body section; metadata starts at the tail.
self.data_size = 0;
} else {
self.read_body_v2_meta_only(&bytes[body_start..body_end])?;
}
// Read tail but skip CRC validation (no data to check against)
self.read_tail_meta_only(&bytes[body_end..], version)?;
Ok(())
}
/// Paged meta-only parse: accepts the 20-byte header+DataSize prefix and the
/// meta tail bytes (everything after the data payload). This avoids reading
/// the data payload from disk at all, matching Go's `ReadNeedleMeta`.
pub fn read_paged_meta(
&mut self,
header_bytes: &[u8], // first 20 bytes: NEEDLE_HEADER_SIZE + DATA_SIZE_SIZE
meta_bytes: &[u8], // tail: non-data body metadata + checksum + timestamp + padding
offset: i64,
expected_size: Size,
version: Version,
) -> Result<(), NeedleError> {
// Parse the 16-byte header
self.read_header(header_bytes);
if self.size != expected_size {
return Err(NeedleError::SizeMismatch {
offset,
id: self.id,
found: self.size,
expected: expected_size,
});
}
if version == VERSION_1 {
self.data_size = self.size.0 as u32;
} else if self.size.0 == 0 {
// Tombstone
self.data_size = 0;
} else {
// Extract DataSize from bytes 16..20
self.data_size = u32::from_be_bytes([
header_bytes[NEEDLE_HEADER_SIZE],
header_bytes[NEEDLE_HEADER_SIZE + 1],
header_bytes[NEEDLE_HEADER_SIZE + 2],
header_bytes[NEEDLE_HEADER_SIZE + 3],
]);
// meta_bytes starts with the non-data body metadata (flags, name, mime, etc.)
// followed by the tail (checksum + timestamp + padding).
// readNeedleDataVersion2NonData returns the index where it stopped.
let index = self.read_body_v2_non_data(meta_bytes)?;
self.read_tail_meta_only(&meta_bytes[index..], version)?;
return Ok(());
}
// For VERSION_1 or tombstones, meta_bytes IS the tail
self.read_tail_meta_only(meta_bytes, version)?;
Ok(())
}
/// Read tail without CRC validation (used when data was not read).
fn read_tail_meta_only(
&mut self,
tail_bytes: &[u8],
version: Version,
) -> Result<(), NeedleError> {
if tail_bytes.len() < NEEDLE_CHECKSUM_SIZE {
return Err(NeedleError::TailTooShort);
}
self.checksum = CRC(u32::from_be_bytes([
tail_bytes[0],
tail_bytes[1],
tail_bytes[2],
tail_bytes[3],
]));
if version == VERSION_3 {
let ts_offset = NEEDLE_CHECKSUM_SIZE;
if tail_bytes.len() < ts_offset + TIMESTAMP_SIZE {
return Err(NeedleError::TailTooShort);
}
self.append_at_ns = u64::from_be_bytes([
tail_bytes[ts_offset],
tail_bytes[ts_offset + 1],
tail_bytes[ts_offset + 2],
tail_bytes[ts_offset + 3],
tail_bytes[ts_offset + 4],
tail_bytes[ts_offset + 5],
tail_bytes[ts_offset + 6],
tail_bytes[ts_offset + 7],
]);
}
Ok(())
}
/// Read the version 2/3 body data from bytes (size bytes starting after header).
/// Returns IndexOutOfRange errors for truncated data (matching Go's readNeedleDataVersion2).
pub fn read_body_v2(&mut self, bytes: &[u8]) -> Result<(), NeedleError> {
let len_bytes = bytes.len();
let mut index = 0;
// DataSize (4 bytes)
if index + 4 > len_bytes {
return Ok(()); // tolerate EOF
}
self.data_size = u32::from_be_bytes([
bytes[index],
bytes[index + 1],
bytes[index + 2],
bytes[index + 3],
]);
index += 4;
// Data
if index + self.data_size as usize > len_bytes {
return Err(NeedleError::IndexOutOfRange(1));
}
self.data = bytes[index..index + self.data_size as usize].to_vec();
index += self.data_size as usize;
// Read non-data metadata
self.read_body_v2_non_data(&bytes[index..])?;
Ok(())
}
/// Read version 2/3 metadata fields (everything after Data).
/// Returns IndexOutOfRange errors for truncated data (matching Go's readNeedleDataVersion2).
fn read_body_v2_non_data(&mut self, bytes: &[u8]) -> Result<usize, NeedleError> {
let len_bytes = bytes.len();
let mut index = 0;
// Flags (1 byte)
if index < len_bytes {
self.flags = bytes[index];
index += 1;
} else {
return Ok(index);
}
// Name
if index < len_bytes && self.has_name() {
self.name_size = bytes[index];
index += 1;
if index + self.name_size as usize > len_bytes {
return Err(NeedleError::IndexOutOfRange(2));
}
self.name = bytes[index..index + self.name_size as usize].to_vec();
index += self.name_size as usize;
}
// Mime
if index < len_bytes && self.has_mime() {
self.mime_size = bytes[index];
index += 1;
if index + self.mime_size as usize > len_bytes {
return Err(NeedleError::IndexOutOfRange(3));
}
self.mime = bytes[index..index + self.mime_size as usize].to_vec();
index += self.mime_size as usize;
}
// LastModified (5 bytes)
if index < len_bytes && self.has_last_modified_date() {
if index + LAST_MODIFIED_BYTES_LENGTH > len_bytes {
return Err(NeedleError::IndexOutOfRange(4));
}
self.last_modified = bytes_to_u64_5(&bytes[index..index + LAST_MODIFIED_BYTES_LENGTH]);
index += LAST_MODIFIED_BYTES_LENGTH;
}
// TTL (2 bytes)
if index < len_bytes && self.has_ttl() {
if index + TTL_BYTES_LENGTH > len_bytes {
return Err(NeedleError::IndexOutOfRange(5));
}
self.ttl = Some(TTL::from_bytes(&bytes[index..index + TTL_BYTES_LENGTH]));
index += TTL_BYTES_LENGTH;
}
// Pairs
if index < len_bytes && self.has_pairs() {
if index + 2 > len_bytes {
return Err(NeedleError::IndexOutOfRange(6));
}
self.pairs_size = u16::from_be_bytes([bytes[index], bytes[index + 1]]);
index += 2;
if index + self.pairs_size as usize > len_bytes {
return Err(NeedleError::IndexOutOfRange(7));
}
self.pairs = bytes[index..index + self.pairs_size as usize].to_vec();
index += self.pairs_size as usize;
}
Ok(index)
}
// ---- Tail reading ----
/// Read the needle tail (checksum + optional timestamp + padding).
pub fn read_tail(&mut self, tail_bytes: &[u8], version: Version) -> Result<(), NeedleError> {
if tail_bytes.len() < NEEDLE_CHECKSUM_SIZE {
return Err(NeedleError::TailTooShort);
}
let expected_checksum = CRC(u32::from_be_bytes([
tail_bytes[0],
tail_bytes[1],
tail_bytes[2],
tail_bytes[3],
]));
if !self.data.is_empty() {
let data_checksum = CRC::new(&self.data);
// Go double-checks: n.Checksum != crc && uint32(n.Checksum) != crc.Value()
// The crc.Value() path is a deprecated legacy transform for backward compat
// with seaweed versions prior to commit 056c480eb.
if expected_checksum != data_checksum
&& expected_checksum.0 != data_checksum.legacy_value()
{
return Err(NeedleError::CrcMismatch {
needle_id: self.id,
got: data_checksum.0,
want: expected_checksum.0,
});
}
self.checksum = data_checksum;
} else {
self.checksum = expected_checksum;
}
if version == VERSION_3 {
let ts_offset = NEEDLE_CHECKSUM_SIZE;
if tail_bytes.len() < ts_offset + TIMESTAMP_SIZE {
return Err(NeedleError::TailTooShort);
}
self.append_at_ns = u64::from_be_bytes([
tail_bytes[ts_offset],
tail_bytes[ts_offset + 1],
tail_bytes[ts_offset + 2],
tail_bytes[ts_offset + 3],
tail_bytes[ts_offset + 4],
tail_bytes[ts_offset + 5],
tail_bytes[ts_offset + 6],
tail_bytes[ts_offset + 7],
]);
}
Ok(())
}
// ---- Full read from bytes ----
/// Read a complete needle from its raw bytes (header + body + tail).
pub fn read_bytes(
&mut self,
bytes: &[u8],
offset: i64,
expected_size: Size,
version: Version,
) -> Result<(), NeedleError> {
self.read_header(bytes);
if self.size != expected_size {
return Err(NeedleError::SizeMismatch {
offset,
id: self.id,
found: self.size,
expected: expected_size,
});
}
let body_start = NEEDLE_HEADER_SIZE;
let body_end = body_start + self.size.0 as usize;
if version == VERSION_1 {
self.data = bytes[body_start..body_end].to_vec();
} else {
self.read_body_v2(&bytes[body_start..body_end])?;
}
self.read_tail(&bytes[body_end..], version)?;
Ok(())
}
// ---- Write (serialize) ----
/// Serialize the needle to bytes for writing to a .dat file (Version 2/3).
pub fn write_bytes(&mut self, version: Version) -> Vec<u8> {
let mut buf = Vec::with_capacity(256);
// Compute sizes (matching Go writeNeedleCommon)
if self.name.len() >= 255 {
self.name_size = 255;
} else {
self.name_size = self.name.len() as u8;
}
self.data_size = self.data.len() as u32;
self.mime_size = self.mime.len() as u8;
// Compute n.Size (body size, excluding header)
if self.data_size > 0 {
let mut s: i32 = 4 + self.data_size as i32 + 1; // DataSize + Data + Flags
if self.has_name() {
s += 1 + self.name_size as i32;
}
if self.has_mime() {
s += 1 + self.mime_size as i32;
}
if self.has_last_modified_date() {
s += LAST_MODIFIED_BYTES_LENGTH as i32;
}
if self.has_ttl() {
s += TTL_BYTES_LENGTH as i32;
}
if self.has_pairs() {
s += 2 + self.pairs_size as i32;
}
self.size = Size(s);
} else {
self.size = Size(0);
}
// Header: Cookie(4) + NeedleId(8) + Size(4) = 16 bytes
let mut header = [0u8; NEEDLE_HEADER_SIZE];
self.cookie.to_bytes(&mut header[0..COOKIE_SIZE]);
self.id
.to_bytes(&mut header[COOKIE_SIZE..COOKIE_SIZE + NEEDLE_ID_SIZE]);
self.size
.to_bytes(&mut header[COOKIE_SIZE + NEEDLE_ID_SIZE..NEEDLE_HEADER_SIZE]);
buf.extend_from_slice(&header);
// Body
if self.data_size > 0 {
buf.extend_from_slice(&self.data_size.to_be_bytes());
buf.extend_from_slice(&self.data);
buf.push(self.flags);
if self.has_name() {
buf.push(self.name_size);
buf.extend_from_slice(&self.name[..self.name_size as usize]);
}
if self.has_mime() {
buf.push(self.mime_size);
buf.extend_from_slice(&self.mime);
}
if self.has_last_modified_date() {
// Write 5 bytes of last_modified (lower 5 bytes of u64 big-endian)
let lm_bytes = self.last_modified.to_be_bytes();
buf.extend_from_slice(&lm_bytes[8 - LAST_MODIFIED_BYTES_LENGTH..8]);
}
if self.has_ttl() {
if let Some(ref ttl) = self.ttl {
let mut ttl_buf = [0u8; 2];
ttl.to_bytes(&mut ttl_buf);
buf.extend_from_slice(&ttl_buf);
} else {
buf.extend_from_slice(&[0u8; 2]);
}
}
if self.has_pairs() {
buf.extend_from_slice(&self.pairs_size.to_be_bytes());
buf.extend_from_slice(&self.pairs);
}
}
// Compute checksum
self.checksum = CRC::new(&self.data);
// Tail: Checksum + [V3: AppendAtNs] + Padding
buf.extend_from_slice(&self.checksum.0.to_be_bytes());
if version == VERSION_3 {
buf.extend_from_slice(&self.append_at_ns.to_be_bytes());
}
// Padding to 8-byte alignment
let padding = padding_length(self.size, version).0 as usize;
buf.extend(std::iter::repeat_n(0u8, padding));
buf
}
/// Total disk size of this needle including header, body, checksum, timestamp, and padding.
pub fn disk_size(&self, version: Version) -> i64 {
get_actual_size(self.size, version)
}
/// Compute ETag string from checksum (matching Go).
pub fn etag(&self) -> String {
etag_from_checksum(self.checksum.0)
}
}
// ============================================================================
// Helper functions (matching Go)
// ============================================================================
/// Compute padding to align needle to NEEDLE_PADDING_SIZE (8 bytes).
pub fn padding_length(needle_size: Size, version: Version) -> Size {
if version == VERSION_3 {
Size(
NEEDLE_PADDING_SIZE as i32
- ((NEEDLE_HEADER_SIZE as i32
+ needle_size.0
+ NEEDLE_CHECKSUM_SIZE as i32
+ TIMESTAMP_SIZE as i32)
% NEEDLE_PADDING_SIZE as i32),
)
} else {
Size(
NEEDLE_PADDING_SIZE as i32
- ((NEEDLE_HEADER_SIZE as i32 + needle_size.0 + NEEDLE_CHECKSUM_SIZE as i32)
% NEEDLE_PADDING_SIZE as i32),
)
}
}
/// Body length = Size + Checksum + [Timestamp] + Padding.
pub fn needle_body_length(needle_size: Size, version: Version) -> i64 {
if version == VERSION_3 {
needle_size.0 as i64
+ NEEDLE_CHECKSUM_SIZE as i64
+ TIMESTAMP_SIZE as i64
+ padding_length(needle_size, version).0 as i64
} else {
needle_size.0 as i64
+ NEEDLE_CHECKSUM_SIZE as i64
+ padding_length(needle_size, version).0 as i64
}
}
/// Total actual size on disk: Header + Body.
pub fn get_actual_size(size: Size, version: Version) -> i64 {
NEEDLE_HEADER_SIZE as i64 + needle_body_length(size, version)
}
/// Read 5 bytes as a u64 (big-endian, zero-padded high bytes).
fn bytes_to_u64_5(bytes: &[u8]) -> u64 {
assert!(bytes.len() >= 5);
// The 5 bytes are the LOWER 5 bytes of a big-endian u64.
// In Go: util.BytesToUint64(bytes[index : index+5]) reads into a uint64
// Go's BytesToUint64 copies into the LAST 5 bytes of an 8-byte array (big-endian).
let mut buf = [0u8; 8];
buf[3..8].copy_from_slice(&bytes[..5]);
u64::from_be_bytes(buf)
}
/// ETag formatted as Go: hex of big-endian u32 bytes.
pub fn etag_from_checksum(checksum: u32) -> String {
let bits = checksum.to_be_bytes();
format!(
"{:02x}{:02x}{:02x}{:02x}",
bits[0], bits[1], bits[2], bits[3]
)
}
// ============================================================================
// FileId
// ============================================================================
/// FileId = VolumeId + NeedleId + Cookie.
/// String format: "<volume_id>,<needle_id_hex><cookie_hex>"
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FileId {
pub volume_id: VolumeId,
pub key: NeedleId,
pub cookie: Cookie,
}
impl FileId {
pub fn new(volume_id: VolumeId, key: NeedleId, cookie: Cookie) -> Self {
FileId {
volume_id,
key,
cookie,
}
}
/// Parse "volume_id,needle_id_cookie" or "volume_id/needle_id_cookie".
pub fn parse(s: &str) -> Result<Self, String> {
let (vid_str, rest) = if let Some(pos) = s.find(',') {
(&s[..pos], &s[pos + 1..])
} else if let Some(pos) = s.find('/') {
(&s[..pos], &s[pos + 1..])
} else {
return Err(format!("invalid file id: {}", s));
};
let volume_id =
VolumeId::parse(vid_str).map_err(|e| format!("invalid volume id: {}", e))?;
let (key, cookie) = parse_needle_id_cookie(rest)?;
Ok(FileId {
volume_id,
key,
cookie,
})
}
/// Format the needle_id + cookie part as a hex string (stripping leading zeros).
pub fn needle_id_cookie_string(&self) -> String {
format_needle_id_cookie(self.key, self.cookie)
}
}
impl std::fmt::Display for FileId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{},{}", self.volume_id, self.needle_id_cookie_string())
}
}
/// Format NeedleId + Cookie as hex, stripping leading zero bytes from NeedleId only.
/// Matches Go: strips leading zero bytes up to NeedleIdSize (8), so cookie is always present.
fn format_needle_id_cookie(key: NeedleId, cookie: Cookie) -> String {
// Encode 12 bytes: 8 for NeedleId + 4 for Cookie
let mut bytes = [0u8; 12];
key.to_bytes(&mut bytes[0..8]);
cookie.to_bytes(&mut bytes[8..12]);
// Strip leading zero bytes, but only within NeedleId portion (first 8 bytes)
let mut nonzero_index = 0;
while nonzero_index < NEEDLE_ID_SIZE && bytes[nonzero_index] == 0 {
nonzero_index += 1;
}
hex::encode(&bytes[nonzero_index..])
}
/// Parse "needle_id_cookie_hex" or "needle_id_cookie_hex_delta" into (NeedleId, Cookie).
/// Matches Go's ParsePath + ParseNeedleIdCookie: supports an optional `_delta` suffix
/// where delta is a decimal number added to the NeedleId (used for sub-file addressing).
/// Rejects strings that are too short or too long, and deltas that overflow the needle id.
pub fn parse_needle_id_cookie(s: &str) -> Result<(NeedleId, Cookie), String> {
// Go ParsePath: check for "_" suffix containing a decimal delta
let (hex_part, delta) = if let Some(underscore_pos) = s.rfind('_') {
if underscore_pos > 0 {
let delta_str = &s[underscore_pos + 1..];
let d: u64 = delta_str
.parse()
.map_err(|e| format!("Parse delta error: {}", e))?;
(&s[..underscore_pos], Some(d))
} else {
(s, None)
}
} else {
(s, None)
};
// Go: len(key_hash_string) <= CookieSize*2 => error (must be > 8 hex chars)
if hex_part.len() <= COOKIE_SIZE * 2 {
return Err("KeyHash is too short.".to_string());
}
// Go: len(key_hash_string) > (NeedleIdSize+CookieSize)*2 => error (must be <= 24 hex chars)
if hex_part.len() > (NEEDLE_ID_SIZE + COOKIE_SIZE) * 2 {
return Err("KeyHash is too long.".to_string());
}
// Split: last CookieSize*2 hex chars are cookie, rest is needle id
let split = hex_part.len() - COOKIE_SIZE * 2;
let needle_id_hex = &hex_part[..split];
let cookie_hex = &hex_part[split..];
// Go parses both with strconv.ParseUint, which accepts odd-length hex
let mut key = NeedleId(
u64::from_str_radix(needle_id_hex, 16)
.map_err(|e| format!("Parse needleId error: {}", e))?,
);
let cookie = Cookie(
u32::from_str_radix(cookie_hex, 16).map_err(|e| format!("Parse cookie error: {}", e))?,
);
// Apply delta if present (Go: n.Id += Uint64ToNeedleId(d))
if let Some(d) = delta {
key = NeedleId(
key.0
.checked_add(d)
.ok_or_else(|| format!("delta {} overflows needle id {:x}", d, key.0))?,
);
}
Ok((key, cookie))
}
// ============================================================================
// Errors
// ============================================================================
#[derive(Debug, thiserror::Error)]
pub enum NeedleError {
#[error(
"size mismatch at offset {offset}: found id={id} size={found:?}, expected size={expected:?}"
)]
SizeMismatch {
offset: i64,
id: NeedleId,
found: Size,
expected: Size,
},
#[error("CRC mismatch for needle {needle_id}: got {got:08x}, want {want:08x}")]
CrcMismatch {
needle_id: NeedleId,
got: u32,
want: u32,
},
#[error("index out of range ({0})")]
IndexOutOfRange(u32),
#[error("needle tail too short")]
TailTooShort,
#[error("unsupported version: {0}")]
UnsupportedVersion(u8),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_header() {
let mut buf = [0u8; NEEDLE_HEADER_SIZE];
let cookie = Cookie(0xdeadbeef);
let id = NeedleId(0x123456789abcdef0);
let size = Size(1024);
cookie.to_bytes(&mut buf[0..4]);
id.to_bytes(&mut buf[4..12]);
size.to_bytes(&mut buf[12..16]);
let (c, i, s) = Needle::parse_header(&buf);
assert_eq!(c, cookie);
assert_eq!(i, id);
assert_eq!(s, size);
}
#[test]
fn test_needle_write_read_round_trip_v3() {
let mut n = Needle {
cookie: Cookie(42),
id: NeedleId(100),
data: b"hello world".to_vec(),
flags: 0,
..Needle::default()
};
n.set_has_name();
n.name = b"test.txt".to_vec();
n.name_size = 8;
n.set_has_mime();
n.mime = b"text/plain".to_vec();
n.mime_size = 10;
n.set_has_last_modified_date();
n.last_modified = 1234567890;
n.set_has_ttl();
n.ttl = Some(TTL {
count: 5,
unit: super::super::ttl::TTL_UNIT_DAY,
});
n.append_at_ns = 999999999;
let bytes = n.write_bytes(VERSION_3);
// Verify total size matches expected
let expected_size = get_actual_size(n.size, VERSION_3);
assert_eq!(bytes.len() as i64, expected_size);
// Read it back
let mut n2 = Needle::default();
n2.read_bytes(&bytes, 0, n.size, VERSION_3).unwrap();
assert_eq!(n2.cookie, n.cookie);
assert_eq!(n2.id, n.id);
assert_eq!(n2.data, n.data);
assert_eq!(n2.name, n.name);
assert_eq!(n2.mime, n.mime);
assert_eq!(n2.last_modified, n.last_modified);
assert_eq!(n2.ttl, n.ttl);
assert_eq!(n2.checksum, n.checksum);
assert_eq!(n2.append_at_ns, n.append_at_ns);
}
#[test]
fn test_needle_write_read_round_trip_v2() {
let mut n = Needle {
cookie: Cookie(77),
id: NeedleId(200),
data: b"data v2".to_vec(),
flags: 0,
..Needle::default()
};
let bytes = n.write_bytes(VERSION_2);
let expected_size = get_actual_size(n.size, VERSION_2);
assert_eq!(bytes.len() as i64, expected_size);
let mut n2 = Needle::default();
n2.read_bytes(&bytes, 0, n.size, VERSION_2).unwrap();
assert_eq!(n2.data, n.data);
assert_eq!(n2.checksum, n.checksum);
}
#[test]
fn test_read_bytes_meta_only_handles_tombstone_v3() {
let mut tombstone = Needle {
cookie: Cookie(0x1234abcd),
id: NeedleId(300),
append_at_ns: 999_999,
..Needle::default()
};
let bytes = tombstone.write_bytes(VERSION_3);
let mut meta = Needle::default();
meta.read_bytes_meta_only(&bytes, 0, Size(0), VERSION_3)
.unwrap();
assert_eq!(meta.cookie, tombstone.cookie);
assert_eq!(meta.id, tombstone.id);
assert_eq!(meta.size, Size(0));
assert_eq!(meta.data_size, 0);
assert_eq!(meta.append_at_ns, tombstone.append_at_ns);
assert_eq!(meta.checksum, tombstone.checksum);
}
#[test]
fn test_padding_alignment() {
// All actual sizes should be multiples of 8
for size_val in 0..50 {
let s = Size(size_val);
let actual_v2 = get_actual_size(s, VERSION_2);
let actual_v3 = get_actual_size(s, VERSION_3);
assert_eq!(actual_v2 % 8, 0, "V2 size {} not aligned", size_val);
assert_eq!(actual_v3 % 8, 0, "V3 size {} not aligned", size_val);
}
}
#[test]
fn test_file_id_parse() {
let fid = FileId::parse("3,01637037d6").unwrap();
assert_eq!(fid.volume_id, VolumeId(3));
// The hex "01637037d6" is 5 bytes = 0x0163..., padded to 12 bytes
assert!(!fid.key.is_empty() || !fid.cookie.0 == 0);
}
#[test]
fn test_file_id_round_trip() {
let fid = FileId::new(VolumeId(5), NeedleId(0x123456), Cookie(0xabcd));
let s = fid.to_string();
let fid2 = FileId::parse(&s).unwrap();
assert_eq!(fid, fid2);
}
#[test]
fn test_needle_id_cookie_format() {
let s = format_needle_id_cookie(NeedleId(1), Cookie(0x12345678));
let (key, cookie) = parse_needle_id_cookie(&s).unwrap();
assert_eq!(key, NeedleId(1));
assert_eq!(cookie, Cookie(0x12345678));
}
#[test]
fn test_parse_rejects_delta_overflow() {
assert!(parse_needle_id_cookie("ffffffffffffffff00000000_1").is_err());
let (key, _) = parse_needle_id_cookie("fffffffffffffffe00000000_1").unwrap();
assert_eq!(key, NeedleId(u64::MAX));
}
#[test]
fn test_parse_odd_length_needle_id() {
// Go's shell formats purge fids with strconv.FormatUint, which does not
// pad the needle id hex to an even length.
let (key, cookie) = parse_needle_id_cookie("100000000").unwrap();
assert_eq!(key, NeedleId(1));
assert_eq!(cookie, Cookie(0));
let fid = FileId::parse("3,12300000000").unwrap();
assert_eq!(fid.volume_id, VolumeId(3));
assert_eq!(fid.key, NeedleId(0x123));
assert_eq!(fid.cookie, Cookie(0));
}
}

View File

@@ -0,0 +1,383 @@
//! Time-to-live encoding for needles.
//!
//! TTL is stored as 2 bytes: Count(1) + Unit(1).
//! Supported units: minute(m), hour(h), day(d), week(w), month(M), year(y).
use std::fmt;
/// TTL unit constants (matching Go).
pub const TTL_UNIT_EMPTY: u8 = 0;
pub const TTL_UNIT_MINUTE: u8 = 1;
pub const TTL_UNIT_HOUR: u8 = 2;
pub const TTL_UNIT_DAY: u8 = 3;
pub const TTL_UNIT_WEEK: u8 = 4;
pub const TTL_UNIT_MONTH: u8 = 5;
pub const TTL_UNIT_YEAR: u8 = 6;
pub const TTL_BYTES_LENGTH: usize = 2;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct TTL {
pub count: u8,
pub unit: u8,
}
impl TTL {
pub const EMPTY: TTL = TTL { count: 0, unit: 0 };
pub fn is_empty(&self) -> bool {
self.count == 0 && self.unit == 0
}
/// Load from 2 bytes.
pub fn from_bytes(input: &[u8]) -> Self {
if input.len() < 2 {
return TTL::EMPTY;
}
TTL {
count: input[0],
unit: input[1],
}
}
/// Serialize to 2 bytes.
pub fn to_bytes(&self, output: &mut [u8]) {
assert!(output.len() >= 2);
output[0] = self.count;
output[1] = self.unit;
}
/// Encode as u32: (count << 8) | unit.
pub fn to_u32(&self) -> u32 {
if self.count == 0 {
return 0;
}
((self.count as u32) << 8) + (self.unit as u32)
}
/// Decode from u32.
pub fn from_u32(v: u32) -> Self {
if v == 0 {
return TTL::EMPTY;
}
TTL {
count: (v >> 8) as u8,
unit: (v & 0xFF) as u8,
}
}
/// Convert to total seconds.
pub fn to_seconds(&self) -> u64 {
unit_to_seconds(self.count as u64, self.unit)
}
/// Parse from string like "3m", "4h", "5d", "6w", "7M", "8y".
/// If the string is all digits (no unit suffix), defaults to minutes.
/// Matches Go's ReadTTL which calls fitTtlCount to normalize:
/// e.g. "120m" -> 2h, "7d" -> 1w, "24h" -> 1d.
pub fn read(s: &str) -> Result<Self, String> {
let s = s.trim();
if s.is_empty() {
return Ok(TTL::EMPTY);
}
let last_byte = s.as_bytes()[s.len() - 1];
let (num_str, unit_byte) = if last_byte.is_ascii_digit() {
// All digits — default to minutes (matching Go)
(s, b'm')
} else {
(&s[..s.len() - 1], last_byte)
};
let count: u32 = num_str
.parse()
.map_err(|e| format!("invalid TTL count: {}", e))?;
let unit = match unit_byte {
b'm' => TTL_UNIT_MINUTE,
b'h' => TTL_UNIT_HOUR,
b'd' => TTL_UNIT_DAY,
b'w' => TTL_UNIT_WEEK,
b'M' => TTL_UNIT_MONTH,
b'y' => TTL_UNIT_YEAR,
_ => return Err(format!("unknown TTL unit: {}", unit_byte as char)),
};
// Match Go's ReadTTL: normalize via fitTtlCount
Ok(fit_ttl_count(count, unit))
}
/// Minutes representation.
pub fn minutes(&self) -> u32 {
(self.to_seconds() / 60) as u32
}
}
fn unit_to_seconds(count: u64, unit: u8) -> u64 {
match unit {
TTL_UNIT_EMPTY => 0,
TTL_UNIT_MINUTE => count * 60,
TTL_UNIT_HOUR => count * 60 * 60,
TTL_UNIT_DAY => count * 60 * 60 * 24,
TTL_UNIT_WEEK => count * 60 * 60 * 24 * 7,
TTL_UNIT_MONTH => count * 60 * 60 * 24 * 30,
TTL_UNIT_YEAR => count * 60 * 60 * 24 * 365,
_ => 0,
}
}
/// Fit a count+unit into a TTL that fits in a single byte count.
/// Converts to seconds first, then finds the coarsest unit that fits.
/// Matches Go's fitTtlCount called from ReadTTL.
fn fit_ttl_count(count: u32, unit: u8) -> TTL {
if count == 0 || unit == TTL_UNIT_EMPTY {
return TTL::EMPTY;
}
// Always convert to seconds and normalize (matches Go).
let seconds = unit_to_seconds(count as u64, unit);
if seconds == 0 {
return TTL::EMPTY;
}
const YEAR_SECS: u64 = 3600 * 24 * 365;
const MONTH_SECS: u64 = 3600 * 24 * 30;
const WEEK_SECS: u64 = 3600 * 24 * 7;
const DAY_SECS: u64 = 3600 * 24;
const HOUR_SECS: u64 = 3600;
const MINUTE_SECS: u64 = 60;
// First pass: try exact fits from largest to smallest
if seconds.is_multiple_of(YEAR_SECS) && seconds / YEAR_SECS < 256 {
return TTL {
count: (seconds / YEAR_SECS) as u8,
unit: TTL_UNIT_YEAR,
};
}
if seconds.is_multiple_of(MONTH_SECS) && seconds / MONTH_SECS < 256 {
return TTL {
count: (seconds / MONTH_SECS) as u8,
unit: TTL_UNIT_MONTH,
};
}
if seconds.is_multiple_of(WEEK_SECS) && seconds / WEEK_SECS < 256 {
return TTL {
count: (seconds / WEEK_SECS) as u8,
unit: TTL_UNIT_WEEK,
};
}
if seconds.is_multiple_of(DAY_SECS) && seconds / DAY_SECS < 256 {
return TTL {
count: (seconds / DAY_SECS) as u8,
unit: TTL_UNIT_DAY,
};
}
if seconds.is_multiple_of(HOUR_SECS) && seconds / HOUR_SECS < 256 {
return TTL {
count: (seconds / HOUR_SECS) as u8,
unit: TTL_UNIT_HOUR,
};
}
// Minutes: truncating division
if seconds / MINUTE_SECS < 256 {
return TTL {
count: (seconds / MINUTE_SECS) as u8,
unit: TTL_UNIT_MINUTE,
};
}
// Second pass: truncating division from smallest to largest
if seconds / HOUR_SECS < 256 {
return TTL {
count: (seconds / HOUR_SECS) as u8,
unit: TTL_UNIT_HOUR,
};
}
if seconds / DAY_SECS < 256 {
return TTL {
count: (seconds / DAY_SECS) as u8,
unit: TTL_UNIT_DAY,
};
}
if seconds / WEEK_SECS < 256 {
return TTL {
count: (seconds / WEEK_SECS) as u8,
unit: TTL_UNIT_WEEK,
};
}
if seconds / MONTH_SECS < 256 {
return TTL {
count: (seconds / MONTH_SECS) as u8,
unit: TTL_UNIT_MONTH,
};
}
if seconds / YEAR_SECS < 256 {
return TTL {
count: (seconds / YEAR_SECS) as u8,
unit: TTL_UNIT_YEAR,
};
}
TTL::EMPTY
}
fn unit_to_char(unit: u8) -> char {
match unit {
TTL_UNIT_MINUTE => 'm',
TTL_UNIT_HOUR => 'h',
TTL_UNIT_DAY => 'd',
TTL_UNIT_WEEK => 'w',
TTL_UNIT_MONTH => 'M',
TTL_UNIT_YEAR => 'y',
_ => ' ',
}
}
impl fmt::Display for TTL {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.count == 0 || self.unit == TTL_UNIT_EMPTY {
return write!(f, "");
}
write!(f, "{}{}", self.count, unit_to_char(self.unit))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ttl_parse() {
let ttl = TTL::read("3m").unwrap();
assert_eq!(
ttl,
TTL {
count: 3,
unit: TTL_UNIT_MINUTE
}
);
assert_eq!(ttl.to_seconds(), 180);
}
#[test]
fn test_ttl_parse_hours() {
// 24h normalizes to 1d via fitTtlCount
let ttl = TTL::read("24h").unwrap();
assert_eq!(ttl.to_seconds(), 86400);
assert_eq!(
ttl,
TTL {
count: 1,
unit: TTL_UNIT_DAY
}
);
}
#[test]
fn test_ttl_display() {
let ttl = TTL {
count: 5,
unit: TTL_UNIT_DAY,
};
assert_eq!(ttl.to_string(), "5d");
}
#[test]
fn test_ttl_bytes_round_trip() {
let ttl = TTL {
count: 10,
unit: TTL_UNIT_WEEK,
};
let mut buf = [0u8; 2];
ttl.to_bytes(&mut buf);
let ttl2 = TTL::from_bytes(&buf);
assert_eq!(ttl, ttl2);
}
#[test]
fn test_ttl_u32_round_trip() {
let ttl = TTL {
count: 42,
unit: TTL_UNIT_HOUR,
};
let v = ttl.to_u32();
let ttl2 = TTL::from_u32(v);
assert_eq!(ttl, ttl2);
}
#[test]
fn test_ttl_empty() {
assert!(TTL::EMPTY.is_empty());
assert_eq!(TTL::EMPTY.to_seconds(), 0);
assert_eq!(TTL::EMPTY.to_u32(), 0);
}
#[test]
fn test_ttl_overflow_normalizes() {
// Go's ReadTTL calls fitTtlCount: 300m = 18000s = 5h (exact fit)
let ttl = TTL::read("300m").unwrap();
assert_eq!(
ttl,
TTL {
count: 5,
unit: TTL_UNIT_HOUR
}
);
// 256h = 921600s. Doesn't fit in hours (256 >= 256), doesn't fit exact in days.
// Second pass: 921600/86400 = 10 (truncated) < 256 -> 10d
let ttl = TTL::read("256h").unwrap();
assert_eq!(
ttl,
TTL {
count: 10,
unit: TTL_UNIT_DAY
}
);
}
#[test]
fn test_ttl_normalizes_unit() {
// Go's ReadTTL calls fitTtlCount which normalizes to coarsest unit.
// 120m -> 2h, 7d -> 1w, 24h -> 1d.
let ttl = TTL::read("120m").unwrap();
assert_eq!(
ttl,
TTL {
count: 2,
unit: TTL_UNIT_HOUR
}
);
let ttl = TTL::read("7d").unwrap();
assert_eq!(
ttl,
TTL {
count: 1,
unit: TTL_UNIT_WEEK
}
);
let ttl = TTL::read("24h").unwrap();
assert_eq!(
ttl,
TTL {
count: 1,
unit: TTL_UNIT_DAY
}
);
// Values that don't simplify stay as-is
let ttl = TTL::read("5d").unwrap();
assert_eq!(
ttl,
TTL {
count: 5,
unit: TTL_UNIT_DAY
}
);
let ttl = TTL::read("3m").unwrap();
assert_eq!(
ttl,
TTL {
count: 3,
unit: TTL_UNIT_MINUTE
}
);
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,375 @@
//! CompactMap: memory-efficient in-memory map of NeedleId -> (Offset, Size).
//!
//! Port of Go's CompactMap from weed/storage/needle_map/compact_map.go.
//! Uses segmented sorted arrays with compressed keys (u16 instead of u64)
//! to achieve ~10 bytes per entry instead of ~40-48 bytes with HashMap.
//!
//! NeedleId is split into: chunk = id / SEGMENT_CHUNK_SIZE, compact_key = id % SEGMENT_CHUNK_SIZE.
//! Each segment stores up to SEGMENT_CHUNK_SIZE entries in a sorted Vec, searched via binary search.
//! Best case (ordered inserts): O(1). Worst case: O(log n) per segment.
use std::collections::HashMap;
use super::NeedleValue;
use crate::storage::types::*;
/// Maximum entries per segment. Must be <= u16::MAX (65535).
const SEGMENT_CHUNK_SIZE: u64 = 50_000;
/// Compact key: only the low bits of NeedleId within a segment.
type CompactKey = u16;
/// Segment chunk identifier: NeedleId / SEGMENT_CHUNK_SIZE.
type Chunk = u64;
/// Compact entry: 10 bytes (2 + 4 + 4) vs 16 bytes for full NeedleId + NeedleValue.
#[derive(Clone, Copy)]
struct CompactEntry {
key: CompactKey, // 2 bytes
offset: [u8; OFFSET_SIZE], // 4 bytes
size: Size, // 4 bytes
}
impl CompactEntry {
fn to_needle_value(self) -> NeedleValue {
NeedleValue {
offset: Offset::from_bytes(&self.offset),
size: self.size,
}
}
}
/// A sorted segment of compact entries for a given chunk.
struct Segment {
list: Vec<CompactEntry>,
chunk: Chunk,
first_key: CompactKey,
last_key: CompactKey,
}
impl Segment {
fn new(chunk: Chunk) -> Self {
Segment {
list: Vec::new(),
chunk,
first_key: u16::MAX,
last_key: 0,
}
}
fn compact_key(&self, id: NeedleId) -> CompactKey {
(id.0 - SEGMENT_CHUNK_SIZE * self.chunk) as CompactKey
}
/// Binary search for a compact key. Returns (index, found).
/// If not found, index is the insertion point.
fn bsearch(&self, id: NeedleId) -> (usize, bool) {
let ck = self.compact_key(id);
if self.list.is_empty() {
return (0, false);
}
if ck == self.first_key {
return (0, true);
}
if ck < self.first_key {
return (0, false);
}
if ck == self.last_key {
return (self.list.len() - 1, true);
}
if ck > self.last_key {
return (self.list.len(), false);
}
let i = self.list.partition_point(|e| e.key < ck);
if i < self.list.len() && self.list[i].key == ck {
(i, true)
} else {
(i, false)
}
}
/// Insert or update. Returns old NeedleValue if updating.
fn set(&mut self, id: NeedleId, offset: Offset, size: Size) -> Option<NeedleValue> {
let (i, found) = self.bsearch(id);
if found {
let old = self.list[i].to_needle_value();
let mut offset_bytes = [0u8; OFFSET_SIZE];
offset.to_bytes(&mut offset_bytes);
self.list[i].offset = offset_bytes;
self.list[i].size = size;
return Some(old);
}
// Insert at sorted position
let ck = self.compact_key(id);
let mut offset_bytes = [0u8; OFFSET_SIZE];
offset.to_bytes(&mut offset_bytes);
let entry = CompactEntry {
key: ck,
offset: offset_bytes,
size,
};
// Match Go panic: don't exceed segment capacity
if self.list.len() >= SEGMENT_CHUNK_SIZE as usize {
panic!(
"attempted to write more than {} entries on CompactMapSegment",
SEGMENT_CHUNK_SIZE
);
}
if self.list.len() == SEGMENT_CHUNK_SIZE as usize - 1 {
// Pin capacity to exact size when maxing out
let mut new_list = Vec::with_capacity(SEGMENT_CHUNK_SIZE as usize);
new_list.extend_from_slice(&self.list[..i]);
new_list.push(entry);
new_list.extend_from_slice(&self.list[i..]);
self.list = new_list;
} else {
self.list.insert(i, entry);
}
if ck < self.first_key {
self.first_key = ck;
}
if ck > self.last_key {
self.last_key = ck;
}
None
}
fn get(&self, id: NeedleId) -> Option<NeedleValue> {
let (i, found) = self.bsearch(id);
if found {
Some(self.list[i].to_needle_value())
} else {
None
}
}
/// Mark as deleted by negating size. Returns previous size if not already deleted.
/// Matches Go behavior: checks !IsDeleted() (i.e., size >= 0).
fn delete(&mut self, id: NeedleId) -> Option<Size> {
let (i, found) = self.bsearch(id);
if found && !self.list[i].size.is_deleted() {
let old_size = self.list[i].size;
if self.list[i].size.0 == 0 {
self.list[i].size = TOMBSTONE_FILE_SIZE;
} else {
self.list[i].size = Size(-self.list[i].size.0);
}
Some(old_size)
} else {
None
}
}
}
/// Memory-efficient map of NeedleId -> (Offset, Size).
/// Segments NeedleIds into chunks of 50,000 and stores compact 10-byte entries
/// in sorted arrays, using only 2 bytes for the key within each segment.
pub struct CompactMap {
segments: HashMap<Chunk, Segment>,
}
impl CompactMap {
pub fn new() -> Self {
CompactMap {
segments: HashMap::new(),
}
}
fn _segment_for_key(&mut self, id: NeedleId) -> &mut Segment {
let chunk = id.0 / SEGMENT_CHUNK_SIZE;
self.segments
.entry(chunk)
.or_insert_with(|| Segment::new(chunk))
}
/// Insert or update. Returns old NeedleValue if updating.
pub fn set(&mut self, id: NeedleId, offset: Offset, size: Size) -> Option<NeedleValue> {
let chunk = id.0 / SEGMENT_CHUNK_SIZE;
let segment = self
.segments
.entry(chunk)
.or_insert_with(|| Segment::new(chunk));
segment.set(id, offset, size)
}
pub fn get(&self, id: NeedleId) -> Option<NeedleValue> {
let chunk = id.0 / SEGMENT_CHUNK_SIZE;
self.segments.get(&chunk)?.get(id)
}
/// Mark as deleted. Returns previous size if was valid.
pub fn delete(&mut self, id: NeedleId) -> Option<Size> {
let chunk = id.0 / SEGMENT_CHUNK_SIZE;
self.segments.get_mut(&chunk)?.delete(id)
}
/// Remove entry entirely (used during idx loading).
pub fn remove(&mut self, id: NeedleId) -> Option<NeedleValue> {
let chunk = id.0 / SEGMENT_CHUNK_SIZE;
let segment = self.segments.get_mut(&chunk)?;
let (i, found) = segment.bsearch(id);
if found {
let entry = segment.list.remove(i);
// Update first/last keys
if segment.list.is_empty() {
segment.first_key = u16::MAX;
segment.last_key = 0;
} else {
segment.first_key = segment.list[0].key;
segment.last_key = segment.list[segment.list.len() - 1].key;
}
Some(entry.to_needle_value())
} else {
None
}
}
/// Iterate all entries in ascending NeedleId order.
pub fn ascending_visit<F, E>(&self, mut f: F) -> Result<(), E>
where
F: FnMut(NeedleId, &NeedleValue) -> Result<(), E>,
{
let mut chunks: Vec<Chunk> = self.segments.keys().copied().collect();
chunks.sort_unstable();
for chunk in chunks {
let segment = &self.segments[&chunk];
for entry in &segment.list {
let id = NeedleId(SEGMENT_CHUNK_SIZE * segment.chunk + entry.key as u64);
let nv = entry.to_needle_value();
f(id, &nv)?;
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn offset(v: u32) -> Offset {
Offset::from_actual_offset(v as i64 * NEEDLE_PADDING_SIZE as i64)
}
#[test]
fn test_compact_map_basic() {
let mut m = CompactMap::new();
// Insert
assert!(m.set(NeedleId(1), offset(100), Size(50)).is_none());
assert!(m.set(NeedleId(2), offset(200), Size(60)).is_none());
// Get
let nv = m.get(NeedleId(1)).unwrap();
assert_eq!(nv.size, Size(50));
// Update returns old value
let old = m.set(NeedleId(1), offset(300), Size(70)).unwrap();
assert_eq!(old.size, Size(50));
// Get updated value
let nv = m.get(NeedleId(1)).unwrap();
assert_eq!(nv.size, Size(70));
// Miss
assert!(m.get(NeedleId(999)).is_none());
}
#[test]
fn test_compact_map_delete() {
let mut m = CompactMap::new();
m.set(NeedleId(1), offset(100), Size(50));
// Delete returns old size
let old = m.delete(NeedleId(1)).unwrap();
assert_eq!(old, Size(50));
// Get returns deleted (negative size)
let nv = m.get(NeedleId(1)).unwrap();
assert!(nv.size.is_deleted());
// Delete again returns None (already deleted)
assert!(m.delete(NeedleId(1)).is_none());
}
#[test]
fn test_compact_map_zero_size_delete() {
let mut m = CompactMap::new();
m.set(NeedleId(1), offset(100), Size(0));
let old = m.delete(NeedleId(1)).unwrap();
assert_eq!(old, Size(0));
let nv = m.get(NeedleId(1)).unwrap();
assert_eq!(nv.size, TOMBSTONE_FILE_SIZE);
}
#[test]
fn test_compact_map_cross_segment() {
let mut m = CompactMap::new();
// Insert across multiple segments
m.set(NeedleId(1), offset(1), Size(1));
m.set(NeedleId(50_000), offset(2), Size(2));
m.set(NeedleId(100_000), offset(3), Size(3));
assert_eq!(m.get(NeedleId(1)).unwrap().size, Size(1));
assert_eq!(m.get(NeedleId(50_000)).unwrap().size, Size(2));
assert_eq!(m.get(NeedleId(100_000)).unwrap().size, Size(3));
}
#[test]
fn test_compact_map_ascending_visit() {
let mut m = CompactMap::new();
m.set(NeedleId(100_005), offset(3), Size(3));
m.set(NeedleId(5), offset(1), Size(1));
m.set(NeedleId(50_005), offset(2), Size(2));
let mut visited = Vec::new();
m.ascending_visit(|id, nv| {
visited.push((id, nv.size));
Ok::<_, String>(())
})
.unwrap();
assert_eq!(visited.len(), 3);
assert_eq!(visited[0].0, NeedleId(5));
assert_eq!(visited[1].0, NeedleId(50_005));
assert_eq!(visited[2].0, NeedleId(100_005));
}
#[test]
fn test_compact_map_remove() {
let mut m = CompactMap::new();
m.set(NeedleId(1), offset(100), Size(50));
m.set(NeedleId(2), offset(200), Size(60));
let removed = m.remove(NeedleId(1)).unwrap();
assert_eq!(removed.size, Size(50));
assert!(m.get(NeedleId(1)).is_none());
assert_eq!(m.get(NeedleId(2)).unwrap().size, Size(60));
}
#[test]
fn test_compact_map_reverse_insert_order() {
let mut m = CompactMap::new();
// Insert in reverse order to test sorted insert
for i in (0..100).rev() {
m.set(NeedleId(i), offset(i as u32), Size(i as i32));
}
for i in 0..100 {
assert_eq!(m.get(NeedleId(i)).unwrap().size, Size(i as i32));
}
}
}

View File

@@ -0,0 +1,252 @@
//! Bounded pool of open index-file descriptors.
//!
//! Read-only volumes — cloud-tiered ones above all — outnumber writable ones by
//! orders of magnitude on a large server, and a volume that pins its `.idx` and
//! `.sdx` for the life of the process costs two descriptors whether or not
//! anybody reads it. At ~600K volumes per server that alone exhausts any fd
//! limit. Neither file is needed except while a lookup is in flight, so
//! [`SortedFileNeedleMap`](super::sorted_file::SortedFileNeedleMap) borrows them
//! from this pool: an idle volume holds nothing, a busy one keeps its handles
//! hot rather than paying an `open()` per needle.
//!
//! Mirrors Go's `weed/storage/needle_map_file_pool.go`. Handles are handed out
//! as `Arc<File>`, so an eviction cannot close a descriptor a reader still
//! holds — the file closes when the last borrower drops its `Arc`.
use std::collections::{BTreeMap, HashMap};
use std::fs::{File, OpenOptions};
use std::io;
use std::sync::{Arc, Mutex, OnceLock};
use crate::storage::volume_open::open_volume_file;
/// Descriptors the pool keeps open. Matches Go's `maxPooledIndexFiles`.
pub const MAX_POOLED_INDEX_FILES: usize = 1024;
struct Entry {
file: Arc<File>,
tick: u64,
}
#[derive(Default)]
struct Inner {
entries: HashMap<String, Entry>,
/// Recency order, oldest tick first, so eviction is a `pop_first`.
order: BTreeMap<u64, String>,
next_tick: u64,
}
pub struct IndexFilePool {
capacity: usize,
inner: Mutex<Inner>,
}
/// Writable and read-only handles for the same path are pooled separately so a
/// read never depends on the file being openable for write — a volume served
/// off a read-only mount still answers lookups.
fn pool_key(path: &str, writable: bool) -> String {
if writable {
format!("{path}\0rw")
} else {
path.to_string()
}
}
impl IndexFilePool {
pub fn new(capacity: usize) -> Self {
IndexFilePool {
capacity: capacity.max(1),
inner: Mutex::new(Inner::default()),
}
}
/// Hand out an open handle for `path`, reusing the pooled one when there is
/// one. The descriptor lives as long as the returned `Arc`.
pub fn borrow(&self, path: &str, writable: bool) -> io::Result<Arc<File>> {
let key = pool_key(path, writable);
if let Some(file) = self.touch(&key) {
return Ok(file);
}
// Opened outside the lock: a cold open blocks on disk, and holding a
// process-wide mutex across it would serialize every volume's lookups.
let file = Arc::new(open_volume_file(
OpenOptions::new().read(true).write(writable),
path,
)?);
Ok(self.insert(key, file))
}
/// Forget the pooled handles for `path`, so a later rename or delete of that
/// path cannot be served from a descriptor on the old inode.
pub fn discard(&self, path: &str) {
let mut inner = self.inner.lock().unwrap();
for key in [pool_key(path, false), pool_key(path, true)] {
if let Some(entry) = inner.entries.remove(&key) {
inner.order.remove(&entry.tick);
}
}
}
/// Descriptors currently pooled. Test-only visibility into the bound.
#[cfg(test)]
pub fn pooled_count(&self) -> usize {
self.inner.lock().unwrap().entries.len()
}
fn touch(&self, key: &str) -> Option<Arc<File>> {
let mut inner = self.inner.lock().unwrap();
let tick = inner.next_tick;
let entry = inner.entries.get_mut(key)?;
let file = entry.file.clone();
let old_tick = std::mem::replace(&mut entry.tick, tick);
inner.order.remove(&old_tick);
inner.order.insert(tick, key.to_string());
inner.next_tick += 1;
Some(file)
}
fn insert(&self, key: String, file: Arc<File>) -> Arc<File> {
let mut inner = self.inner.lock().unwrap();
if let Some(entry) = inner.entries.get(&key) {
// Another borrower opened the same path first; keep one descriptor.
return entry.file.clone();
}
let tick = inner.next_tick;
inner.next_tick += 1;
inner.order.insert(tick, key.clone());
inner.entries.insert(
key,
Entry {
file: file.clone(),
tick,
},
);
while inner.entries.len() > self.capacity {
let Some((_, oldest)) = inner.order.pop_first() else {
break;
};
inner.entries.remove(&oldest);
}
file
}
}
/// Process-wide pool shared by every read-only volume on this server.
pub fn pooled_index_files() -> &'static IndexFilePool {
static POOL: OnceLock<IndexFilePool> = OnceLock::new();
POOL.get_or_init(|| IndexFilePool::new(MAX_POOLED_INDEX_FILES))
}
/// Descriptors this process holds on `.idx`/`.sdx` files under `dir`, read from
/// `/proc/self/fd` where it exists and from `lsof` otherwise. `None` when
/// neither is available, so a caller can skip rather than assert vacuously.
#[cfg(test)]
pub(crate) fn open_index_fds(dir: &std::path::Path) -> Option<usize> {
let prefix = std::fs::canonicalize(dir).unwrap_or_else(|_| dir.to_path_buf());
let is_index = |target: &std::path::Path| {
target.starts_with(&prefix)
&& matches!(
target.extension().and_then(|e| e.to_str()),
Some("idx") | Some("sdx")
)
};
if let Ok(entries) = std::fs::read_dir("/proc/self/fd") {
return Some(
entries
.filter_map(|e| std::fs::read_link(e.ok()?.path()).ok())
.filter(|target| is_index(target))
.count(),
);
}
let out = std::process::Command::new("lsof")
.args(["-p", &std::process::id().to_string(), "-F", "n"])
.output()
.ok()?;
Some(
String::from_utf8_lossy(&out.stdout)
.lines()
.filter_map(|line| line.strip_prefix('n'))
.filter(|line| is_index(std::path::Path::new(line)))
.count(),
)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
fn write_file(dir: &std::path::Path, name: &str, contents: &[u8]) -> String {
let path = dir.join(name);
let mut f = File::create(&path).unwrap();
f.write_all(contents).unwrap();
path.to_str().unwrap().to_string()
}
#[test]
fn evicted_handle_stays_usable_for_its_borrower() {
let dir = tempfile::tempdir().unwrap();
let first = write_file(dir.path(), "first", b"first");
let second = write_file(dir.path(), "second", b"second");
let pool = IndexFilePool::new(1);
let borrowed = pool.borrow(&first, false).unwrap();
// Pushes the single slot over, evicting the entry still in use.
let _other = pool.borrow(&second, false).unwrap();
assert_eq!(pool.pooled_count(), 1);
let mut buf = [0u8; 5];
#[cfg(unix)]
{
use std::os::unix::fs::FileExt;
borrowed.read_exact_at(&mut buf, 0).unwrap();
}
assert_eq!(&buf, b"first");
}
#[test]
fn borrow_reuses_the_pooled_handle() {
let dir = tempfile::tempdir().unwrap();
let path = write_file(dir.path(), "idx", b"x");
let pool = IndexFilePool::new(4);
let a = pool.borrow(&path, false).unwrap();
let b = pool.borrow(&path, false).unwrap();
assert!(Arc::ptr_eq(&a, &b));
assert_eq!(pool.pooled_count(), 1);
// A writable handle is pooled separately from the read-only one.
let w = pool.borrow(&path, true).unwrap();
assert!(!Arc::ptr_eq(&a, &w));
assert_eq!(pool.pooled_count(), 2);
}
#[test]
fn discard_drops_both_handles() {
let dir = tempfile::tempdir().unwrap();
let path = write_file(dir.path(), "idx", b"x");
let pool = IndexFilePool::new(4);
let _r = pool.borrow(&path, false).unwrap();
let _w = pool.borrow(&path, true).unwrap();
assert_eq!(pool.pooled_count(), 2);
pool.discard(&path);
assert_eq!(pool.pooled_count(), 0);
}
#[test]
fn pool_stays_within_capacity() {
let dir = tempfile::tempdir().unwrap();
let pool = IndexFilePool::new(3);
for i in 0..10 {
let path = write_file(dir.path(), &format!("f{i}"), b"x");
let _ = pool.borrow(&path, false).unwrap();
}
assert_eq!(pool.pooled_count(), 3);
}
}

View File

@@ -0,0 +1,159 @@
//! Rebuild a needle map's counters from the whole `.idx` history.
//!
//! The volume's garbage ratio is `deleted_size / content_size`, and both are
//! additive over the life of the volume: an overwritten needle keeps its
//! bytes in `content_size` and adds them to `deleted_size`. A backend whose
//! table only holds the final value per key (redb) cannot recover that from
//! the table, so on load the counters come from the `.idx` file instead.
//!
//! This mirrors Go's `needleMapMetricFromIndexFile`: walk the index newest
//! entry first with a bloom filter of the keys already seen, so the memory
//! cost is a few bits per entry instead of a map of every key. The counting
//! rule reproduces what the live `on_put`/`on_delete` path accumulates:
//!
//! - every live entry is one put: `file_count`, `file_byte_count`;
//! - a live entry with a newer entry for the same key was overwritten or
//! deleted later, so it is also one deletion: `deletion_count`,
//! `deletion_byte_count`;
//! - a tombstone only marks its key as seen.
//!
//! A bloom false positive (0.1%) can only add a spurious deletion, which
//! over-reports garbage slightly; it never hides any.
use std::io::{self, Read, Seek, SeekFrom};
use std::sync::atomic::Ordering;
use xxhash_rust::xxh64::xxh64;
use super::NeedleMapMetric;
use crate::storage::types::*;
/// Entries read per batch while walking backwards (64 KiB of index).
const BATCH_ENTRIES: usize = 4096;
/// Same target false-positive rate as the Go server's filter.
const FALSE_POSITIVE_RATE: f64 = 0.001;
/// Minimal bloom filter over needle ids, double hashing with xxh64.
struct SeenKeys {
bits: Vec<u64>,
bit_count: u64,
hashes: u64,
}
impl SeenKeys {
fn new(expected: u64, false_positive_rate: f64) -> Self {
let n = expected.max(1) as f64;
let ln2 = std::f64::consts::LN_2;
let bit_count = (-(n * false_positive_rate.ln()) / (ln2 * ln2))
.ceil()
.max(64.0) as u64;
let hashes = ((bit_count as f64 / n) * ln2).round().clamp(1.0, 16.0) as u64;
SeenKeys {
bits: vec![0u64; bit_count.div_ceil(64) as usize],
bit_count,
hashes,
}
}
/// Whether `key` was (probably) seen before; marks it seen either way.
fn test_and_add(&mut self, key: u64) -> bool {
let bytes = key.to_le_bytes();
let h1 = xxh64(&bytes, 0);
let h2 = xxh64(&bytes, 0x9E37_79B9_7F4A_7C15) | 1;
let mut seen = true;
for i in 0..self.hashes {
let bit = h1.wrapping_add(i.wrapping_mul(h2)) % self.bit_count;
let word = (bit / 64) as usize;
let mask = 1u64 << (bit % 64);
if self.bits[word] & mask == 0 {
seen = false;
self.bits[word] |= mask;
}
}
seen
}
}
/// Walk `reader` (an `.idx` file) newest entry first and return the counters
/// a live volume would hold after applying the same history. A torn partial
/// entry at the tail is ignored, as `walk_index_file` does. The reader is
/// left positioned at the start of the file.
pub(super) fn metrics_from_idx<R: Read + Seek>(
reader: &mut R,
version: Version,
) -> io::Result<NeedleMapMetric> {
let metric = NeedleMapMetric::default();
let file_size = reader.seek(SeekFrom::End(0))?;
let entry_count = file_size / NEEDLE_MAP_ENTRY_SIZE as u64;
let mut seen = SeenKeys::new(entry_count, FALSE_POSITIVE_RATE);
let mut buf = vec![0u8; NEEDLE_MAP_ENTRY_SIZE * BATCH_ENTRIES];
let mut remaining = entry_count;
while remaining > 0 {
let batch = remaining.min(BATCH_ENTRIES as u64) as usize;
let first_entry = remaining - batch as u64;
let len = batch * NEEDLE_MAP_ENTRY_SIZE;
reader.seek(SeekFrom::Start(first_entry * NEEDLE_MAP_ENTRY_SIZE as u64))?;
reader.read_exact(&mut buf[..len])?;
for i in (0..batch).rev() {
let entry = &buf[i * NEEDLE_MAP_ENTRY_SIZE..(i + 1) * NEEDLE_MAP_ENTRY_SIZE];
let (key, offset, size) = idx_entry_from_bytes(entry);
metric.maybe_set_max_needle_end(offset, size, version);
let superseded = seen.test_and_add(key.into());
if offset.is_zero() || size.is_deleted() {
// Tombstone: reserves no bytes, only marks the key as seen.
continue;
}
metric.maybe_set_max_file_key(key);
metric.file_count.fetch_add(1, Ordering::Relaxed);
metric
.file_byte_count
.fetch_add(size.0 as u64, Ordering::Relaxed);
if superseded && size.0 > 0 {
metric.deletion_count.fetch_add(1, Ordering::Relaxed);
metric
.deletion_byte_count
.fetch_add(size.0 as u64, Ordering::Relaxed);
}
}
remaining = first_entry;
}
reader.seek(SeekFrom::Start(0))?;
Ok(metric)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_seen_keys_reports_repeats_and_not_fresh_keys() {
let mut seen = SeenKeys::new(10_000, FALSE_POSITIVE_RATE);
// Fresh keys may occasionally collide (that is the false-positive
// rate), but only rarely.
let fresh_reported_seen = (0..10_000u64).filter(|&key| seen.test_and_add(key)).count();
assert!(
fresh_reported_seen < 50,
"fresh keys reported seen: {fresh_reported_seen}"
);
// A repeated key is never reported fresh: no false negatives.
for key in 0..10_000u64 {
assert!(seen.test_and_add(key), "repeated key {key} reported fresh");
}
// Over 100k never-inserted keys the false-positive rate stays near
// the 0.1% target; allow a generous margin.
let false_positives = (1_000_000..1_100_000u64)
.filter(|k| {
let bytes = k.to_le_bytes();
let h1 = xxh64(&bytes, 0);
let h2 = xxh64(&bytes, 0x9E37_79B9_7F4A_7C15) | 1;
(0..seen.hashes).all(|i| {
let bit = h1.wrapping_add(i.wrapping_mul(h2)) % seen.bit_count;
seen.bits[(bit / 64) as usize] & (1u64 << (bit % 64)) != 0
})
})
.count();
assert!(false_positives < 500, "false positives: {false_positives}");
}
}

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,420 @@
//! Physical EC sidecar mirroring across disks of the same volume
//! server. Mirrors `weed/storage/store_ec_mirror.go`.
use std::collections::HashMap;
use std::fs;
use std::io::{self, Read, Write};
use std::path::Path;
use tracing::{info, warn};
use crate::storage::disk_location::{DiskLocation, parse_collection_volume_id_pub};
use crate::storage::store::Store;
use crate::storage::types::VolumeId;
// Listed in `EcVolume::new`'s open order.
const EC_MIRRORED_SIDECARS: &[&str] = &[".ecx", ".ecj", ".vif"];
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct EcKey {
collection: String,
vid: VolumeId,
}
#[derive(Clone, Debug)]
struct EcxOwner {
location: usize,
idx_dir: String,
data_dir: String,
}
impl Store {
/// Mirror EC sidecars onto every shard-bearing disk that lacks
/// them, so each disk mounts self-contained. Runs before
/// `reconcile_ec_shards_across_disks` so the orphan pass can
/// prefer the local idx_directory.
pub fn mirror_ec_metadata_to_shard_disks(&mut self) {
if self.locations.len() < 2 {
return;
}
let owners = self.index_ecx_owners_for_mirror();
if owners.is_empty() {
return;
}
// Two-pass: gather work under an immutable borrow, then apply
// copies under independent mutable borrows.
struct Mirror<'a> {
target_idx: usize,
owner: &'a EcxOwner,
collection: String,
vid: VolumeId,
}
let mut mirrors: Vec<Mirror> = Vec::new();
for (loc_idx, loc) in self.locations.iter().enumerate() {
let orphans = collect_shard_disk_volumes(loc);
for (key, _shards) in orphans {
let Some(owner) = owners.get(&key) else {
continue;
};
if owner.location == loc_idx {
continue;
}
if disk_has_all_sidecars(loc, &key.collection, key.vid) {
continue;
}
mirrors.push(Mirror {
target_idx: loc_idx,
owner,
collection: key.collection,
vid: key.vid,
});
}
}
for m in mirrors {
let loc_dir = self.locations[m.target_idx].directory.clone();
let loc_idx_dir = self.locations[m.target_idx].idx_directory.clone();
match mirror_sidecars_for_volume(
&m.owner.idx_dir,
&m.owner.data_dir,
&loc_dir,
&loc_idx_dir,
&m.collection,
m.vid,
) {
Ok(0) => {}
Ok(copied) => {
info!(
volume_id = m.vid.0,
collection = %m.collection,
from = %m.owner.data_dir,
to = %loc_dir,
copied,
"mirrored EC sidecar(s) for same-disk invariant",
);
}
Err(e) => {
warn!(
volume_id = m.vid.0,
collection = %m.collection,
from = %m.owner.data_dir,
to = %loc_dir,
error = %e,
"mirror EC sidecars failed; cross-disk fallback will handle this volume",
);
}
}
}
}
// Records both idx_dir (where .ecx was found) and data_dir, so
// the mirror can resolve .vif from data_dir even when .ecx lives
// in idx_directory.
fn index_ecx_owners_for_mirror(&self) -> HashMap<EcKey, EcxOwner> {
let mut owners: HashMap<EcKey, EcxOwner> = HashMap::new();
for (loc_idx, loc) in self.locations.iter().enumerate() {
let mut seen: Vec<&str> = Vec::with_capacity(2);
for scan in [loc.idx_directory.as_str(), loc.directory.as_str()] {
if scan.is_empty() || seen.contains(&scan) {
continue;
}
seen.push(scan);
let Ok(read) = fs::read_dir(scan) else {
continue;
};
for ent in read.flatten() {
if ent.file_type().map(|ft| ft.is_dir()).unwrap_or(false) {
continue;
}
let name = ent.file_name().to_string_lossy().into_owned();
let Some(base) = name.strip_suffix(".ecx") else {
continue;
};
let Some((collection, vid)) = parse_collection_volume_id_pub(base) else {
continue;
};
owners
.entry(EcKey { collection, vid })
.or_insert_with(|| EcxOwner {
location: loc_idx,
idx_dir: scan.to_string(),
data_dir: loc.directory.clone(),
});
}
}
}
owners
}
}
// Checks the modern routing and the opposite directory — without
// that fallback, a destination with a legacy pre-`-dir.idx` .ecx in
// its data dir would be re-mirrored into idx_directory.
fn disk_has_all_sidecars(loc: &DiskLocation, collection: &str, vid: VolumeId) -> bool {
for ext in EC_MIRRORED_SIDECARS {
let primary = sidecar_dest_path(&loc.directory, &loc.idx_directory, collection, vid, ext);
if path_is_regular_file(&primary) {
continue;
}
if loc.idx_directory != loc.directory {
let (fallback_data, fallback_idx) = if *ext == ".vif" {
(loc.idx_directory.as_str(), loc.directory.as_str())
} else {
(loc.directory.as_str(), loc.directory.as_str())
};
let fallback = sidecar_dest_path(fallback_data, fallback_idx, collection, vid, ext);
if path_is_regular_file(&fallback) {
continue;
}
}
return false;
}
true
}
fn path_is_regular_file(path: &str) -> bool {
fs::metadata(path).map(|m| !m.is_dir()).unwrap_or(false)
}
// `.ecx`/`.ecj` route to idx_directory, `.vif` to directory.
fn sidecar_dest_path(
data_dir: &str,
idx_dir: &str,
collection: &str,
vid: VolumeId,
ext: &str,
) -> String {
let dir = if ext == ".vif" { data_dir } else { idx_dir };
if collection.is_empty() {
format!("{}/{}{}", dir, vid.0, ext)
} else {
format!("{}/{}_{}{}", dir, collection, vid.0, ext)
}
}
fn mirror_sidecars_for_volume(
src_idx_dir: &str,
src_data_dir: &str,
dst_data_dir: &str,
dst_idx_dir: &str,
collection: &str,
vid: VolumeId,
) -> io::Result<usize> {
let mut copied = 0usize;
for ext in EC_MIRRORED_SIDECARS {
let dst = sidecar_dest_path(dst_data_dir, dst_idx_dir, collection, vid, ext);
// An existing local copy is authoritative — it may be newer
// than the owner's after a delete journal append.
if fs::metadata(&dst).is_ok() {
continue;
}
let candidates = [
sidecar_dest_path(src_data_dir, src_idx_dir, collection, vid, ext),
sidecar_dest_path(src_idx_dir, src_data_dir, collection, vid, ext),
];
let mut src_path: Option<String> = None;
for c in candidates.iter() {
if fs::metadata(c).map(|m| !m.is_dir()).unwrap_or(false) {
src_path = Some(c.clone());
break;
}
}
let Some(src) = src_path else {
continue;
};
copy_sidecar_atomic(Path::new(&src), Path::new(&dst))?;
copied += 1;
}
Ok(copied)
}
fn copy_sidecar_atomic(src: &Path, dst: &Path) -> io::Result<()> {
if let Some(parent) = dst.parent() {
fs::create_dir_all(parent)?;
}
let mut src_file = fs::File::open(src)?;
let tmp = {
let mut s = dst.as_os_str().to_owned();
s.push(".mirror.tmp");
std::path::PathBuf::from(s)
};
let _ = fs::remove_file(&tmp);
let mut dst_file = fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&tmp)?;
let mut buf = [0u8; 64 * 1024];
loop {
let n = src_file.read(&mut buf)?;
if n == 0 {
break;
}
dst_file.write_all(&buf[..n])?;
}
dst_file.sync_all()?;
drop(dst_file);
if let Err(e) = fs::rename(&tmp, dst) {
let _ = fs::remove_file(&tmp);
return Err(e);
}
Ok(())
}
fn collect_shard_disk_volumes(loc: &DiskLocation) -> HashMap<EcKey, Vec<String>> {
let mut out: HashMap<EcKey, Vec<String>> = HashMap::new();
let Ok(read) = fs::read_dir(&loc.directory) else {
return out;
};
for ent in read.flatten() {
if ent.file_type().map(|ft| ft.is_dir()).unwrap_or(false) {
continue;
}
let name = ent.file_name().to_string_lossy().into_owned();
let Some(dot) = name.rfind('.') else {
continue;
};
let (base, ext) = name.split_at(dot);
if crate::storage::disk_location::is_ec_shard_extension(ext).is_none() {
continue;
}
match ent.metadata() {
Ok(meta) if meta.len() > 0 => {}
_ => continue,
}
let Some((collection, vid)) = parse_collection_volume_id_pub(base) else {
continue;
};
out.entry(EcKey { collection, vid }).or_default().push(name);
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::MinFreeSpace;
use crate::storage::needle_map::NeedleMapKind;
use crate::storage::types::DiskType;
use crate::storage::volume::{VifEcShardConfig, VifVolumeInfo};
use tempfile::TempDir;
fn plant_shard(dir: &Path, collection: &str, vid: u32, shard_id: u8) {
let path = if collection.is_empty() {
dir.join(format!("{}.ec{:02}", vid, shard_id))
} else {
dir.join(format!("{}_{}.ec{:02}", collection, vid, shard_id))
};
fs::write(&path, b"shard data nonempty").unwrap();
}
fn plant_ecx(dir: &Path, collection: &str, vid: u32, bytes: &[u8]) {
let path = dir.join(format!("{}_{}.ecx", collection, vid));
fs::write(&path, bytes).unwrap();
}
fn plant_ecj(dir: &Path, collection: &str, vid: u32, bytes: &[u8]) {
let path = dir.join(format!("{}_{}.ecj", collection, vid));
fs::write(&path, bytes).unwrap();
}
fn plant_vif(dir: &Path, collection: &str, vid: u32, data_shards: u32, parity_shards: u32) {
let vif = VifVolumeInfo {
version: 3,
ec_shard_config: Some(VifEcShardConfig {
data_shards,
parity_shards,
..Default::default()
}),
..Default::default()
};
let path = dir.join(format!("{}_{}.vif", collection, vid));
fs::write(&path, serde_json::to_string(&vif).unwrap()).unwrap();
}
fn add_loc(store: &mut Store, dir: &Path) {
store
.add_location(
dir.to_str().unwrap(),
dir.to_str().unwrap(),
100,
DiskType::HardDrive,
MinFreeSpace::Percent(0.0),
Vec::new(),
)
.unwrap();
}
#[test]
fn mirror_copies_sidecars_to_shard_only_disk() {
let tmp = TempDir::new().unwrap();
let dir0 = tmp.path().join("data0");
let dir1 = tmp.path().join("data1");
fs::create_dir_all(&dir0).unwrap();
fs::create_dir_all(&dir1).unwrap();
let collection = "video-recordings";
let vid = 4121u32;
plant_shard(&dir0, collection, vid, 0);
plant_shard(&dir0, collection, vid, 12);
plant_shard(&dir1, collection, vid, 1);
let ecx = vec![0xA1u8; 20];
let ecj = vec![0xB2u8; 16];
plant_ecx(&dir1, collection, vid, &ecx);
plant_ecj(&dir1, collection, vid, &ecj);
plant_vif(&dir1, collection, vid, 10, 4);
let mut store = Store::new(NeedleMapKind::InMemory);
add_loc(&mut store, &dir0);
add_loc(&mut store, &dir1);
for ext in [".ecx", ".ecj", ".vif"] {
let dst = dir0.join(format!("{}_{}{}", collection, vid, ext));
assert!(
dst.exists(),
"mirror did not install sidecar {} on dir0",
ext
);
}
let ecx_dst = fs::read(dir0.join(format!("{}_{}.ecx", collection, vid))).unwrap();
assert_eq!(ecx_dst, ecx, ".ecx mirrored bytes differ from source");
let ecj_dst = fs::read(dir0.join(format!("{}_{}.ecj", collection, vid))).unwrap();
assert_eq!(ecj_dst, ecj, ".ecj mirrored bytes differ from source");
}
#[test]
fn mirror_preserves_existing_destination_sidecars() {
let tmp = TempDir::new().unwrap();
let dir0 = tmp.path().join("data0");
let dir1 = tmp.path().join("data1");
fs::create_dir_all(&dir0).unwrap();
fs::create_dir_all(&dir1).unwrap();
let collection = "video-recordings";
let vid = 7777u32;
plant_shard(&dir0, collection, vid, 0);
plant_shard(&dir0, collection, vid, 12);
plant_shard(&dir1, collection, vid, 1);
let ecx_owner = vec![0xC3u8; 20];
let ecx_local = vec![0x5Au8; 20];
let ecj_bytes = vec![0xD4u8; 16];
plant_ecx(&dir1, collection, vid, &ecx_owner);
plant_ecj(&dir1, collection, vid, &ecj_bytes);
plant_vif(&dir1, collection, vid, 10, 4);
plant_ecx(&dir0, collection, vid, &ecx_local);
plant_ecj(&dir0, collection, vid, &ecj_bytes);
plant_vif(&dir0, collection, vid, 10, 4);
let mut store = Store::new(NeedleMapKind::InMemory);
add_loc(&mut store, &dir0);
add_loc(&mut store, &dir1);
let post = fs::read(dir0.join(format!("{}_{}.ecx", collection, vid))).unwrap();
assert_eq!(post, ecx_local, "mirror overwrote dir0's existing .ecx");
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,289 @@
//! SuperBlock: the 8-byte (+ optional extra) header at the start of every .dat file.
//!
//! Byte layout:
//! [0] Version
//! [1] ReplicaPlacement byte
//! [2..4] TTL (2 bytes)
//! [4..6] CompactionRevision (u16 big-endian)
//! [6..8] ExtraSize (u16 big-endian)
//! [8..] Extra data (protobuf, ExtraSize bytes) — only for Version 2/3
use crate::storage::needle::ttl::TTL;
use crate::storage::types::Version;
pub const SUPER_BLOCK_SIZE: usize = 8;
/// SuperBlock metadata at the start of a volume .dat file.
#[derive(Debug, Clone)]
pub struct SuperBlock {
pub version: Version,
pub replica_placement: ReplicaPlacement,
pub ttl: TTL,
pub compaction_revision: u16,
pub extra_size: u16,
pub extra_data: Vec<u8>, // raw protobuf bytes (SuperBlockExtra)
}
impl SuperBlock {
/// Total block size on disk (base 8 + extra).
pub fn block_size(&self) -> usize {
match self.version.0 {
2 | 3 => SUPER_BLOCK_SIZE + self.extra_size as usize,
_ => SUPER_BLOCK_SIZE,
}
}
/// Serialize to bytes.
pub fn to_bytes(&self) -> Vec<u8> {
let mut header = vec![0u8; SUPER_BLOCK_SIZE];
header[0] = self.version.0;
header[1] = self.replica_placement.to_byte();
self.ttl.to_bytes(&mut header[2..4]);
header[4..6].copy_from_slice(&self.compaction_revision.to_be_bytes());
if !self.extra_data.is_empty() {
// Go checks extraSize > 256*256-2 and calls glog.Fatalf; guard against u16 overflow.
assert!(
self.extra_data.len() <= 65534,
"super block extra data too large: {} > 65534",
self.extra_data.len()
);
let extra_size = self.extra_data.len() as u16;
header[6..8].copy_from_slice(&extra_size.to_be_bytes());
header.extend_from_slice(&self.extra_data);
}
header
}
/// Parse from bytes (must be at least SUPER_BLOCK_SIZE bytes).
pub fn from_bytes(bytes: &[u8]) -> Result<Self, SuperBlockError> {
if bytes.len() < SUPER_BLOCK_SIZE {
return Err(SuperBlockError::TooShort(bytes.len()));
}
let version = Version(bytes[0]);
let replica_placement = ReplicaPlacement::from_byte(bytes[1])?;
let ttl = TTL::from_bytes(&bytes[2..4]);
let compaction_revision = u16::from_be_bytes([bytes[4], bytes[5]]);
let extra_size = u16::from_be_bytes([bytes[6], bytes[7]]);
let extra_data = if extra_size > 0 && bytes.len() >= SUPER_BLOCK_SIZE + extra_size as usize
{
bytes[SUPER_BLOCK_SIZE..SUPER_BLOCK_SIZE + extra_size as usize].to_vec()
} else {
vec![]
};
Ok(SuperBlock {
version,
replica_placement,
ttl,
compaction_revision,
extra_size,
extra_data,
})
}
pub fn initialized(&self) -> bool {
true // ReplicaPlacement and TTL are always valid after construction
}
}
impl Default for SuperBlock {
fn default() -> Self {
SuperBlock {
version: Version::current(),
replica_placement: ReplicaPlacement::default(),
ttl: TTL::EMPTY,
compaction_revision: 0,
extra_size: 0,
extra_data: vec![],
}
}
}
// ============================================================================
// ReplicaPlacement
// ============================================================================
/// Replication strategy encoded as a single byte.
///
/// Byte value = DiffDataCenterCount * 100 + DiffRackCount * 10 + SameRackCount
///
/// Examples:
/// "000" → no replication (1 copy total)
/// "010" → 1 copy in different rack (2 copies total)
/// "100" → 1 copy in different datacenter
/// "200" → 2 copies in different datacenters
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct ReplicaPlacement {
pub same_rack_count: u8,
pub diff_rack_count: u8,
pub diff_data_center_count: u8,
}
impl ReplicaPlacement {
/// Parse from a string like "000", "010", "100".
/// Accepts 0-3 character strings, padding with leading zeros to match Go behavior.
/// E.g. "" -> "000", "1" -> "001", "01" -> "001", "010" -> "010"
pub fn from_string(s: &str) -> Result<Self, SuperBlockError> {
let s = s.trim();
if s.is_empty() {
return Ok(ReplicaPlacement::default());
}
// Pad with leading zeros to 3 chars, matching Go's NewReplicaPlacementFromString
let padded = match s.len() {
1 => format!("00{}", s),
2 => format!("0{}", s),
3 => s.to_string(),
_ => return Err(SuperBlockError::InvalidReplicaPlacement(s.to_string())),
};
let chars: Vec<char> = padded.chars().collect();
let dc = chars[0]
.to_digit(10)
.ok_or_else(|| SuperBlockError::InvalidReplicaPlacement(s.to_string()))?
as u8;
let rack = chars[1]
.to_digit(10)
.ok_or_else(|| SuperBlockError::InvalidReplicaPlacement(s.to_string()))?
as u8;
let same = chars[2]
.to_digit(10)
.ok_or_else(|| SuperBlockError::InvalidReplicaPlacement(s.to_string()))?
as u8;
// Go validates: value = dc*100 + rack*10 + same must fit in a byte
let value = dc as u16 * 100 + rack as u16 * 10 + same as u16;
if value > 255 {
return Err(SuperBlockError::InvalidReplicaPlacement(s.to_string()));
}
Ok(ReplicaPlacement {
diff_data_center_count: dc,
diff_rack_count: rack,
same_rack_count: same,
})
}
/// Parse from a single byte.
pub fn from_byte(b: u8) -> Result<Self, SuperBlockError> {
Ok(ReplicaPlacement {
diff_data_center_count: b / 100,
diff_rack_count: (b % 100) / 10,
same_rack_count: b % 10,
})
}
/// Encode as a single byte.
pub fn to_byte(&self) -> u8 {
self.diff_data_center_count * 100 + self.diff_rack_count * 10 + self.same_rack_count
}
/// Total number of copies (including the original).
pub fn get_copy_count(&self) -> u8 {
self.diff_data_center_count + self.diff_rack_count + self.same_rack_count + 1
}
/// Whether this placement requires replication (more than 1 copy).
pub fn has_replication(&self) -> bool {
self.get_copy_count() > 1
}
}
impl std::fmt::Display for ReplicaPlacement {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}{}{}",
self.diff_data_center_count, self.diff_rack_count, self.same_rack_count
)
}
}
// ============================================================================
// Errors
// ============================================================================
#[derive(Debug, thiserror::Error)]
pub enum SuperBlockError {
#[error("super block too short: {0} bytes")]
TooShort(usize),
#[error("invalid replica placement: {0}")]
InvalidReplicaPlacement(String),
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::types::*;
#[test]
fn test_super_block_round_trip() {
let sb = SuperBlock {
version: VERSION_3,
replica_placement: ReplicaPlacement::from_string("010").unwrap(),
ttl: TTL { count: 5, unit: 3 },
compaction_revision: 42,
extra_size: 0,
extra_data: vec![],
};
let bytes = sb.to_bytes();
assert_eq!(bytes.len(), SUPER_BLOCK_SIZE);
let sb2 = SuperBlock::from_bytes(&bytes).unwrap();
assert_eq!(sb2.version, sb.version);
assert_eq!(sb2.replica_placement, sb.replica_placement);
assert_eq!(sb2.ttl, sb.ttl);
assert_eq!(sb2.compaction_revision, sb.compaction_revision);
}
#[test]
fn test_super_block_with_extra() {
let sb = SuperBlock {
version: VERSION_3,
replica_placement: ReplicaPlacement::default(),
ttl: TTL::EMPTY,
compaction_revision: 0,
extra_size: 3,
extra_data: vec![1, 2, 3],
};
let bytes = sb.to_bytes();
assert_eq!(bytes.len(), SUPER_BLOCK_SIZE + 3);
let sb2 = SuperBlock::from_bytes(&bytes).unwrap();
assert_eq!(sb2.extra_data, vec![1, 2, 3]);
}
#[test]
fn test_replica_placement_byte_round_trip() {
let rp = ReplicaPlacement::from_string("123").unwrap();
assert_eq!(rp.diff_data_center_count, 1);
assert_eq!(rp.diff_rack_count, 2);
assert_eq!(rp.same_rack_count, 3);
assert_eq!(rp.to_byte(), 123);
assert_eq!(rp.get_copy_count(), 7); // 1+2+3+1
let rp2 = ReplicaPlacement::from_byte(123).unwrap();
assert_eq!(rp, rp2);
}
#[test]
fn test_replica_placement_no_replication() {
let rp = ReplicaPlacement::from_string("000").unwrap();
assert!(!rp.has_replication());
assert_eq!(rp.get_copy_count(), 1);
}
#[test]
fn test_replica_placement_display() {
let rp = ReplicaPlacement::from_string("010").unwrap();
assert_eq!(rp.to_string(), "010");
assert!(rp.has_replication());
}
}

View File

@@ -0,0 +1,659 @@
//! Core storage types: NeedleId, Offset, Size, Cookie, DiskType.
//!
//! These types define the binary-compatible on-disk format matching the Go implementation.
//! CRITICAL: Byte layout must match exactly for cross-compatibility.
use std::fmt;
// ============================================================================
// Constants
// ============================================================================
pub const NEEDLE_ID_SIZE: usize = 8;
pub const NEEDLE_ID_EMPTY: u64 = 0;
pub const COOKIE_SIZE: usize = 4;
pub const SIZE_SIZE: usize = 4;
pub const NEEDLE_HEADER_SIZE: usize = COOKIE_SIZE + NEEDLE_ID_SIZE + SIZE_SIZE; // 16
pub const DATA_SIZE_SIZE: usize = 4;
pub const TIMESTAMP_SIZE: usize = 8;
pub const NEEDLE_PADDING_SIZE: usize = 8;
pub const NEEDLE_CHECKSUM_SIZE: usize = 4;
/// 5-byte offset mode (matching Go production builds with `-tags 5BytesOffset`).
/// Max volume size: 8TB. Index entry: 17 bytes (8 + 5 + 4).
#[cfg(feature = "5bytes")]
pub const OFFSET_SIZE: usize = 5;
#[cfg(feature = "5bytes")]
pub const MAX_POSSIBLE_VOLUME_SIZE: u64 = 4 * 1024 * 1024 * 1024 * 8 * 256; // 8TB
/// 4-byte offset mode (matching Go default build without `5BytesOffset`).
/// Max volume size: 32GB. Index entry: 16 bytes (8 + 4 + 4).
#[cfg(not(feature = "5bytes"))]
pub const OFFSET_SIZE: usize = 4;
#[cfg(not(feature = "5bytes"))]
pub const MAX_POSSIBLE_VOLUME_SIZE: u64 = 4 * 1024 * 1024 * 1024 * 8; // 32GB
pub const NEEDLE_MAP_ENTRY_SIZE: usize = NEEDLE_ID_SIZE + OFFSET_SIZE + SIZE_SIZE;
// ============================================================================
// NeedleId
// ============================================================================
/// 64-bit unique identifier for a needle within a volume.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Default)]
pub struct NeedleId(pub u64);
impl NeedleId {
pub fn to_bytes(&self, bytes: &mut [u8]) {
assert!(bytes.len() >= NEEDLE_ID_SIZE);
bytes[0..8].copy_from_slice(&self.0.to_be_bytes());
}
pub fn from_bytes(bytes: &[u8]) -> Self {
assert!(bytes.len() >= NEEDLE_ID_SIZE);
NeedleId(u64::from_be_bytes([
bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
]))
}
pub fn is_empty(&self) -> bool {
self.0 == 0
}
/// Parse a hex string into a NeedleId.
pub fn parse(s: &str) -> Result<Self, std::num::ParseIntError> {
u64::from_str_radix(s, 16).map(NeedleId)
}
}
impl fmt::Display for NeedleId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:x}", self.0)
}
}
impl From<u64> for NeedleId {
fn from(v: u64) -> Self {
NeedleId(v)
}
}
impl From<NeedleId> for u64 {
fn from(v: NeedleId) -> Self {
v.0
}
}
// ============================================================================
// Cookie
// ============================================================================
/// Random 32-bit value to mitigate brute-force lookups.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub struct Cookie(pub u32);
impl Cookie {
pub fn to_bytes(&self, bytes: &mut [u8]) {
assert!(bytes.len() >= COOKIE_SIZE);
bytes[0..4].copy_from_slice(&self.0.to_be_bytes());
}
pub fn from_bytes(bytes: &[u8]) -> Self {
assert!(bytes.len() >= COOKIE_SIZE);
Cookie(u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
}
/// Parse a hex string into a Cookie.
pub fn parse(s: &str) -> Result<Self, std::num::ParseIntError> {
u32::from_str_radix(s, 16).map(Cookie)
}
}
impl fmt::Display for Cookie {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:x}", self.0)
}
}
impl From<u32> for Cookie {
fn from(v: u32) -> Self {
Cookie(v)
}
}
// ============================================================================
// Size
// ============================================================================
/// Needle size as stored in the index. Negative = deleted.
///
/// - Positive: valid needle with that many bytes of body content
/// - TombstoneFileSize (-1): tombstone marker
/// - Other negative: deleted, absolute value was the original size
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub struct Size(pub i32);
/// Special marker for a tombstone (deletion marker) entry.
pub const TOMBSTONE_FILE_SIZE: Size = Size(-1);
impl Size {
pub fn is_tombstone(&self) -> bool {
self.0 == TOMBSTONE_FILE_SIZE.0
}
pub fn is_deleted(&self) -> bool {
self.0 < 0 || self.0 == TOMBSTONE_FILE_SIZE.0
}
pub fn is_valid(&self) -> bool {
self.0 > 0 && !self.is_tombstone()
}
/// Raw storage size. For tombstones returns 0; for negative returns abs value.
pub fn raw(&self) -> u32 {
if self.is_tombstone() {
return 0;
}
if self.0 < 0 {
return -self.0 as u32;
}
self.0 as u32
}
pub fn to_bytes(&self, bytes: &mut [u8]) {
assert!(bytes.len() >= SIZE_SIZE);
bytes[0..4].copy_from_slice(&(self.0 as u32).to_be_bytes());
}
pub fn from_bytes(bytes: &[u8]) -> Self {
assert!(bytes.len() >= SIZE_SIZE);
let v = u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
Size(v as i32)
}
}
impl From<i32> for Size {
fn from(v: i32) -> Self {
Size(v)
}
}
impl From<Size> for i32 {
fn from(v: Size) -> Self {
v.0
}
}
// ============================================================================
// Offset
// ============================================================================
/// Offset encoding for needle positions in .dat files.
///
/// The offset is stored divided by NEEDLE_PADDING_SIZE (8).
///
/// With `5bytes` feature (default, matching Go production builds):
/// 5 bytes can address up to 8TB.
/// On-disk layout: [b3][b2][b1][b0][b4] (big-endian 4 bytes + 1 high byte)
///
/// Without `5bytes` feature (matching Go default build):
/// 4 bytes can address up to 32GB.
/// On-disk layout: [b3][b2][b1][b0] (big-endian 4 bytes)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
pub struct Offset {
pub b0: u8,
pub b1: u8,
pub b2: u8,
pub b3: u8,
#[cfg(feature = "5bytes")]
pub b4: u8,
}
impl Offset {
/// Convert to the actual byte offset in the .dat file.
pub fn to_actual_offset(&self) -> i64 {
let stored = self.b0 as i64
+ (self.b1 as i64) * 256
+ (self.b2 as i64) * 65536
+ (self.b3 as i64) * 16777216;
#[cfg(feature = "5bytes")]
let stored = stored + (self.b4 as i64) * 4294967296; // 1 << 32
stored * NEEDLE_PADDING_SIZE as i64
}
/// Create an Offset from an actual byte offset.
pub fn from_actual_offset(offset: i64) -> Self {
let smaller = offset / NEEDLE_PADDING_SIZE as i64;
Offset {
b0: smaller as u8,
b1: (smaller >> 8) as u8,
b2: (smaller >> 16) as u8,
b3: (smaller >> 24) as u8,
#[cfg(feature = "5bytes")]
b4: (smaller >> 32) as u8,
}
}
/// Serialize to bytes in the .idx file format.
/// 5-byte layout: [b3][b2][b1][b0][b4]
/// 4-byte layout: [b3][b2][b1][b0]
pub fn to_bytes(&self, bytes: &mut [u8]) {
assert!(bytes.len() >= OFFSET_SIZE);
bytes[0] = self.b3;
bytes[1] = self.b2;
bytes[2] = self.b1;
bytes[3] = self.b0;
#[cfg(feature = "5bytes")]
{
bytes[4] = self.b4;
}
}
/// Deserialize from bytes in the .idx file format.
pub fn from_bytes(bytes: &[u8]) -> Self {
assert!(bytes.len() >= OFFSET_SIZE);
Offset {
b3: bytes[0],
b2: bytes[1],
b1: bytes[2],
b0: bytes[3],
#[cfg(feature = "5bytes")]
b4: bytes[4],
}
}
pub fn is_zero(&self) -> bool {
#[cfg(feature = "5bytes")]
{
self.b0 == 0 && self.b1 == 0 && self.b2 == 0 && self.b3 == 0 && self.b4 == 0
}
#[cfg(not(feature = "5bytes"))]
{
self.b0 == 0 && self.b1 == 0 && self.b2 == 0 && self.b3 == 0
}
}
}
impl fmt::Display for Offset {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.to_actual_offset())
}
}
// ============================================================================
// DiskType
// ============================================================================
#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
pub enum DiskType {
#[default]
HardDrive,
Ssd,
Custom(String),
}
impl DiskType {
pub fn from_string(s: &str) -> Self {
match s.to_lowercase().as_str() {
"" | "hdd" => DiskType::HardDrive,
"ssd" => DiskType::Ssd,
other => DiskType::Custom(other.to_string()),
}
}
pub fn readable_string(&self) -> &str {
match self {
DiskType::HardDrive => "hdd",
DiskType::Ssd => "ssd",
DiskType::Custom(s) => s,
}
}
}
impl fmt::Display for DiskType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
DiskType::HardDrive => write!(f, ""),
DiskType::Ssd => write!(f, "ssd"),
DiskType::Custom(s) => write!(f, "{}", s),
}
}
}
// ============================================================================
// VolumeId
// ============================================================================
/// Volume identifier, stored as u32.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Default)]
pub struct VolumeId(pub u32);
impl VolumeId {
pub fn parse(s: &str) -> Result<Self, std::num::ParseIntError> {
s.parse::<u32>().map(VolumeId)
}
pub fn next(&self) -> VolumeId {
VolumeId(self.0 + 1)
}
}
impl fmt::Display for VolumeId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<u32> for VolumeId {
fn from(v: u32) -> Self {
VolumeId(v)
}
}
// ============================================================================
// Version
// ============================================================================
/// Needle storage format version.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Version(pub u8);
pub const VERSION_1: Version = Version(1);
pub const VERSION_2: Version = Version(2);
pub const VERSION_3: Version = Version(3);
impl Version {
pub fn current() -> Self {
VERSION_3
}
pub fn is_supported(&self) -> bool {
self.0 >= 1 && self.0 <= 3
}
}
impl Default for Version {
fn default() -> Self {
VERSION_3
}
}
impl From<u8> for Version {
fn from(v: u8) -> Self {
Version(v)
}
}
// ============================================================================
// ReadOption
// ============================================================================
/// Options controlling needle read behavior, matching Go's `ReadOption` in store.go.
///
/// Fields are split into request-side options (set by the caller) and response-side
/// flags (set during the read to communicate status back).
#[derive(Debug, Clone, Default)]
pub struct ReadOption {
// -- request --
/// If true, allow reading needles that have been soft-deleted.
pub read_deleted: bool,
/// If true, attempt to read only metadata for large needles (> PagedReadLimit).
pub attempt_meta_only: bool,
/// If true, the caller requires metadata only (no data payload).
pub must_meta_only: bool,
// -- response --
/// Set to true when the read actually returned metadata only.
pub is_meta_only: bool,
/// Compaction revision at the time of the read (for consistency during streaming).
pub volume_revision: u16,
/// Set to true when the offset exceeded MaxPossibleVolumeSize (4-byte offset wrap).
pub is_out_of_range: bool,
// -- slow-read / streaming --
/// When true, the read lock is acquired and released per chunk instead of held
/// for the entire read, reducing write latency at the cost of higher read P99.
pub has_slow_read: bool,
/// Buffer size for chunked streaming reads (used with `has_slow_read`).
pub read_buffer_size: i32,
}
// ============================================================================
// NeedleMapEntry helpers (for .idx file)
// ============================================================================
/// Parse a single .idx file entry (17 bytes) into (NeedleId, Offset, Size).
pub fn idx_entry_from_bytes(bytes: &[u8]) -> (NeedleId, Offset, Size) {
assert!(bytes.len() >= NEEDLE_MAP_ENTRY_SIZE);
let key = NeedleId::from_bytes(&bytes[..NEEDLE_ID_SIZE]);
let offset = Offset::from_bytes(&bytes[NEEDLE_ID_SIZE..NEEDLE_ID_SIZE + OFFSET_SIZE]);
let size = Size::from_bytes(
&bytes[NEEDLE_ID_SIZE + OFFSET_SIZE..NEEDLE_ID_SIZE + OFFSET_SIZE + SIZE_SIZE],
);
(key, offset, size)
}
/// Write a single .idx file entry (17 bytes).
pub fn idx_entry_to_bytes(bytes: &mut [u8], key: NeedleId, offset: Offset, size: Size) {
assert!(bytes.len() >= NEEDLE_MAP_ENTRY_SIZE);
key.to_bytes(&mut bytes[..NEEDLE_ID_SIZE]);
offset.to_bytes(&mut bytes[NEEDLE_ID_SIZE..NEEDLE_ID_SIZE + OFFSET_SIZE]);
size.to_bytes(
&mut bytes[NEEDLE_ID_SIZE + OFFSET_SIZE..NEEDLE_ID_SIZE + OFFSET_SIZE + SIZE_SIZE],
);
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_needle_id_round_trip() {
let id = NeedleId(0x123456789abcdef0);
let mut buf = [0u8; 8];
id.to_bytes(&mut buf);
let id2 = NeedleId::from_bytes(&buf);
assert_eq!(id, id2);
}
#[test]
fn test_needle_id_display() {
let id = NeedleId(255);
assert_eq!(id.to_string(), "ff");
}
#[test]
fn test_needle_id_parse() {
let id = NeedleId::parse("ff").unwrap();
assert_eq!(id, NeedleId(255));
}
#[test]
fn test_cookie_round_trip() {
let cookie = Cookie(0xdeadbeef);
let mut buf = [0u8; 4];
cookie.to_bytes(&mut buf);
let cookie2 = Cookie::from_bytes(&buf);
assert_eq!(cookie, cookie2);
}
#[test]
fn test_size_semantics() {
assert!(Size(100).is_valid());
assert!(!Size(100).is_deleted());
assert!(!Size(100).is_tombstone());
assert_eq!(Size(100).raw(), 100);
assert!(Size(-50).is_deleted());
assert!(!Size(-50).is_tombstone());
assert_eq!(Size(-50).raw(), 50);
assert!(TOMBSTONE_FILE_SIZE.is_deleted());
assert!(TOMBSTONE_FILE_SIZE.is_tombstone());
assert_eq!(TOMBSTONE_FILE_SIZE.raw(), 0);
assert!(!Size(0).is_valid());
assert!(!Size(0).is_deleted());
}
#[test]
fn test_size_round_trip() {
let size = Size(12345);
let mut buf = [0u8; 4];
size.to_bytes(&mut buf);
let size2 = Size::from_bytes(&buf);
assert_eq!(size, size2);
}
#[test]
fn test_size_negative_round_trip() {
// Negative sizes round-trip through u32 bit pattern
let size = Size(-50);
let mut buf = [0u8; 4];
size.to_bytes(&mut buf);
let size2 = Size::from_bytes(&buf);
assert_eq!(size, size2);
}
#[test]
fn test_offset_round_trip() {
// Test with a known actual offset
let actual_offset: i64 = 8 * 1000000; // must be multiple of 8
let offset = Offset::from_actual_offset(actual_offset);
assert_eq!(offset.to_actual_offset(), actual_offset);
// Test byte serialization
let mut buf = [0u8; 5];
offset.to_bytes(&mut buf);
let offset2 = Offset::from_bytes(&buf);
assert_eq!(offset.to_actual_offset(), offset2.to_actual_offset());
}
#[test]
fn test_offset_zero() {
let offset = Offset::default();
assert!(offset.is_zero());
assert_eq!(offset.to_actual_offset(), 0);
}
#[test]
fn test_offset_max() {
// Max stored value depends on offset size
#[cfg(feature = "5bytes")]
let max_stored: i64 = (1i64 << 40) - 1; // 5-byte max
#[cfg(not(feature = "5bytes"))]
let max_stored: i64 = (1i64 << 32) - 1; // 4-byte max
let max_actual = max_stored * NEEDLE_PADDING_SIZE as i64;
let offset = Offset::from_actual_offset(max_actual);
assert_eq!(offset.to_actual_offset(), max_actual);
}
#[test]
fn test_offset_size_constants() {
#[cfg(feature = "5bytes")]
{
assert_eq!(OFFSET_SIZE, 5);
assert_eq!(NEEDLE_MAP_ENTRY_SIZE, 17); // 8 + 5 + 4
assert_eq!(MAX_POSSIBLE_VOLUME_SIZE, 4 * 1024 * 1024 * 1024 * 8 * 256);
// 8TB
}
#[cfg(not(feature = "5bytes"))]
{
assert_eq!(OFFSET_SIZE, 4);
assert_eq!(NEEDLE_MAP_ENTRY_SIZE, 16); // 8 + 4 + 4
assert_eq!(MAX_POSSIBLE_VOLUME_SIZE, 4 * 1024 * 1024 * 1024 * 8); // 32GB
}
}
#[test]
fn test_idx_entry_round_trip() {
let key = NeedleId(0xdeadbeef12345678);
let offset = Offset::from_actual_offset(8 * 999);
let size = Size(4096);
let mut buf = [0u8; NEEDLE_MAP_ENTRY_SIZE];
idx_entry_to_bytes(&mut buf, key, offset, size);
let (key2, offset2, size2) = idx_entry_from_bytes(&buf);
assert_eq!(key, key2);
assert_eq!(offset.to_actual_offset(), offset2.to_actual_offset());
assert_eq!(size, size2);
}
#[test]
fn test_volume_id() {
let vid = VolumeId::parse("42").unwrap();
assert_eq!(vid, VolumeId(42));
assert_eq!(vid.to_string(), "42");
assert_eq!(vid.next(), VolumeId(43));
}
#[test]
fn test_version() {
assert!(VERSION_1.is_supported());
assert!(VERSION_2.is_supported());
assert!(VERSION_3.is_supported());
assert!(!Version(0).is_supported());
assert!(!Version(4).is_supported());
assert_eq!(Version::current(), VERSION_3);
}
#[test]
fn test_disk_type() {
assert_eq!(DiskType::from_string(""), DiskType::HardDrive);
assert_eq!(DiskType::from_string("hdd"), DiskType::HardDrive);
assert_eq!(DiskType::from_string("SSD"), DiskType::Ssd);
assert_eq!(
DiskType::from_string("nvme"),
DiskType::Custom("nvme".to_string())
);
assert_eq!(DiskType::HardDrive.readable_string(), "hdd");
assert_eq!(DiskType::Ssd.readable_string(), "ssd");
}
#[test]
fn test_read_option_default() {
let ro = ReadOption::default();
assert!(!ro.read_deleted);
assert!(!ro.attempt_meta_only);
assert!(!ro.must_meta_only);
assert!(!ro.is_meta_only);
assert_eq!(ro.volume_revision, 0);
assert!(!ro.is_out_of_range);
assert!(!ro.has_slow_read);
assert_eq!(ro.read_buffer_size, 0);
}
#[test]
fn test_read_option_custom() {
let ro = ReadOption {
read_deleted: true,
attempt_meta_only: true,
has_slow_read: true,
read_buffer_size: 1024 * 1024,
..ReadOption::default()
};
assert!(ro.read_deleted);
assert!(ro.attempt_meta_only);
assert!(!ro.must_meta_only);
assert!(!ro.is_meta_only);
assert!(ro.has_slow_read);
assert_eq!(ro.read_buffer_size, 1024 * 1024);
}
#[test]
fn test_read_option_clone() {
let ro = ReadOption {
is_out_of_range: true,
volume_revision: 42,
..ReadOption::default()
};
let ro2 = ro.clone();
assert!(ro2.is_out_of_range);
assert_eq!(ro2.volume_revision, 42);
}
}

View File

@@ -0,0 +1,344 @@
//! Rebuild a missing .idx from the .dat it indexes. Mirrors
//! `weed/storage/volume_idx_rebuild.go`.
use std::fs::{self, OpenOptions};
use std::io::{BufWriter, Write};
use std::path::Path;
use crate::storage::idx;
use crate::storage::needle::Needle;
use crate::storage::super_block::SuperBlock;
use crate::storage::types::*;
use crate::storage::volume::{
Volume, VolumeError, VolumeFileVisitor, fsync_dir, needle_disk_end, scan_volume_file,
};
/// Writes one .idx row per .dat record, in .dat append order, which is the
/// shape the volume server's own writes leave behind.
struct VolumeFileScanner4RebuildIdx<W: Write> {
writer: W,
dat_size: i64,
version: Version,
stopped: bool,
}
impl<W: Write> VolumeFileVisitor for VolumeFileScanner4RebuildIdx<W> {
fn visit_super_block(&mut self, _sb: &SuperBlock) -> Result<(), VolumeError> {
Ok(())
}
fn read_needle_body(&self) -> bool {
false
}
fn visit_needle(&mut self, n: &Needle, offset: i64) -> Result<(), VolumeError> {
// A record reaching past the end of .dat is a torn append or a corrupt
// header: nothing beyond it is indexable, and a row pointing past EOF
// would fail every read of that needle. The all-zero header case ends
// the walk upstream.
if self.stopped {
return Ok(());
}
if needle_disk_end(Offset::from_actual_offset(offset), n.size, self.version) > self.dat_size
{
self.stopped = true;
return Ok(());
}
let size = if n.size.is_valid() {
n.size
} else {
TOMBSTONE_FILE_SIZE
};
idx::write_index_entry(
&mut self.writer,
n.id,
Offset::from_actual_offset(offset),
size,
)?;
Ok(())
}
}
impl Volume {
/// Regenerate the volume's .idx from its .dat. The whole index is derivable
/// from the data file, so a volume whose index directory has no .idx -- a
/// --dir.idx pointed at an empty directory, or a lost index -- comes back on
/// its own instead of mounting with every needle invisible. The rows go to a
/// temp file that is renamed in, so an interrupted rebuild leaves no partial
/// index behind.
pub(crate) fn rebuild_idx_file(&self) -> Result<(), VolumeError> {
let idx_path = self.file_name(".idx");
let dat_path = self.file_name(".dat");
let tmp_path = format!("{idx_path}.tmp");
let rebuild = || -> Result<(), VolumeError> {
if let Some(parent) = Path::new(&idx_path).parent() {
fs::create_dir_all(parent)?;
}
let tmp_file = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(&tmp_path)?;
let mut scanner = VolumeFileScanner4RebuildIdx {
writer: BufWriter::new(&tmp_file),
dat_size: fs::metadata(&dat_path)?.len() as i64,
version: self.version(),
stopped: false,
};
scan_volume_file(&dat_path, &mut scanner)?;
scanner.writer.flush()?;
drop(scanner);
tmp_file.sync_all()?;
fs::rename(&tmp_path, &idx_path)?;
fsync_dir(&idx_path)?;
Ok(())
};
let result = rebuild();
if result.is_err() {
let _ = fs::remove_file(&tmp_path);
}
result
}
}
#[cfg(test)]
mod tests {
use crate::storage::needle::Needle;
use crate::storage::needle::crc::CRC;
use crate::storage::needle_map::NeedleMapKind;
use crate::storage::types::*;
use crate::storage::volume::{Volume, VolumeSpec};
use std::fs;
use std::path::Path;
use tempfile::TempDir;
fn needle(id: u64) -> Needle {
let data = format!("payload-{id}").into_bytes();
Needle {
id: NeedleId(id),
cookie: Cookie(0x55),
data_size: data.len() as u32,
checksum: CRC::new(&data),
data,
..Needle::default()
}
}
// Pointing --dir.idx at a directory with no .idx used to mount the volume on
// an empty index; the index is derivable from the .dat, so it must be
// rebuilt in place instead.
#[test]
fn test_load_moved_idx_directory_rebuilds_idx() {
let root = TempDir::new().unwrap();
let data_dir = root.path().join("data");
let old_idx_dir = root.path().join("idxA");
let new_idx_dir = root.path().join("idxB");
for dir in [&data_dir, &old_idx_dir, &new_idx_dir] {
fs::create_dir_all(dir).unwrap();
}
let data = data_dir.to_str().unwrap();
let old_idx = old_idx_dir.to_str().unwrap();
let new_idx = new_idx_dir.to_str().unwrap();
let mut v = Volume::new(
data,
old_idx,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
for id in 1..=3 {
v.write_needle(&mut needle(id), true, false).unwrap();
}
v.delete_needle(&mut needle(2)).unwrap();
v.sync_to_disk().unwrap();
let (want_count, want_deleted) = (v.file_count(), v.deleted_count());
drop(v);
let seeded = fs::read(format!("{old_idx}/1.idx")).unwrap();
let reopened = Volume::new(
data,
new_idx,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
assert!(
Path::new(&format!("{new_idx}/1.idx")).exists(),
"idx not rebuilt in the new idx dir"
);
assert_eq!(reopened.file_count(), want_count);
assert_eq!(reopened.deleted_count(), want_deleted);
assert_eq!(
fs::read(format!("{new_idx}/1.idx")).unwrap(),
seeded,
"rebuilt idx differs from the one the server wrote"
);
for id in [1, 3] {
let mut got = needle(id);
got.data.clear();
reopened.read_needle(&mut got).unwrap();
assert_eq!(got.data, format!("payload-{id}").into_bytes());
}
}
// A .dat padded with zeros must not be indexed as needle 0 rows: the walk
// stops where the records do.
#[test]
fn test_rebuild_idx_stops_at_zero_padded_dat_tail() {
let root = TempDir::new().unwrap();
let dir = root.path().to_str().unwrap();
let mut v = Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
v.write_needle(&mut needle(1), true, false).unwrap();
v.sync_to_disk().unwrap();
drop(v);
let seeded = fs::read(format!("{dir}/1.idx")).unwrap();
let dat = fs::OpenOptions::new()
.write(true)
.open(format!("{dir}/1.dat"))
.unwrap();
let dat_size = dat.metadata().unwrap().len();
dat.set_len(dat_size + 4096).unwrap();
drop(dat);
fs::remove_file(format!("{dir}/1.idx")).unwrap();
let reopened = Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
drop(reopened);
assert_eq!(
fs::read(format!("{dir}/1.idx")).unwrap(),
seeded,
"the zero-padded tail leaked into the rebuilt idx"
);
}
// A .dat whose last append was torn mid-body must not gain an index row
// that points past the end of the file.
#[test]
fn test_rebuild_idx_skips_truncated_dat_tail() {
let root = TempDir::new().unwrap();
let dir = root.path().to_str().unwrap();
let mut v = Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
v.write_needle(&mut needle(1), true, false).unwrap();
v.sync_to_disk().unwrap();
let kept = fs::read(format!("{dir}/1.idx")).unwrap();
v.write_needle(&mut needle(2), true, false).unwrap();
v.sync_to_disk().unwrap();
drop(v);
// Chop the second needle's body, leaving its header intact.
let dat = fs::OpenOptions::new()
.write(true)
.open(format!("{dir}/1.dat"))
.unwrap();
let torn_size = dat.metadata().unwrap().len() - 8;
dat.set_len(torn_size).unwrap();
drop(dat);
fs::remove_file(format!("{dir}/1.idx")).unwrap();
let reopened = Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
drop(reopened);
assert_eq!(
fs::read(format!("{dir}/1.idx")).unwrap(),
kept,
"the torn record leaked into the rebuilt idx"
);
}
// A corrupt header carrying a negative size advances the .dat walk
// backwards, which cycles forever between it and the record before it.
#[test]
fn test_rebuild_idx_stops_at_negative_size_header() {
let root = TempDir::new().unwrap();
let dir = root.path().to_str().unwrap();
let mut v = Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
v.write_needle(&mut needle(1), true, false).unwrap();
v.sync_to_disk().unwrap();
let kept = fs::read(format!("{dir}/1.idx")).unwrap();
v.write_needle(&mut needle(2), true, false).unwrap();
v.sync_to_disk().unwrap();
drop(v);
let rows = fs::read(format!("{dir}/1.idx")).unwrap();
let (_, offset, _) = idx_entry_from_bytes(&rows[NEEDLE_MAP_ENTRY_SIZE..]);
let corrupt_at = offset.to_actual_offset();
// Overwrite the second needle's size field with a negative i32.
use std::io::{Seek, SeekFrom, Write};
let mut dat = fs::OpenOptions::new()
.write(true)
.open(format!("{dir}/1.dat"))
.unwrap();
dat.seek(SeekFrom::Start(
corrupt_at as u64 + COOKIE_SIZE as u64 + NEEDLE_ID_SIZE as u64,
))
.unwrap();
dat.write_all(&[0xff, 0xff, 0xf0, 0x00]).unwrap();
drop(dat);
fs::remove_file(format!("{dir}/1.idx")).unwrap();
// Pre-fix the rebuild walked backwards from here and never terminated.
let reopened = Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap();
drop(reopened);
assert_eq!(
fs::read(format!("{dir}/1.idx")).unwrap(),
kept,
"the corrupt record leaked into the rebuilt idx"
);
}
}

View File

@@ -0,0 +1,418 @@
//! Recovery for .idx rows that deletes on a tiered read-only volume
//! overwrote. Mirrors `weed/storage/volume_idx_repair.go`.
use std::collections::HashMap;
use std::fs::{self, File, OpenOptions};
use std::io::{self, BufReader, Write};
use std::path::Path;
use tracing::info;
use crate::storage::idx;
use crate::storage::needle::Needle;
use crate::storage::needle::needle::needle_body_length;
use crate::storage::types::*;
use crate::storage::volume::{Volume, VolumeError, fsync_dir};
/// Needles found in the head of .dat, keyed by id, plus the ids in .dat order.
type DatHeadNeedles = (HashMap<NeedleId, (Offset, Size)>, Vec<NeedleId>);
impl Volume {
/// Restore the .idx rows that deletes on a tiered read-only volume used to
/// overwrite.
///
/// The sorted-file needle map opened .idx read-write but never seeded its
/// write position, so a delete wrote its (key, offset 0, tombstone) row at
/// .idx offset 0 and advanced one row at a time instead of appending. Every
/// delete therefore replaced one more row at the front of .idx, and the rows
/// it replaced -- the Put rows indexing the first needles in .dat -- were
/// lost. Reads for those needles return not-found even though .dat still
/// holds them intact.
///
/// The damage leaves a fingerprint: .idx begins with a run of offset-0
/// tombstones. A healthy .idx never does. Its first row is the Put for the
/// first needle in .dat, and an offset-0 tombstone -- a delete against a
/// tiered volume, which appends no .dat record and so has no extent to point
/// at -- can only ever land at the tail.
///
/// .idx and .dat grow in lockstep, so the clobbered rows indexed exactly the
/// first N records of .dat. Re-deriving them is a header-only walk over the
/// head of .dat, which stays cheap even when .dat is served from a remote
/// tier.
///
/// The recovered rows go back in front, where the rows they replace used to
/// sit, and every existing row keeps its relative order behind them. That
/// restores .idx to .dat append order, which several readers lean on: the
/// fingerprint is gone so a later load stops after one row, and the last row
/// is the .dat-tail needle again.
pub(crate) fn repair_idx_head_tombstones(&self) -> Result<usize, VolumeError> {
if !self.has_data_backend() {
return Ok(0);
}
let version = self.version();
if !version.is_supported() {
return Ok(0);
}
let first_needle_offset = self.super_block.block_size() as i64;
let idx_path = self.file_name(".idx");
let clobbered = idx_head_tombstone_count(&idx_path)?;
if clobbered == 0 {
return Ok(0);
}
info!(
volume_id = self.id.0,
idx = %idx_path,
clobbered,
"idx starts with offset-0 tombstones, recovering the rows they overwrote from .dat"
);
let (mut lost, order) = self.scan_dat_head(version, first_needle_offset, clobbered)?;
drop_indexed_keys(&idx_path, &mut lost)?;
if lost.is_empty() {
return Ok(0);
}
let mut rows = Vec::with_capacity(lost.len() * NEEDLE_MAP_ENTRY_SIZE);
let mut restored = 0;
for key in order {
let Some((offset, size)) = lost.get(&key).copied() else {
continue;
};
idx::write_index_entry(&mut rows, key, offset, size)?;
restored += 1;
}
prepend_idx_rows(&idx_path, &rows)?;
Ok(restored)
}
/// Read the headers of the first `limit` records of .dat and return the
/// needles they hold, in .dat order. A record with an invalid size is a
/// delete marker, so the key it names drops out of the result rather than
/// being resurrected.
fn scan_dat_head(
&self,
version: Version,
first_needle_offset: i64,
limit: usize,
) -> Result<DatHeadNeedles, VolumeError> {
let mut found: HashMap<NeedleId, (Offset, Size)> = HashMap::new();
let mut order: Vec<NeedleId> = Vec::new();
let mut offset = first_needle_offset;
for _ in 0..limit {
let mut header = [0u8; NEEDLE_HEADER_SIZE];
match self.read_exact_at_backend(&mut header, offset as u64) {
Ok(()) => {}
Err(VolumeError::Io(ref e)) if e.kind() == io::ErrorKind::UnexpectedEof => break,
Err(e) => return Err(e),
}
let (_cookie, id, size) = Needle::parse_header(&header);
if size.0 == 0 && id.is_empty() {
break;
}
if size.is_valid() {
if found
.insert(id, (Offset::from_actual_offset(offset), size))
.is_none()
{
order.push(id);
}
} else {
found.remove(&id);
}
offset += NEEDLE_HEADER_SIZE as i64 + needle_body_length(size, version);
}
Ok((found, order))
}
}
/// Report how many rows at the front of .idx are offset-0 tombstones. A healthy
/// .idx answers 0 on its first row.
fn idx_head_tombstone_count(idx_path: &str) -> Result<usize, VolumeError> {
if !Path::new(idx_path).exists() {
return Ok(0);
}
let mut reader = BufReader::new(File::open(idx_path)?);
let mut clobbered = 0usize;
let walk = idx::walk_index_file(&mut reader, 0, |_key, offset, size| {
if !offset.is_zero() || !size.is_tombstone() {
return Err(stop_walk());
}
clobbered += 1;
Ok(())
});
finish_walk(walk)?;
Ok(clobbered)
}
/// Remove every candidate the .idx already names, whether by a Put row or a
/// tombstone. What is left is only what the clobbered rows held.
fn drop_indexed_keys(
idx_path: &str,
candidates: &mut HashMap<NeedleId, (Offset, Size)>,
) -> Result<(), VolumeError> {
if candidates.is_empty() {
return Ok(());
}
let mut reader = BufReader::new(File::open(idx_path)?);
let walk = idx::walk_index_file(&mut reader, 0, |key, _offset, _size| {
candidates.remove(&key);
if candidates.is_empty() {
return Err(stop_walk());
}
Ok(())
});
finish_walk(walk)
}
/// Rewrite .idx as `rows` followed by its current contents, through a temp file
/// and a rename so a crash mid-write never leaves a partial index at the live
/// name.
fn prepend_idx_rows(idx_path: &str, rows: &[u8]) -> Result<(), VolumeError> {
let tmp_path = format!("{}.tmp", idx_path);
let mut src = File::open(idx_path)?;
let commit = (|| -> io::Result<()> {
let src_perm = src.metadata()?.permissions();
let mut dst = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(&tmp_path)?;
// The rename replaces .idx with this file, so it has to carry the mode
// the index already had rather than whatever the umask allows.
dst.set_permissions(src_perm)?;
dst.write_all(rows)?;
io::copy(&mut src, &mut dst)?;
dst.sync_all()?;
drop(dst);
fs::rename(&tmp_path, idx_path)?;
fsync_dir(idx_path)
})();
if commit.is_err() {
let _ = fs::remove_file(&tmp_path);
}
Ok(commit?)
}
/// Sentinel that ends an `idx::walk_index_file` early once the answer is known.
fn stop_walk() -> io::Error {
io::Error::new(io::ErrorKind::Interrupted, "stop idx walk")
}
fn finish_walk(walk: io::Result<()>) -> Result<(), VolumeError> {
match walk {
Err(ref e) if e.kind() == io::ErrorKind::Interrupted => Ok(()),
other => Ok(other?),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::needle::crc::CRC;
use crate::storage::needle_map::NeedleMapKind;
use crate::storage::volume::VolumeSpec;
use std::os::unix::fs::{FileExt, PermissionsExt};
use tempfile::TempDir;
fn open_volume(dir: &str) -> Volume {
Volume::new(
dir,
dir,
VolumeId(1),
NeedleMapKind::InMemory,
&VolumeSpec::default(),
)
.unwrap()
}
fn write_test_volume(dir: &str, needle_count: u64) -> Vec<Vec<u8>> {
let mut v = open_volume(dir);
let mut written = Vec::new();
for i in 1..=needle_count {
let data = format!("needle-{}-payload", i).into_bytes();
let mut n = Needle {
id: NeedleId(i),
cookie: Cookie(0x1234_0000 + i as u32),
data_size: data.len() as u32,
checksum: CRC::new(&data),
data: data.clone(),
..Needle::default()
};
v.write_needle(&mut n, true, false).unwrap();
written.push(data);
}
written
}
/// Reproduce the damage a delete on a tiered read-only volume used to do: it
/// writes (key, offset 0, tombstone) rows over the front of .idx instead of
/// appending them.
fn clobber_idx_head(idx_path: &str, keys: &[u64]) {
let file = OpenOptions::new().write(true).open(idx_path).unwrap();
for (i, key) in keys.iter().enumerate() {
let mut row = Vec::new();
idx::write_index_entry(
&mut row,
NeedleId(*key),
Offset::from_actual_offset(0),
TOMBSTONE_FILE_SIZE,
)
.unwrap();
file.write_at(&row, (i * NEEDLE_MAP_ENTRY_SIZE) as u64)
.unwrap();
}
}
fn idx_size(idx_path: &str) -> u64 {
std::fs::metadata(idx_path).unwrap().len()
}
fn idx_rows(idx_path: &str) -> Vec<(NeedleId, i64, Size)> {
let mut reader = BufReader::new(File::open(idx_path).unwrap());
let mut rows = Vec::new();
idx::walk_index_file(&mut reader, 0, |key, offset, size| {
rows.push((key, offset.to_actual_offset(), size));
Ok(())
})
.unwrap();
rows
}
fn read_needle_data(v: &Volume, id: u64) -> Result<Vec<u8>, VolumeError> {
let mut n = Needle {
id: NeedleId(id),
..Needle::default()
};
v.read_needle(&mut n)?;
Ok(n.data)
}
#[test]
fn test_repair_idx_head_tombstones_restores_clobbered_rows() {
let tmp = TempDir::new().unwrap();
let dir = tmp.path().to_str().unwrap();
let idx_path = format!("{}/1.idx", dir);
let written = write_test_volume(dir, 12);
let size_before = idx_size(&idx_path);
// The rewrite replaces .idx wholesale, so it must not widen the mode.
fs::set_permissions(&idx_path, fs::Permissions::from_mode(0o600)).unwrap();
// Deletes against needles 9..12 land on the front of .idx and take the
// rows indexing needles 1..4 with them.
clobber_idx_head(&idx_path, &[9, 10, 11, 12]);
let v = open_volume(dir);
for i in 1..=8u64 {
assert_eq!(
read_needle_data(&v, i).unwrap(),
written[(i - 1) as usize],
"needle {} not recovered",
i
);
}
let want = size_before + 4 * NEEDLE_MAP_ENTRY_SIZE as u64;
assert_eq!(idx_size(&idx_path), want, "idx size after recovery");
assert_eq!(
fs::metadata(&idx_path).unwrap().permissions().mode() & 0o777,
0o600,
"idx mode after recovery"
);
// The recovered rows go back in front, so .idx is in .dat append order
// again: the fingerprint is gone and the last row is still the .dat tail.
let rows = idx_rows(&idx_path);
assert_eq!(
rows[0].1,
v.super_block.block_size() as i64,
"first row does not index the first needle in .dat"
);
for i in 1..4 {
assert!(
rows[i - 1].1 < rows[i].1,
"recovered rows are not in .dat order: {:?} then {:?}",
rows[i - 1],
rows[i]
);
}
// The recovery is idempotent: a second load finds nothing left to restore.
drop(v);
let _v = open_volume(dir);
assert_eq!(idx_size(&idx_path), want, "idx grew on second load");
}
#[test]
fn test_repair_idx_head_tombstones_leaves_healthy_idx_alone() {
let tmp = TempDir::new().unwrap();
let dir = tmp.path().to_str().unwrap();
let idx_path = format!("{}/1.idx", dir);
write_test_volume(dir, 6);
let mut v = open_volume(dir);
v.delete_needle(&mut Needle {
id: NeedleId(3),
..Needle::default()
})
.unwrap();
drop(v);
let size_before = idx_size(&idx_path);
let v = open_volume(dir);
assert_eq!(
idx_size(&idx_path),
size_before,
"healthy idx was rewritten"
);
assert!(
matches!(
read_needle_data(&v, 3),
Err(VolumeError::Deleted) | Err(VolumeError::NotFound)
),
"needle 3 should stay deleted"
);
}
/// A needle deleted before the damage must not be resurrected: its tombstone
/// row survives at the tail of .idx, so the key is still indexed and stays
/// out of the recovery.
#[test]
fn test_repair_idx_head_tombstones_keeps_deleted_needles_deleted() {
let tmp = TempDir::new().unwrap();
let dir = tmp.path().to_str().unwrap();
let idx_path = format!("{}/1.idx", dir);
write_test_volume(dir, 8);
let mut v = open_volume(dir);
v.delete_needle(&mut Needle {
id: NeedleId(2),
..Needle::default()
})
.unwrap();
drop(v);
clobber_idx_head(&idx_path, &[7, 8]);
let v = open_volume(dir);
assert!(
matches!(
read_needle_data(&v, 2),
Err(VolumeError::Deleted) | Err(VolumeError::NotFound)
),
"needle 2 was resurrected"
);
assert!(
read_needle_data(&v, 1).is_ok(),
"needle 1 should have been recovered"
);
}
}

View File

@@ -0,0 +1,27 @@
//! Opening volume data and index files without access-time updates.
use std::fs::{File, OpenOptions};
use std::io;
use std::path::Path;
/// Open a volume data or index file with `O_NOATIME`. Nothing reads these
/// files' atime, but without the flag every needle read dirties the inode —
/// even relatime writes atime on the first read after each write, so an
/// actively written volume pays a metadata write per read/write cycle.
/// Matches Go's `backend.OpenVolumeFile`.
pub fn open_volume_file(opts: &OpenOptions, path: impl AsRef<Path>) -> io::Result<File> {
let path = path.as_ref();
#[cfg(target_os = "linux")]
{
use std::os::unix::fs::OpenOptionsExt;
let mut noatime = opts.clone();
noatime.custom_flags(libc::O_NOATIME);
match noatime.open(path) {
// O_NOATIME is refused unless we own the file or hold CAP_FOWNER.
Err(e) if e.raw_os_error() == Some(libc::EPERM) => opts.open(path),
result => result,
}
}
#[cfg(not(target_os = "linux"))]
opts.open(path)
}

View File

@@ -0,0 +1,130 @@
//! Mirror of `weed/storage/store_volume_report.go`.
use std::collections::HashMap;
use std::sync::Mutex;
/// Identifies one reported copy. Keyed by disk as well as id because a volume
/// id can be mounted on two disks, and reporting one of them would leave the
/// other's changes untold.
pub type VolumeReportKey = (u32, u32);
/// What the master was told about one volume copy: the hash that detects
/// change, and the heartbeat pass that last found the copy held.
#[derive(Clone, Copy)]
struct ReportedVolume {
hash: u64,
pass: u64,
}
/// Everything a heartbeat reads and writes about what the master was told,
/// under one lock. The full-list flag and the generation that answers it have
/// to move together: split across atomics, a request landing between two of
/// them is answered by a heartbeat that never carried a list. Go holds a single
/// mutex over the same fields.
#[derive(Default)]
struct ReportState {
/// Set once the master says it compares digests. Until then the whole list
/// goes every time, which is what an older master needs.
deltas_accepted: bool,
full_list_needed: bool,
/// Counts requests for the whole list, so one arriving while a heartbeat is
/// being built is not marked satisfied by it.
full_list_generation: u64,
/// Numbers heartbeats, so one can mark the copies it finds held without
/// building a second map of them.
pass: u64,
last_reported: HashMap<VolumeReportKey, ReportedVolume>,
}
/// Remembers what the master was last told about each volume, so a heartbeat
/// can carry only what moved since.
///
/// Per-connection: a server that reconnects, or reaches a different master,
/// knows nothing about what that master holds and starts again from the full
/// list. The default has told no master anything, so it sends the whole list
/// until one accepts changes.
#[derive(Default)]
pub struct VolumeReportState {
state: Mutex<ReportState>,
}
impl VolumeReportState {
/// Drops everything known about the master's view.
pub fn reset(&self) {
let mut state = self.state.lock().unwrap();
state.deltas_accepted = false;
state.full_list_needed = true;
state.full_list_generation += 1;
state.last_reported.clear();
}
pub fn accept_deltas(&self) {
self.state.lock().unwrap().deltas_accepted = true;
}
pub fn request_full_list(&self) {
let mut state = self.state.lock().unwrap();
state.full_list_needed = true;
state.full_list_generation += 1;
}
/// Opens a heartbeat: whether it must carry the whole list, the request it
/// answers, and the pass number that marks the copies it finds still held.
pub fn begin(&self) -> (bool, u64, u64) {
let mut state = self.state.lock().unwrap();
state.pass += 1;
(
state.full_list_needed || !state.deltas_accepted,
state.full_list_generation,
state.pass,
)
}
/// Reports whether the master needs telling about this volume, given what
/// it was last told. For a caller that is only taking a snapshot and so
/// must leave the reporting state alone; a heartbeat calls `record`.
pub fn changed(&self, key: VolumeReportKey, hash: u64) -> bool {
self.state
.lock()
.unwrap()
.last_reported
.get(&key)
.is_none_or(|previous| previous.hash != hash)
}
/// Marks one volume copy as held by the heartbeat being built, and reports
/// whether the master needs telling about it. It updates the entry already
/// held rather than build a second map beside it, so a server whose volumes
/// are quiet allocates nothing per volume per heartbeat.
pub fn record(&self, key: VolumeReportKey, hash: u64, pass: u64) -> bool {
let mut state = self.state.lock().unwrap();
match state.last_reported.get_mut(&key) {
Some(previous) => {
let changed = previous.hash != hash;
previous.hash = hash;
previous.pass = pass;
changed
}
None => {
state
.last_reported
.insert(key, ReportedVolume { hash, pass });
true
}
}
}
/// Closes the heartbeat. Copies this pass did not find are forgotten, so one
/// that comes back is reported again.
pub fn commit(&self, pass: u64, generation: u64) {
let mut state = self.state.lock().unwrap();
state
.last_reported
.retain(|_, reported| reported.pass == pass);
// A request that arrived while this heartbeat was being built asked
// about a later state than it carries, so it stands.
if state.full_list_generation == generation {
state.full_list_needed = false;
}
}
}

View File

@@ -0,0 +1,111 @@
//! Mirror of `weed/storage/volume_report_hash.go`.
//!
//! The master compares the digest a volume server reports against one it
//! computes itself, so this has to agree with the Go implementation
//! byte-for-byte. `report_hash_vectors` pins that against values produced by
//! the Go side; do not change the layout without regenerating them there.
use xxhash_rust::xxh64::xxh64;
use crate::pb::master_pb;
/// Digests everything a volume server reports about a volume.
///
/// It must cover every field of `VolumeInformationMessage`: a change the hash
/// misses is a change the master would never be told about.
pub fn report_hash(m: &master_pb::VolumeInformationMessage) -> u64 {
let mut buf = [0u8; 57];
buf[0..4].copy_from_slice(&m.id.to_le_bytes());
buf[4..12].copy_from_slice(&m.size.to_le_bytes());
buf[12..20].copy_from_slice(&m.file_count.to_le_bytes());
buf[20..28].copy_from_slice(&m.delete_count.to_le_bytes());
buf[28..36].copy_from_slice(&m.deleted_byte_count.to_le_bytes());
// The master stores these narrowed, so hash what it will hold, not what the
// wire type could carry.
buf[36..40].copy_from_slice(&((m.replica_placement as u8) as u32).to_le_bytes());
buf[40..44].copy_from_slice(&((m.version as u8) as u32).to_le_bytes());
buf[44..48].copy_from_slice(&normalize_ttl(m.ttl).to_le_bytes());
buf[48..52].copy_from_slice(&m.compact_revision.to_le_bytes());
buf[52..56].copy_from_slice(&m.disk_id.to_le_bytes());
if m.read_only {
buf[56] = 1;
}
if m.read_only_can_delete {
buf[56] |= 2;
}
let mut h = xxh64(&buf, 0);
h = fold(h, xxh64(&(m.modified_at_second as u64).to_le_bytes(), 0));
h = fold(h, xxh64(m.collection.as_bytes(), 0));
h = fold(h, xxh64(m.disk_type.as_bytes(), 0));
// Not the remote storage key: the master does not keep it, so a change to
// it alters nothing its copy holds.
h = fold(h, xxh64(m.remote_storage_name.as_bytes(), 0));
h
}
/// A ttl whose count is zero encodes as zero however the unit is set, matching
/// what the master stores after decoding it.
fn normalize_ttl(ttl: u32) -> u32 {
let count = (ttl >> 8) & 0xff;
if count == 0 {
return 0;
}
(count << 8) | (ttl & 0xff)
}
/// Combines two hashes order-dependently, so swapping two string fields is not
/// invisible.
fn fold(h: u64, x: u64) -> u64 {
let h = (h ^ x).wrapping_mul(0x9E37_79B9_7F4A_7C15);
h ^ (h >> 29)
}
#[cfg(test)]
mod tests {
use super::*;
// Produced by the Go implementation. If these drift, every volume server
// running this build reports a digest the master can never match, and falls
// back to sending its whole volume list forever.
#[test]
fn report_hash_vectors() {
let empty = master_pb::VolumeInformationMessage::default();
assert_eq!(report_hash(&empty), 10988706248825469653);
let one = master_pb::VolumeInformationMessage {
id: 1,
..Default::default()
};
assert_eq!(report_hash(&one), 2035849960016744285);
let full = master_pb::VolumeInformationMessage {
id: 42,
size: 1 << 30,
collection: "c".to_string(),
file_count: 7,
delete_count: 2,
deleted_byte_count: 99,
read_only: true,
read_only_can_delete: false,
replica_placement: 10,
version: 3,
ttl: 3 << 8,
compact_revision: 5,
modified_at_second: 1700000000,
remote_storage_name: "s3".to_string(),
remote_storage_key: "k/1.dat".to_string(),
disk_type: "ssd".to_string(),
disk_id: 2,
};
assert_eq!(report_hash(&full), 2748844479819636032);
}
#[test]
fn ttl_with_no_count_is_dropped() {
assert_eq!(normalize_ttl(0), 0);
assert_eq!(normalize_ttl(3), 0);
assert_eq!(normalize_ttl(3 << 8), 3 << 8);
assert_eq!(normalize_ttl((3 << 8) | 4), (3 << 8) | 4);
}
}

View File

@@ -0,0 +1,79 @@
//! Version helpers aligned with Go's util/version package.
use std::sync::OnceLock;
#[cfg(feature = "5bytes")]
const SIZE_LIMIT: &str = "8000GB"; // Matches Go production builds (5BytesOffset)
#[cfg(not(feature = "5bytes"))]
const SIZE_LIMIT: &str = "30GB"; // Matches Go default build (!5BytesOffset)
pub fn size_limit() -> &'static str {
SIZE_LIMIT
}
pub fn commit() -> &'static str {
option_env!("SEAWEEDFS_COMMIT")
.or(option_env!("GIT_COMMIT"))
.or(option_env!("GIT_SHA"))
.unwrap_or("")
}
pub fn version_number() -> &'static str {
static VERSION_NUMBER: OnceLock<String> = OnceLock::new();
VERSION_NUMBER
.get_or_init(|| {
parse_go_version_number().unwrap_or_else(|| env!("CARGO_PKG_VERSION").to_string())
})
.as_str()
}
pub fn version() -> &'static str {
static VERSION: OnceLock<String> = OnceLock::new();
VERSION
.get_or_init(|| format!("{} {}", size_limit(), version_number()))
.as_str()
}
pub fn full_version() -> &'static str {
static FULL: OnceLock<String> = OnceLock::new();
FULL.get_or_init(|| format!("{} {}", version(), commit()))
.as_str()
}
pub fn server_header() -> &'static str {
static HEADER: OnceLock<String> = OnceLock::new();
HEADER
.get_or_init(|| format!("SeaweedFS Volume {}", version()))
.as_str()
}
fn parse_go_version_number() -> Option<String> {
let src = include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../weed/util/version/constants.go"
));
let mut major: Option<u32> = None;
let mut minor: Option<u32> = None;
for line in src.lines() {
let l = line.trim();
if l.starts_with("MAJOR_VERSION") {
major = parse_int32_line(l);
} else if l.starts_with("MINOR_VERSION") {
minor = parse_int32_line(l);
}
if major.is_some() && minor.is_some() {
break;
}
}
match (major, minor) {
(Some(maj), Some(min)) => Some(format!("{}.{:02}", maj, min)),
_ => None,
}
}
fn parse_int32_line(line: &str) -> Option<u32> {
let start = line.find("int32(")? + "int32(".len();
let rest = &line[start..];
let end = rest.find(')')?;
rest[..end].trim().parse::<u32>().ok()
}