use bytes::{Bytes, BytesMut};
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TrySendError;
use tokio::task::JoinHandle;
use tonic::Streaming;
use tracing::{debug, trace, warn};
use crate::client::WorkerClient;
use crate::config::GoosefsConfig;
use crate::error::{Error, Result};
use crate::metrics::name;
use crate::proto::grpc::block::{ReadRequest, ReadResponse};
use crate::proto::proto::dataserver::OpenUfsBlockOptions;
#[derive(Debug, Clone, Copy)]
pub struct ReadTuning {
pub prefetch_window: i32,
pub buffer_messages: usize,
pub ack_interval_bytes: i64,
pub ack_interval_chunks: u32,
}
impl ReadTuning {
pub fn from_config(config: &GoosefsConfig) -> Self {
Self {
prefetch_window: config.prefetch_window,
buffer_messages: config.read_buffer_messages.max(1),
ack_interval_bytes: config.ack_interval_bytes.max(0),
ack_interval_chunks: config.ack_interval_chunks.max(1),
}
}
}
enum StreamItem {
Data(ReadResponse),
End,
Error(Error),
}
enum ChunkSource {
Direct(Streaming<ReadResponse>),
Buffered {
rx: mpsc::Receiver<StreamItem>,
task: JoinHandle<()>,
},
}
pub struct GrpcBlockReader {
block_id: i64,
offset: i64,
length: i64,
bytes_received: i64,
request_tx: mpsc::Sender<ReadRequest>,
source: ChunkSource,
bytes_since_last_ack: i64,
chunks_since_last_ack: u32,
ack_interval_bytes: i64,
ack_interval_chunks: u32,
}
#[derive(Debug, PartialEq, Eq)]
enum ChunkAction {
Eof,
KeepReading,
Deliver(Bytes),
}
fn classify_response(resp: Option<ReadResponse>) -> ChunkAction {
match resp {
None => ChunkAction::Eof,
Some(r) => {
let data = r.chunk.and_then(|c| c.data).unwrap_or_default();
if data.is_empty() {
ChunkAction::KeepReading
} else {
ChunkAction::Deliver(Bytes::from(data))
}
}
}
}
fn check_positioned_read_complete(block_id: i64, bytes_received: i64, length: i64) -> Result<()> {
if bytes_received < length {
return Err(Error::Internal {
message: format!(
"short positioned read on block {}: received {} of {} bytes \
(server half-closed early)",
block_id, bytes_received, length
),
source: None,
});
}
Ok(())
}
fn should_send_ack(
bytes_since_last_ack: i64,
chunks_since_last_ack: u32,
ack_interval_bytes: i64,
ack_interval_chunks: u32,
bytes_received: i64,
length: i64,
) -> bool {
bytes_since_last_ack >= ack_interval_bytes
|| chunks_since_last_ack >= ack_interval_chunks
|| bytes_received >= length
}
impl GrpcBlockReader {
pub async fn open(
worker: &WorkerClient,
block_id: i64,
offset: i64,
length: i64,
chunk_size: i64,
open_ufs_block_options: Option<OpenUfsBlockOptions>,
) -> Result<Self> {
let (request_tx, response_rx) = worker
.read_block(
block_id,
offset,
length,
chunk_size,
None,
open_ufs_block_options,
)
.await?;
crate::metrics::gauge(name::CLIENT_BLOCKS_READ_IN_PROGRESS)
.set(crate::metrics::gauge(name::CLIENT_BLOCKS_READ_IN_PROGRESS).get() + 1);
debug!(
block_id = block_id,
offset = offset,
length = length,
"opened GrpcBlockReader"
);
Ok(Self {
block_id,
offset,
length,
bytes_received: 0,
request_tx,
source: ChunkSource::Direct(response_rx),
bytes_since_last_ack: 0,
chunks_since_last_ack: 0,
ack_interval_bytes: 0,
ack_interval_chunks: 1,
})
}
pub async fn open_sequential(
worker: &WorkerClient,
block_id: i64,
offset: i64,
length: i64,
chunk_size: i64,
open_ufs_block_options: Option<OpenUfsBlockOptions>,
tuning: ReadTuning,
) -> Result<Self> {
let (request_tx, response_rx) = worker
.read_block(
block_id,
offset,
length,
chunk_size,
Some(tuning.prefetch_window),
open_ufs_block_options,
)
.await?;
crate::metrics::gauge(name::CLIENT_BLOCKS_READ_IN_PROGRESS)
.set(crate::metrics::gauge(name::CLIENT_BLOCKS_READ_IN_PROGRESS).get() + 1);
let (chunk_tx, chunk_rx) = mpsc::channel::<StreamItem>(tuning.buffer_messages);
let mut stream = response_rx;
let task = tokio::spawn(async move {
loop {
match stream.message().await {
Ok(Some(resp)) => {
if chunk_tx.send(StreamItem::Data(resp)).await.is_err() {
break;
}
}
Ok(None) => {
let _ = chunk_tx.send(StreamItem::End).await;
break;
}
Err(status) => {
let _ = chunk_tx.send(StreamItem::Error(Error::from(status))).await;
break;
}
}
}
});
debug!(
block_id = block_id,
offset = offset,
length = length,
prefetch_window = tuning.prefetch_window,
buffer_messages = tuning.buffer_messages,
"opened GrpcBlockReader (sequential, buffered)"
);
Ok(Self {
block_id,
offset,
length,
bytes_received: 0,
request_tx,
source: ChunkSource::Buffered { rx: chunk_rx, task },
bytes_since_last_ack: 0,
chunks_since_last_ack: 0,
ack_interval_bytes: tuning.ack_interval_bytes,
ack_interval_chunks: tuning.ack_interval_chunks,
})
}
pub async fn read_chunk(&mut self) -> Result<Option<Bytes>> {
if self.bytes_received >= self.length {
return Ok(None);
}
loop {
let resp = match &mut self.source {
ChunkSource::Direct(stream) => match stream.message().await? {
None => {
debug!(
block_id = self.block_id,
bytes_received = self.bytes_received,
"stream ended before all expected data received"
);
return Ok(None);
}
Some(r) => r,
},
ChunkSource::Buffered { rx, .. } => match rx.recv().await {
Some(StreamItem::Data(r)) => r,
Some(StreamItem::End) => {
debug!(
block_id = self.block_id,
bytes_received = self.bytes_received,
"buffered stream reached clean EOF"
);
return Ok(None);
}
Some(StreamItem::Error(e)) => return Err(e),
None => {
return Err(Error::Internal {
message: format!(
"read stream drain task ended unexpectedly on block {} \
({} of {} bytes received)",
self.block_id, self.bytes_received, self.length
),
source: None,
});
}
},
};
match classify_response(Some(resp)) {
ChunkAction::Eof => unreachable!("Some(_) cannot classify as Eof"),
ChunkAction::KeepReading => {
trace!(
block_id = self.block_id,
bytes_received = self.bytes_received,
expected = self.length,
"received empty data frame, awaiting next chunk"
);
continue;
}
ChunkAction::Deliver(data) => {
let len = data.len() as i64;
self.bytes_received += len;
trace!(
block_id = self.block_id,
chunk_len = data.len(),
total_received = self.bytes_received,
"received chunk"
);
crate::metrics::counter(name::CLIENT_BYTES_READ_LOCAL).inc(len);
self.maybe_send_ack(len);
return Ok(Some(data));
}
}
}
}
fn maybe_send_ack(&mut self, delta: i64) {
self.bytes_since_last_ack += delta;
self.chunks_since_last_ack += 1;
let need_ack = should_send_ack(
self.bytes_since_last_ack,
self.chunks_since_last_ack,
self.ack_interval_bytes,
self.ack_interval_chunks,
self.bytes_received,
self.length,
);
if !need_ack {
return;
}
let ack = ReadRequest {
offset_received: Some(self.offset + self.bytes_received),
..Default::default()
};
match self.request_tx.try_send(ack) {
Ok(()) => {
self.bytes_since_last_ack = 0;
self.chunks_since_last_ack = 0;
}
Err(TrySendError::Full(_)) => {
}
Err(TrySendError::Closed(_)) => {
warn!(
block_id = self.block_id,
"ACK channel closed (read may be complete)"
);
}
}
}
pub async fn read_all(&mut self) -> Result<Bytes> {
let mut buf = BytesMut::with_capacity(self.length as usize);
while let Some(chunk) = self.read_chunk().await? {
buf.extend_from_slice(&chunk);
}
crate::metrics::counter(name::CLIENT_BLOCKS_READ_TOTAL).inc(1);
crate::metrics::gauge(name::CLIENT_BLOCKS_READ_IN_PROGRESS)
.set((crate::metrics::gauge(name::CLIENT_BLOCKS_READ_IN_PROGRESS).get() - 1).max(0));
check_positioned_read_complete(self.block_id, self.bytes_received, self.length)?;
Ok(buf.freeze())
}
pub fn block_id(&self) -> i64 {
self.block_id
}
pub fn bytes_received(&self) -> i64 {
self.bytes_received
}
pub fn is_complete(&self) -> bool {
self.bytes_received >= self.length
}
pub async fn positioned_read(
worker: &WorkerClient,
block_id: i64,
offset: i64,
length: i64,
chunk_size: i64,
open_ufs_block_options: Option<OpenUfsBlockOptions>,
) -> Result<Bytes> {
let (request_tx, response_rx) = worker
.read_block_positioned(block_id, offset, length, chunk_size, open_ufs_block_options)
.await?;
debug!(
block_id = block_id,
offset = offset,
length = length,
"positioned_read: opened position_short stream"
);
let mut reader = Self {
block_id,
offset,
length,
bytes_received: 0,
request_tx,
source: ChunkSource::Direct(response_rx),
bytes_since_last_ack: 0,
chunks_since_last_ack: 0,
ack_interval_bytes: 0,
ack_interval_chunks: 1,
};
reader.read_all().await
}
}
impl Drop for GrpcBlockReader {
fn drop(&mut self) {
if let ChunkSource::Buffered { task, .. } = &self.source {
task.abort();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::proto::grpc::block::Chunk;
#[test]
fn metrics_counter_accessible() {
let _counter = crate::metrics::counter(name::CLIENT_BYTES_READ_LOCAL);
}
#[test]
fn classify_response_none_is_eof() {
assert_eq!(classify_response(None), ChunkAction::Eof);
}
#[test]
fn classify_response_no_chunk_is_keep_reading() {
let resp = ReadResponse {
chunk: None,
..Default::default()
};
assert_eq!(classify_response(Some(resp)), ChunkAction::KeepReading);
}
#[test]
fn classify_response_empty_chunk_is_keep_reading() {
let resp = ReadResponse {
chunk: Some(Chunk {
data: Some(Vec::new()),
..Default::default()
}),
..Default::default()
};
assert_eq!(classify_response(Some(resp)), ChunkAction::KeepReading);
}
#[test]
fn classify_response_chunk_with_none_data_is_keep_reading() {
let resp = ReadResponse {
chunk: Some(Chunk {
data: None,
..Default::default()
}),
..Default::default()
};
assert_eq!(classify_response(Some(resp)), ChunkAction::KeepReading);
}
#[test]
fn classify_response_data_is_delivered() {
let payload = b"hello world".to_vec();
let resp = ReadResponse {
chunk: Some(Chunk {
data: Some(payload.clone()),
..Default::default()
}),
..Default::default()
};
match classify_response(Some(resp)) {
ChunkAction::Deliver(b) => assert_eq!(b.as_ref(), payload.as_slice()),
other => panic!("expected Deliver, got {:?}", other),
}
}
#[test]
fn check_positioned_read_complete_short_read_errors() {
let err = check_positioned_read_complete(
16777216, 1024, 4096,
)
.unwrap_err();
let msg = format!("{}", err);
assert!(
msg.contains("short positioned read on block 16777216"),
"expected short-read message, got: {}",
msg
);
assert!(
msg.contains("received 1024 of 4096"),
"expected received/length pair in message, got: {}",
msg
);
}
#[test]
fn check_positioned_read_complete_full_read_ok() {
assert!(check_positioned_read_complete(1, 4096, 4096).is_ok());
}
#[test]
fn check_positioned_read_complete_over_read_ok() {
assert!(check_positioned_read_complete(1, 5000, 4096).is_ok());
}
#[test]
fn should_send_ack_direct_mode_acks_every_chunk() {
assert!(should_send_ack(64, 1, 0, 1, 64, 1_000_000));
}
#[test]
fn should_send_ack_coalesces_until_threshold() {
let interval_bytes = 4 * 1024 * 1024;
let interval_chunks = 4;
assert!(!should_send_ack(
2 * 1024 * 1024,
2,
interval_bytes,
interval_chunks,
2 * 1024 * 1024,
64 * 1024 * 1024
));
assert!(should_send_ack(
interval_bytes,
2,
interval_bytes,
interval_chunks,
interval_bytes,
64 * 1024 * 1024
));
assert!(should_send_ack(
1024,
interval_chunks,
interval_bytes,
interval_chunks,
4096,
64 * 1024 * 1024
));
}
#[test]
fn should_send_ack_forces_final_ack_at_completion() {
assert!(should_send_ack(
64,
1,
4 * 1024 * 1024,
4,
1_000_000,
1_000_000
));
}
#[test]
fn read_tuning_from_config_defaults_and_clamps() {
let mut cfg = crate::config::GoosefsConfig::new("127.0.0.1:9200");
let t = ReadTuning::from_config(&cfg);
assert_eq!(t.prefetch_window, 8);
assert_eq!(t.buffer_messages, 16);
assert_eq!(t.ack_interval_bytes, 0); assert_eq!(t.ack_interval_chunks, 1);
cfg.read_buffer_messages = 0;
cfg.ack_interval_chunks = 0;
cfg.ack_interval_bytes = -1;
let t = ReadTuning::from_config(&cfg);
assert_eq!(t.buffer_messages, 1);
assert_eq!(t.ack_interval_chunks, 1);
assert_eq!(t.ack_interval_bytes, 0);
}
}