335 lines
11 KiB
Rust
335 lines
11 KiB
Rust
//! gRPC service implementation for storage node
|
|
|
|
use crate::proto::{
|
|
node_service_server::NodeService, BatchGetChunksRequest, BatchPutChunksResponse,
|
|
ChunkExistsRequest, ChunkExistsResponse, ChunkSizeRequest, ChunkSizeResponse,
|
|
DeleteChunkRequest, GetChunkRequest, GetChunkResponse, GetStatusRequest, GetStatusResponse,
|
|
PingRequest, PingResponse, PutChunkRequest, PutChunkResponse,
|
|
};
|
|
use crate::storage::LocalChunkStore;
|
|
use crate::NodeConfig;
|
|
use std::sync::Arc;
|
|
use std::time::Instant;
|
|
use tokio::task::JoinSet;
|
|
use tokio_stream::wrappers::ReceiverStream;
|
|
use tonic::{Request, Response, Status, Streaming};
|
|
use tracing::{debug, error};
|
|
|
|
const BATCH_IO_PARALLELISM: usize = 32;
|
|
|
|
/// Implementation of the NodeService gRPC service
|
|
pub struct NodeServiceImpl {
|
|
/// Local chunk storage
|
|
store: Arc<LocalChunkStore>,
|
|
/// Node configuration
|
|
config: Arc<NodeConfig>,
|
|
/// Server start time
|
|
start_time: Instant,
|
|
}
|
|
|
|
impl NodeServiceImpl {
|
|
/// Create a new node service
|
|
pub fn new(store: Arc<LocalChunkStore>, config: Arc<NodeConfig>) -> Self {
|
|
Self {
|
|
store,
|
|
config,
|
|
start_time: Instant::now(),
|
|
}
|
|
}
|
|
|
|
fn chunk_read_status(chunk_id: &str, error: crate::storage::StorageError) -> Status {
|
|
match error {
|
|
crate::storage::StorageError::NotFound(_) => Status::not_found(format!(
|
|
"Chunk not found: {}",
|
|
chunk_id
|
|
)),
|
|
other => Status::internal(other.to_string()),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tonic::async_trait]
|
|
impl NodeService for NodeServiceImpl {
|
|
async fn put_chunk(
|
|
&self,
|
|
request: Request<PutChunkRequest>,
|
|
) -> Result<Response<PutChunkResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
debug!(
|
|
chunk_id = %req.chunk_id,
|
|
shard_index = req.shard_index,
|
|
is_parity = req.is_parity,
|
|
size = req.data.len(),
|
|
"PutChunk request"
|
|
);
|
|
|
|
let size = self
|
|
.store
|
|
.put(&req.chunk_id, &req.data)
|
|
.await
|
|
.map_err(|e| {
|
|
error!(error = ?e, "Failed to put chunk");
|
|
Status::internal(e.to_string())
|
|
})?;
|
|
|
|
metrics::counter!("node_chunks_stored").increment(1);
|
|
metrics::counter!("node_bytes_stored").increment(size);
|
|
|
|
Ok(Response::new(PutChunkResponse { size }))
|
|
}
|
|
|
|
async fn get_chunk(
|
|
&self,
|
|
request: Request<GetChunkRequest>,
|
|
) -> Result<Response<GetChunkResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
debug!(
|
|
chunk_id = %req.chunk_id,
|
|
shard_index = req.shard_index,
|
|
is_parity = req.is_parity,
|
|
"GetChunk request"
|
|
);
|
|
|
|
let data = self.store.get(&req.chunk_id).await.map_err(|e| {
|
|
match &e {
|
|
crate::storage::StorageError::NotFound(_) => {
|
|
debug!(chunk_id = %req.chunk_id, "Chunk not found");
|
|
Status::not_found(e.to_string())
|
|
}
|
|
_ => {
|
|
error!(error = ?e, "Failed to get chunk");
|
|
Status::internal(e.to_string())
|
|
}
|
|
}
|
|
})?;
|
|
|
|
metrics::counter!("node_chunks_retrieved").increment(1);
|
|
metrics::counter!("node_bytes_retrieved").increment(data.len() as u64);
|
|
|
|
Ok(Response::new(GetChunkResponse {
|
|
data: data.into(),
|
|
size: 0, // Size is implicit from data.len()
|
|
}))
|
|
}
|
|
|
|
async fn delete_chunk(
|
|
&self,
|
|
request: Request<DeleteChunkRequest>,
|
|
) -> Result<Response<()>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
debug!(chunk_id = %req.chunk_id, "DeleteChunk request");
|
|
|
|
self.store.delete(&req.chunk_id).await.map_err(|e| {
|
|
error!(error = ?e, "Failed to delete chunk");
|
|
Status::internal(e.to_string())
|
|
})?;
|
|
|
|
metrics::counter!("node_chunks_deleted").increment(1);
|
|
|
|
Ok(Response::new(()))
|
|
}
|
|
|
|
async fn chunk_exists(
|
|
&self,
|
|
request: Request<ChunkExistsRequest>,
|
|
) -> Result<Response<ChunkExistsResponse>, Status> {
|
|
let req = request.into_inner();
|
|
let exists = self.store.exists(&req.chunk_id);
|
|
|
|
Ok(Response::new(ChunkExistsResponse { exists }))
|
|
}
|
|
|
|
async fn chunk_size(
|
|
&self,
|
|
request: Request<ChunkSizeRequest>,
|
|
) -> Result<Response<ChunkSizeResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
match self.store.size(&req.chunk_id) {
|
|
Some(size) => Ok(Response::new(ChunkSizeResponse { size, exists: true })),
|
|
None => Ok(Response::new(ChunkSizeResponse {
|
|
size: 0,
|
|
exists: false,
|
|
})),
|
|
}
|
|
}
|
|
|
|
async fn ping(&self, _request: Request<PingRequest>) -> Result<Response<PingResponse>, Status> {
|
|
let start = Instant::now();
|
|
// Minimal processing - just measure latency
|
|
let latency_us = start.elapsed().as_micros() as u64;
|
|
|
|
Ok(Response::new(PingResponse { latency_us }))
|
|
}
|
|
|
|
async fn get_status(
|
|
&self,
|
|
_request: Request<GetStatusRequest>,
|
|
) -> Result<Response<GetStatusResponse>, Status> {
|
|
let uptime_seconds = self.start_time.elapsed().as_secs();
|
|
|
|
Ok(Response::new(GetStatusResponse {
|
|
node_id: self.config.node_id.clone(),
|
|
endpoint: self.config.grpc_addr.to_string(),
|
|
zone: self.config.zone.clone(),
|
|
region: self.config.region.clone(),
|
|
capacity_bytes: self.store.max_capacity(),
|
|
used_bytes: self.store.total_bytes(),
|
|
chunk_count: self.store.chunk_count(),
|
|
healthy: true,
|
|
uptime_seconds,
|
|
}))
|
|
}
|
|
|
|
async fn batch_put_chunks(
|
|
&self,
|
|
request: Request<Streaming<PutChunkRequest>>,
|
|
) -> Result<Response<BatchPutChunksResponse>, Status> {
|
|
let mut stream = request.into_inner();
|
|
let mut success_count = 0u32;
|
|
let mut failure_count = 0u32;
|
|
let mut errors = Vec::new();
|
|
let mut in_flight = JoinSet::new();
|
|
|
|
while let Some(req) = stream.message().await? {
|
|
while in_flight.len() >= BATCH_IO_PARALLELISM {
|
|
record_batch_put_result(
|
|
in_flight.join_next().await,
|
|
&mut success_count,
|
|
&mut failure_count,
|
|
&mut errors,
|
|
);
|
|
}
|
|
|
|
let store = self.store.clone();
|
|
in_flight.spawn(async move {
|
|
let chunk_id = req.chunk_id;
|
|
let result = store.put(&chunk_id, &req.data).await;
|
|
(chunk_id, result)
|
|
});
|
|
}
|
|
|
|
while !in_flight.is_empty() {
|
|
record_batch_put_result(
|
|
in_flight.join_next().await,
|
|
&mut success_count,
|
|
&mut failure_count,
|
|
&mut errors,
|
|
);
|
|
}
|
|
|
|
Ok(Response::new(BatchPutChunksResponse {
|
|
success_count,
|
|
failure_count,
|
|
errors,
|
|
}))
|
|
}
|
|
|
|
type BatchGetChunksStream = ReceiverStream<Result<GetChunkResponse, Status>>;
|
|
|
|
async fn batch_get_chunks(
|
|
&self,
|
|
request: Request<BatchGetChunksRequest>,
|
|
) -> Result<Response<Self::BatchGetChunksStream>, Status> {
|
|
let req = request.into_inner();
|
|
let (tx, rx) = tokio::sync::mpsc::channel(32);
|
|
let store = self.store.clone();
|
|
|
|
tokio::spawn(async move {
|
|
let chunks = req.chunks;
|
|
let chunk_count = chunks.len();
|
|
let mut results: Vec<Option<Result<GetChunkResponse, Status>>> =
|
|
(0..chunk_count).map(|_| None).collect();
|
|
let mut readers = JoinSet::new();
|
|
let mut next_index = 0usize;
|
|
|
|
let spawn_reader =
|
|
|readers: &mut JoinSet<(usize, Result<GetChunkResponse, Status>)>,
|
|
index: usize,
|
|
chunk_req: crate::proto::GetChunkRequest| {
|
|
let store = store.clone();
|
|
readers.spawn(async move {
|
|
let chunk_id = chunk_req.chunk_id;
|
|
let result = store
|
|
.get(&chunk_id)
|
|
.await
|
|
.map(|data| GetChunkResponse {
|
|
size: data.len() as u64,
|
|
data: data.into(),
|
|
})
|
|
.map_err(|error| Self::chunk_read_status(&chunk_id, error));
|
|
(index, result)
|
|
});
|
|
};
|
|
|
|
while next_index < chunk_count && readers.len() < BATCH_IO_PARALLELISM {
|
|
spawn_reader(&mut readers, next_index, chunks[next_index].clone());
|
|
next_index += 1;
|
|
}
|
|
|
|
while let Some(reader) = readers.join_next().await {
|
|
match reader {
|
|
Ok((index, result)) => {
|
|
results[index] = Some(result);
|
|
}
|
|
Err(join_error) => {
|
|
let status = Status::internal(format!(
|
|
"batch get task failed: {}",
|
|
join_error
|
|
));
|
|
let index = results.iter().position(|entry| entry.is_none()).unwrap_or(0);
|
|
results[index] = Some(Err(status));
|
|
}
|
|
}
|
|
|
|
if next_index < chunk_count {
|
|
spawn_reader(&mut readers, next_index, chunks[next_index].clone());
|
|
next_index += 1;
|
|
}
|
|
}
|
|
|
|
for result in results {
|
|
let Some(result) = result else {
|
|
let _ = tx.send(Err(Status::internal("batch get result missing"))).await;
|
|
break;
|
|
};
|
|
|
|
if tx.send(result).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(Response::new(ReceiverStream::new(rx)))
|
|
}
|
|
}
|
|
|
|
fn record_batch_put_result(
|
|
joined: Option<Result<(String, Result<u64, crate::storage::StorageError>), tokio::task::JoinError>>,
|
|
success_count: &mut u32,
|
|
failure_count: &mut u32,
|
|
errors: &mut Vec<String>,
|
|
) {
|
|
let Some(result) = joined else {
|
|
return;
|
|
};
|
|
|
|
match result {
|
|
Ok((_chunk_id, Ok(size))) => {
|
|
*success_count += 1;
|
|
metrics::counter!("node_chunks_stored").increment(1);
|
|
metrics::counter!("node_bytes_stored").increment(size);
|
|
}
|
|
Ok((chunk_id, Err(error))) => {
|
|
*failure_count += 1;
|
|
errors.push(format!("{}: {}", chunk_id, error));
|
|
}
|
|
Err(join_error) => {
|
|
*failure_count += 1;
|
|
errors.push(format!("join error: {}", join_error));
|
|
}
|
|
}
|
|
}
|