497 lines
16 KiB
Rust
497 lines
16 KiB
Rust
use std::net::SocketAddr;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
|
|
use sq_grpc_interface::{
|
|
data_plane_service_client::DataPlaneServiceClient,
|
|
data_plane_service_server::DataPlaneServiceServer,
|
|
status_service_client::StatusServiceClient,
|
|
status_service_server::StatusServiceServer,
|
|
AckMode, GetStatusRequest, MessageHeader, PublishMessage, PublishRequest, PublishSettings,
|
|
SubscribeRequest,
|
|
};
|
|
use sq_sim::fs::InMemoryFileSystem;
|
|
use sq_sim::SimClock;
|
|
use sq_storage::engine::StorageEngine;
|
|
use tokio::sync::Mutex;
|
|
use tokio_stream::StreamExt;
|
|
|
|
/// A lightweight test harness that starts a gRPC server on a random port
|
|
/// and returns both the server task and connected clients.
|
|
struct TestServer {
|
|
addr: SocketAddr,
|
|
_shutdown: tokio::sync::oneshot::Sender<()>,
|
|
}
|
|
|
|
impl TestServer {
|
|
async fn start() -> Self {
|
|
let fs = Arc::new(InMemoryFileSystem::new());
|
|
let clock = Arc::new(SimClock::new());
|
|
let config = sq_models::WalConfig {
|
|
max_segment_bytes: 1024 * 1024,
|
|
max_segment_age_secs: 3600,
|
|
data_dir: PathBuf::from("/data"),
|
|
..Default::default()
|
|
};
|
|
|
|
let engine = StorageEngine::new(fs, clock, config).unwrap();
|
|
engine.recover().unwrap();
|
|
|
|
let engine = Arc::new(Mutex::new(engine));
|
|
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let addr = listener.local_addr().unwrap();
|
|
|
|
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
|
|
|
|
// Build the health server state-like object inline for tests.
|
|
let node_id = "test-node".to_string();
|
|
|
|
struct TestHealthServer {
|
|
node_id: String,
|
|
}
|
|
|
|
#[tonic::async_trait]
|
|
impl sq_grpc_interface::status_service_server::StatusService for TestHealthServer {
|
|
async fn status(
|
|
&self,
|
|
_request: tonic::Request<GetStatusRequest>,
|
|
) -> Result<tonic::Response<sq_grpc_interface::GetStatusResponse>, tonic::Status> {
|
|
Ok(tonic::Response::new(sq_grpc_interface::GetStatusResponse {
|
|
node_id: self.node_id.clone(),
|
|
cluster: None,
|
|
}))
|
|
}
|
|
}
|
|
|
|
struct TestDataPlaneServer {
|
|
engine: Arc<Mutex<StorageEngine<InMemoryFileSystem, SimClock>>>,
|
|
}
|
|
|
|
#[tonic::async_trait]
|
|
impl sq_grpc_interface::data_plane_service_server::DataPlaneService
|
|
for TestDataPlaneServer
|
|
{
|
|
async fn publish(
|
|
&self,
|
|
request: tonic::Request<PublishRequest>,
|
|
) -> Result<tonic::Response<sq_grpc_interface::PublishResponse>, tonic::Status> {
|
|
let req = request.into_inner();
|
|
|
|
if req.messages.is_empty() {
|
|
return Err(tonic::Status::invalid_argument(
|
|
"messages must not be empty",
|
|
));
|
|
}
|
|
|
|
let mut results = Vec::new();
|
|
let engine = self.engine.lock().await;
|
|
|
|
for msg in &req.messages {
|
|
if msg.topic.is_empty() {
|
|
return Err(tonic::Status::invalid_argument("topic must not be empty"));
|
|
}
|
|
|
|
let headers: Vec<sq_models::Header> = msg
|
|
.headers
|
|
.iter()
|
|
.map(|h| sq_models::Header {
|
|
key: h.key.clone(),
|
|
value: h.value.clone(),
|
|
})
|
|
.collect();
|
|
|
|
let key = if msg.key.is_empty() {
|
|
None
|
|
} else {
|
|
Some(msg.key.as_slice())
|
|
};
|
|
|
|
let offset = engine
|
|
.append(&msg.topic, 0, key, &msg.value, &headers, 0)
|
|
.map_err(|e| tonic::Status::internal(e.to_string()))?;
|
|
|
|
results.push(sq_grpc_interface::PublishResult {
|
|
topic: msg.topic.clone(),
|
|
partition: 0,
|
|
offset,
|
|
});
|
|
}
|
|
|
|
Ok(tonic::Response::new(sq_grpc_interface::PublishResponse {
|
|
results,
|
|
}))
|
|
}
|
|
|
|
type SubscribeStream = std::pin::Pin<
|
|
Box<
|
|
dyn tokio_stream::Stream<
|
|
Item = Result<sq_grpc_interface::SubscribeResponse, tonic::Status>,
|
|
> + Send
|
|
+ 'static,
|
|
>,
|
|
>;
|
|
|
|
async fn subscribe(
|
|
&self,
|
|
request: tonic::Request<SubscribeRequest>,
|
|
) -> Result<tonic::Response<Self::SubscribeStream>, tonic::Status> {
|
|
let req = request.into_inner();
|
|
let batch_size = if req.max_batch_size == 0 {
|
|
100
|
|
} else {
|
|
req.max_batch_size as usize
|
|
};
|
|
let start_offset = req.start_offset.unwrap_or(0);
|
|
let topic = req.topic.clone();
|
|
let partition = req.partition;
|
|
let engine = self.engine.clone();
|
|
|
|
let stream = async_stream::try_stream! {
|
|
let mut current_offset = start_offset;
|
|
let mut empty_polls = 0u32;
|
|
|
|
loop {
|
|
let messages = {
|
|
let eng = engine.lock().await;
|
|
eng.read(&topic, partition, current_offset, batch_size)
|
|
.map_err(|e| tonic::Status::internal(e.to_string()))?
|
|
};
|
|
|
|
if messages.is_empty() {
|
|
empty_polls += 1;
|
|
// In tests, stop after a few empty polls to avoid hanging.
|
|
if empty_polls > 3 {
|
|
break;
|
|
}
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
|
continue;
|
|
}
|
|
|
|
empty_polls = 0;
|
|
|
|
let consumed: Vec<sq_grpc_interface::ConsumedMessage> = messages
|
|
.iter()
|
|
.map(|m| {
|
|
current_offset = m.offset + 1;
|
|
sq_grpc_interface::ConsumedMessage {
|
|
offset: m.offset,
|
|
topic: m.topic.to_string(),
|
|
partition: m.partition,
|
|
key: m.key.clone().unwrap_or_default(),
|
|
value: m.value.clone(),
|
|
headers: m
|
|
.headers
|
|
.iter()
|
|
.map(|h| MessageHeader {
|
|
key: h.key.clone(),
|
|
value: h.value.clone(),
|
|
})
|
|
.collect(),
|
|
timestamp_ms: m.timestamp_ms,
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
yield sq_grpc_interface::SubscribeResponse { messages: consumed };
|
|
}
|
|
};
|
|
|
|
Ok(tonic::Response::new(Box::pin(stream)))
|
|
}
|
|
|
|
async fn ack(
|
|
&self,
|
|
_request: tonic::Request<sq_grpc_interface::AckRequest>,
|
|
) -> Result<tonic::Response<sq_grpc_interface::AckResponse>, tonic::Status> {
|
|
Ok(tonic::Response::new(sq_grpc_interface::AckResponse {}))
|
|
}
|
|
}
|
|
|
|
let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
|
|
|
|
tokio::spawn(async move {
|
|
tonic::transport::Server::builder()
|
|
.add_service(StatusServiceServer::new(TestHealthServer {
|
|
node_id: node_id.clone(),
|
|
}))
|
|
.add_service(DataPlaneServiceServer::new(TestDataPlaneServer {
|
|
engine,
|
|
}))
|
|
.serve_with_incoming_shutdown(incoming, async {
|
|
let _ = shutdown_rx.await;
|
|
})
|
|
.await
|
|
.unwrap();
|
|
});
|
|
|
|
// Give the server a moment to start.
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
|
|
|
TestServer {
|
|
addr,
|
|
_shutdown: shutdown_tx,
|
|
}
|
|
}
|
|
|
|
fn endpoint(&self) -> String {
|
|
format!("http://{}", self.addr)
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_status_rpc() {
|
|
let server = TestServer::start().await;
|
|
let mut client = StatusServiceClient::connect(server.endpoint()).await.unwrap();
|
|
|
|
let response = client
|
|
.status(tonic::Request::new(GetStatusRequest {}))
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.into_inner().node_id, "test-node");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_publish_single_message() {
|
|
let server = TestServer::start().await;
|
|
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
|
.await
|
|
.unwrap();
|
|
|
|
let response = client
|
|
.publish(tonic::Request::new(PublishRequest {
|
|
messages: vec![PublishMessage {
|
|
topic: "orders".to_string(),
|
|
key: vec![],
|
|
value: b"hello world".to_vec(),
|
|
headers: vec![],
|
|
}],
|
|
settings: None,
|
|
producer_id: "test".to_string(),
|
|
}))
|
|
.await
|
|
.unwrap();
|
|
|
|
let results = response.into_inner().results;
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].topic, "orders");
|
|
assert_eq!(results[0].offset, 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_publish_batch_sequential_offsets() {
|
|
let server = TestServer::start().await;
|
|
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
|
.await
|
|
.unwrap();
|
|
|
|
let messages: Vec<PublishMessage> = (0..100)
|
|
.map(|i| PublishMessage {
|
|
topic: "events".to_string(),
|
|
key: vec![],
|
|
value: format!("msg-{i}").into_bytes(),
|
|
headers: vec![],
|
|
})
|
|
.collect();
|
|
|
|
let response = client
|
|
.publish(tonic::Request::new(PublishRequest {
|
|
messages,
|
|
settings: Some(PublishSettings {
|
|
ack_mode: AckMode::All.into(),
|
|
}),
|
|
producer_id: "test".to_string(),
|
|
}))
|
|
.await
|
|
.unwrap();
|
|
|
|
let results = response.into_inner().results;
|
|
assert_eq!(results.len(), 100);
|
|
for (i, r) in results.iter().enumerate() {
|
|
assert_eq!(r.offset, i as u64);
|
|
assert_eq!(r.topic, "events");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_publish_empty_topic_returns_error() {
|
|
let server = TestServer::start().await;
|
|
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
|
.await
|
|
.unwrap();
|
|
|
|
let err = client
|
|
.publish(tonic::Request::new(PublishRequest {
|
|
messages: vec![PublishMessage {
|
|
topic: "".to_string(),
|
|
key: vec![],
|
|
value: b"data".to_vec(),
|
|
headers: vec![],
|
|
}],
|
|
settings: None,
|
|
producer_id: "test".to_string(),
|
|
}))
|
|
.await
|
|
.unwrap_err();
|
|
|
|
assert_eq!(err.code(), tonic::Code::InvalidArgument);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_publish_empty_messages_returns_error() {
|
|
let server = TestServer::start().await;
|
|
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
|
.await
|
|
.unwrap();
|
|
|
|
let err = client
|
|
.publish(tonic::Request::new(PublishRequest {
|
|
messages: vec![],
|
|
settings: None,
|
|
producer_id: "test".to_string(),
|
|
}))
|
|
.await
|
|
.unwrap_err();
|
|
|
|
assert_eq!(err.code(), tonic::Code::InvalidArgument);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_publish_with_key_and_headers() {
|
|
let server = TestServer::start().await;
|
|
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
|
.await
|
|
.unwrap();
|
|
|
|
let response = client
|
|
.publish(tonic::Request::new(PublishRequest {
|
|
messages: vec![PublishMessage {
|
|
topic: "orders".to_string(),
|
|
key: b"order-123".to_vec(),
|
|
value: b"payload".to_vec(),
|
|
headers: vec![MessageHeader {
|
|
key: "trace-id".to_string(),
|
|
value: b"abc-123".to_vec(),
|
|
}],
|
|
}],
|
|
settings: None,
|
|
producer_id: "test".to_string(),
|
|
}))
|
|
.await
|
|
.unwrap();
|
|
|
|
let results = response.into_inner().results;
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].offset, 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_subscribe_from_beginning() {
|
|
let server = TestServer::start().await;
|
|
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
|
.await
|
|
.unwrap();
|
|
|
|
// Publish 10 messages first.
|
|
let messages: Vec<PublishMessage> = (0..10)
|
|
.map(|i| PublishMessage {
|
|
topic: "events".to_string(),
|
|
key: vec![],
|
|
value: format!("msg-{i}").into_bytes(),
|
|
headers: vec![],
|
|
})
|
|
.collect();
|
|
|
|
client
|
|
.publish(tonic::Request::new(PublishRequest {
|
|
messages,
|
|
settings: None,
|
|
producer_id: "test".to_string(),
|
|
}))
|
|
.await
|
|
.unwrap();
|
|
|
|
// Subscribe from offset 0.
|
|
let response = client
|
|
.subscribe(tonic::Request::new(SubscribeRequest {
|
|
topic: "events".to_string(),
|
|
partition: 0,
|
|
consumer_group: "".to_string(),
|
|
start_offset: Some(0),
|
|
max_batch_size: 100,
|
|
}))
|
|
.await
|
|
.unwrap();
|
|
|
|
let mut stream = response.into_inner();
|
|
let mut all_messages = Vec::new();
|
|
|
|
while let Some(Ok(batch)) = stream.next().await {
|
|
all_messages.extend(batch.messages);
|
|
if all_messages.len() >= 10 {
|
|
break;
|
|
}
|
|
}
|
|
|
|
assert_eq!(all_messages.len(), 10);
|
|
for (i, msg) in all_messages.iter().enumerate() {
|
|
assert_eq!(msg.offset, i as u64);
|
|
assert_eq!(msg.value, format!("msg-{i}").as_bytes());
|
|
assert_eq!(msg.topic, "events");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_subscribe_from_middle() {
|
|
let server = TestServer::start().await;
|
|
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
|
.await
|
|
.unwrap();
|
|
|
|
// Publish 10 messages.
|
|
let messages: Vec<PublishMessage> = (0..10)
|
|
.map(|i| PublishMessage {
|
|
topic: "events".to_string(),
|
|
key: vec![],
|
|
value: format!("msg-{i}").into_bytes(),
|
|
headers: vec![],
|
|
})
|
|
.collect();
|
|
|
|
client
|
|
.publish(tonic::Request::new(PublishRequest {
|
|
messages,
|
|
settings: None,
|
|
producer_id: "test".to_string(),
|
|
}))
|
|
.await
|
|
.unwrap();
|
|
|
|
// Subscribe from offset 5.
|
|
let response = client
|
|
.subscribe(tonic::Request::new(SubscribeRequest {
|
|
topic: "events".to_string(),
|
|
partition: 0,
|
|
consumer_group: "".to_string(),
|
|
start_offset: Some(5),
|
|
max_batch_size: 100,
|
|
}))
|
|
.await
|
|
.unwrap();
|
|
|
|
let mut stream = response.into_inner();
|
|
let mut all_messages = Vec::new();
|
|
|
|
while let Some(Ok(batch)) = stream.next().await {
|
|
all_messages.extend(batch.messages);
|
|
if all_messages.len() >= 5 {
|
|
break;
|
|
}
|
|
}
|
|
|
|
assert_eq!(all_messages.len(), 5);
|
|
assert_eq!(all_messages[0].offset, 5);
|
|
assert_eq!(all_messages[4].offset, 9);
|
|
}
|