Browse Source
production sprint 4: EC batch delete, JPEG orientation, write queue, S3 tier move
production sprint 4: EC batch delete, JPEG orientation, write queue, S3 tier move
- BatchDelete now supports EC volumes: looks up needle in .ecx index, journals deletion to .ecj file (local-only, Go handles distributed part) - JPEG EXIF orientation auto-fix on upload using kamadak-exif + image crate, matching Go's FixJpgOrientation behavior (8 orientation transforms) - Async batched write processing via mpsc queue (up to 128 entries per batch), groups writes by volume ID and syncs once per volume per batch - VolumeTierMoveDatToRemote: multipart upload .dat file to S3 with progress streaming, updates .vif with remote file reference - VolumeTierMoveDatFromRemote: downloads .dat from S3 with progress streaming, removes remote file reference from .vif - S3TierRegistry for managing named remote storage backends - VolumeInfo (.vif) JSON persistence matching Go's protojson format - 124 lib + 7 integration = 131 Rust tests pass - All 109 Go integration tests pass (53 HTTP + 56 gRPC)rust-volume-server
Chris Lu
4 days ago
16 changed files with 1545 additions and 89 deletions
-
16seaweed-volume/Cargo.lock
-
1seaweed-volume/Cargo.toml
-
31seaweed-volume/DEV_PLAN.md
-
3seaweed-volume/src/config.rs
-
275seaweed-volume/src/images.rs
-
1seaweed-volume/src/lib.rs
-
12seaweed-volume/src/main.rs
-
1seaweed-volume/src/remote_storage/mod.rs
-
327seaweed-volume/src/remote_storage/s3_tier.rs
-
392seaweed-volume/src/server/grpc_server.rs
-
46seaweed-volume/src/server/handlers.rs
-
1seaweed-volume/src/server/mod.rs
-
5seaweed-volume/src/server/volume_server.rs
-
315seaweed-volume/src/server/write_queue.rs
-
204seaweed-volume/src/storage/volume.rs
-
4seaweed-volume/tests/http_integration.rs
@ -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);
|
||||
|
}
|
||||
|
}
|
||||
@ -0,0 +1,327 @@ |
|||||
|
//! 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::sync::Arc;
|
||||
|
|
||||
|
use aws_sdk_s3::Client;
|
||||
|
use aws_sdk_s3::config::{BehaviorVersion, Credentials, Region};
|
||||
|
use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart};
|
||||
|
use tokio::io::AsyncReadExt;
|
||||
|
|
||||
|
/// 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 progress reporting.
|
||||
|
///
|
||||
|
/// Returns (s3_key, file_size) on success.
|
||||
|
/// The progress callback receives (bytes_uploaded, percentage).
|
||||
|
pub async fn upload_file<F>(
|
||||
|
&self,
|
||||
|
file_path: &str,
|
||||
|
mut progress_fn: F,
|
||||
|
) -> Result<(String, u64), String>
|
||||
|
where
|
||||
|
F: FnMut(i64, f32) + 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
|
||||
|
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();
|
||||
|
|
||||
|
let mut file = tokio::fs::File::open(file_path)
|
||||
|
.await
|
||||
|
.map_err(|e| format!("failed to open file {}: {}", file_path, e))?;
|
||||
|
|
||||
|
let mut completed_parts = Vec::new();
|
||||
|
let mut offset: u64 = 0;
|
||||
|
let mut part_number: i32 = 1;
|
||||
|
|
||||
|
loop {
|
||||
|
let remaining = file_size - offset;
|
||||
|
if remaining == 0 {
|
||||
|
break;
|
||||
|
}
|
||||
|
let this_part_size = std::cmp::min(part_size, remaining) as usize;
|
||||
|
let mut buf = vec![0u8; this_part_size];
|
||||
|
file.read_exact(&mut buf)
|
||||
|
.await
|
||||
|
.map_err(|e| format!("failed to read file at offset {}: {}", offset, e))?;
|
||||
|
|
||||
|
let upload_part_resp = self
|
||||
|
.client
|
||||
|
.upload_part()
|
||||
|
.bucket(&self.bucket)
|
||||
|
.key(&key)
|
||||
|
.upload_id(&upload_id)
|
||||
|
.part_number(part_number)
|
||||
|
.body(buf.into())
|
||||
|
.send()
|
||||
|
.await
|
||||
|
.map_err(|e| {
|
||||
|
format!(
|
||||
|
"failed to upload part {} at offset {}: {}",
|
||||
|
part_number, offset, e
|
||||
|
)
|
||||
|
})?;
|
||||
|
|
||||
|
let e_tag = upload_part_resp.e_tag().unwrap_or_default().to_string();
|
||||
|
completed_parts.push(
|
||||
|
CompletedPart::builder()
|
||||
|
.e_tag(e_tag)
|
||||
|
.part_number(part_number)
|
||||
|
.build(),
|
||||
|
);
|
||||
|
|
||||
|
offset += this_part_size as u64;
|
||||
|
|
||||
|
// Report progress
|
||||
|
let pct = if file_size > 0 {
|
||||
|
(offset as f32 * 100.0) / file_size as f32
|
||||
|
} else {
|
||||
|
100.0
|
||||
|
};
|
||||
|
progress_fn(offset as i64, pct);
|
||||
|
|
||||
|
part_number += 1;
|
||||
|
}
|
||||
|
|
||||
|
// 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((key, file_size))
|
||||
|
}
|
||||
|
|
||||
|
/// Download a file from S3 to a local path with progress reporting.
|
||||
|
///
|
||||
|
/// Returns the file size on success.
|
||||
|
pub async fn download_file<F>(
|
||||
|
&self,
|
||||
|
dest_path: &str,
|
||||
|
key: &str,
|
||||
|
mut progress_fn: F,
|
||||
|
) -> Result<u64, String>
|
||||
|
where
|
||||
|
F: FnMut(i64, f32) + 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;
|
||||
|
|
||||
|
// Download in chunks
|
||||
|
let part_size: u64 = 64 * 1024 * 1024;
|
||||
|
let mut 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))?;
|
||||
|
|
||||
|
let mut offset: u64 = 0;
|
||||
|
|
||||
|
loop {
|
||||
|
if offset >= file_size {
|
||||
|
break;
|
||||
|
}
|
||||
|
let end = std::cmp::min(offset + part_size - 1, file_size - 1);
|
||||
|
let range = format!("bytes={}-{}", offset, end);
|
||||
|
|
||||
|
let get_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 = get_resp
|
||||
|
.body
|
||||
|
.collect()
|
||||
|
.await
|
||||
|
.map_err(|e| format!("failed to read body: {}", e))?;
|
||||
|
let bytes = body.into_bytes();
|
||||
|
|
||||
|
use tokio::io::AsyncWriteExt;
|
||||
|
file.write_all(&bytes)
|
||||
|
.await
|
||||
|
.map_err(|e| format!("failed to write to {}: {}", dest_path, e))?;
|
||||
|
|
||||
|
offset += bytes.len() as u64;
|
||||
|
|
||||
|
let pct = if file_size > 0 {
|
||||
|
(offset as f32 * 100.0) / file_size as f32
|
||||
|
} else {
|
||||
|
100.0
|
||||
|
};
|
||||
|
progress_fn(offset as i64, pct);
|
||||
|
}
|
||||
|
|
||||
|
use tokio::io::AsyncWriteExt;
|
||||
|
file.flush()
|
||||
|
.await
|
||||
|
.map_err(|e| format!("failed to flush {}: {}", dest_path, e))?;
|
||||
|
|
||||
|
Ok(file_size)
|
||||
|
}
|
||||
|
|
||||
|
/// 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(())
|
||||
|
}
|
||||
|
}
|
||||
|
|
||||
|
/// 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()
|
||||
|
}
|
||||
|
}
|
||||
@ -0,0 +1,315 @@ |
|||||
|
//! Async batched write processing for the volume server.
|
||||
|
//!
|
||||
|
//! Instead of each upload handler directly calling `write_needle` and syncing,
|
||||
|
//! writes are submitted to a queue. A background worker drains the queue in
|
||||
|
//! batches (up to 128 entries), groups them by volume ID, processes them
|
||||
|
//! together, and syncs once per volume for the entire batch.
|
||||
|
|
||||
|
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>;
|
||||
|
|
||||
|
/// A request to write a needle, submitted to the write queue.
|
||||
|
pub struct WriteRequest {
|
||||
|
pub volume_id: VolumeId,
|
||||
|
pub needle: Needle,
|
||||
|
pub response_tx: oneshot::Sender<WriteResult>,
|
||||
|
}
|
||||
|
|
||||
|
/// Maximum number of write requests to batch together.
|
||||
|
const MAX_BATCH_SIZE: usize = 128;
|
||||
|
|
||||
|
/// 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,
|
||||
|
) -> WriteResult {
|
||||
|
let (response_tx, response_rx) = oneshot::channel();
|
||||
|
let request = WriteRequest {
|
||||
|
volume_id,
|
||||
|
needle,
|
||||
|
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.
|
||||
|
let mut batch = Vec::with_capacity(MAX_BATCH_SIZE);
|
||||
|
batch.push(first);
|
||||
|
|
||||
|
while batch.len() < MAX_BATCH_SIZE {
|
||||
|
match self.rx.try_recv() {
|
||||
|
Ok(req) => 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, Vec<(Needle, oneshot::Sender<WriteResult>)>)> = 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.response_tx));
|
||||
|
} else {
|
||||
|
groups.push((vid, vec![(req.needle, 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, response_tx) in entries {
|
||||
|
let result = store.write_volume_needle(vid, &mut needle);
|
||||
|
// 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 std::sync::RwLock;
|
||||
|
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU32};
|
||||
|
use crate::security::{Guard, SigningKey};
|
||||
|
use crate::storage::store::Store;
|
||||
|
use crate::storage::needle_map::NeedleMapKind;
|
||||
|
|
||||
|
let store = Store::new(NeedleMapKind::InMemory);
|
||||
|
let guard = Guard::new(
|
||||
|
&[],
|
||||
|
SigningKey(vec![]),
|
||||
|
0,
|
||||
|
SigningKey(vec![]),
|
||||
|
0,
|
||||
|
);
|
||||
|
|
||||
|
Arc::new(VolumeServerState {
|
||||
|
store: RwLock::new(store),
|
||||
|
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,
|
||||
|
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()),
|
||||
|
})
|
||||
|
}
|
||||
|
|
||||
|
#[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).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).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).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).await;
|
||||
|
assert!(result.is_err()); // NotFound is fine -- the point is it doesn't panic
|
||||
|
}
|
||||
|
}
|
||||
Write
Preview
Loading…
Cancel
Save
Reference in new issue