Skip to content

GitLab

Unverified Commit 008683d6 authored by Ryan Olson's avatar Ryan Olson Committed by GitHub
Browse files

feat: adding kvbm-engine (#6773) 


Signed-off-by: default avatarRyan Olson <rolson@nvidia.com>
parent cf79c4fc
Hide whitespace changes
Inline Side-by-side
Showing with 0 additions and 402 deletions
lib/velo-transports/tests/uds_shutdown.rs deleted 100644 → 0
View file @ cf79c4fc
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Integration tests for UDS graceful shutdown
//!
//! These tests verify the 3-phase shutdown behavior:
//! 1. Gate: new Message frames are rejected with ShuttingDown
//! 2. Drain: in-flight work completes, responses/events still flow
//! 3. Teardown: listener and writer tasks exit
#![cfg(unix)]
mod common;
use bytes::Bytes;
use std::time::Duration;
use tokio::time::{sleep, timeout};
use velo_transports::tcp::TcpFrameCodec;
use velo_transports::uds::UdsTransport;
use velo_transports::{MessageType, Transport};
use common::TestTransportHandle;
/// Get the socket path from a UDS transport by parsing its WorkerAddress.
fn get_socket_path(handle: &TestTransportHandle<UdsTransport>) -> std::path::PathBuf {
let addr = handle.transport.address();
let key = handle.transport.key();
let endpoint = addr.get_entry(&key).unwrap().unwrap();
let s = std::str::from_utf8(&endpoint).unwrap();
let s = s.strip_prefix("uds://").unwrap_or(s);
std::path::PathBuf::from(s)
}
/// Helper: connect a raw UDS client to the transport's socket and send a frame.
async fn connect_and_send_frame(
socket_path: &std::path::Path,
msg_type: MessageType,
header: &[u8],
payload: &[u8],
) -> tokio::net::UnixStream {
let mut stream = tokio::net::UnixStream::connect(socket_path).await.unwrap();
TcpFrameCodec::encode_frame(&mut stream, msg_type, header, payload)
.await
.unwrap();
stream
}
/// Helper: read one frame from a raw UDS stream.
async fn read_one_frame(stream: &mut tokio::net::UnixStream) -> (MessageType, Bytes, Bytes) {
use futures::StreamExt;
use tokio_util::codec::Framed;
let mut framed = Framed::new(stream, TcpFrameCodec::new());
framed.next().await.unwrap().unwrap()
}
// --- Test: Drain rejects Message frames ---
#[tokio::test]
async fn test_uds_drain_rejects_messages() {
let handle = TestTransportHandle::new_uds().await.unwrap();
let socket_path = get_socket_path(&handle);
// Begin drain
handle.streams.shutdown_state.begin_drain();
// Give listener time to be ready
sleep(Duration::from_millis(50)).await;
// Connect and send a Message frame
let mut stream = connect_and_send_frame(
&socket_path,
MessageType::Message,
b"req-header",
b"req-payload",
)
.await;
// Should get ShuttingDown back
let (msg_type, header, payload) = read_one_frame(&mut stream).await;
assert_eq!(msg_type, MessageType::ShuttingDown);
assert_eq!(&header[..], b"req-header"); // Original header echoed back
assert_eq!(payload.len(), 0); // Empty payload
handle.streams.shutdown_state.teardown_token().cancel();
}
// --- Test: Drain accepts Response frames ---
#[tokio::test]
async fn test_uds_drain_accepts_responses() {
let handle = TestTransportHandle::new_uds().await.unwrap();
let socket_path = get_socket_path(&handle);
// Begin drain
handle.streams.shutdown_state.begin_drain();
sleep(Duration::from_millis(50)).await;
// Connect and send a Response frame
connect_and_send_frame(
&socket_path,
MessageType::Response,
b"resp-header",
b"resp-payload",
)
.await;
// Should arrive on the response stream
let (header, payload) = timeout(
Duration::from_secs(2),
handle.streams.response_stream.recv_async(),
)
.await
.expect("timeout")
.expect("recv");
assert_eq!(&header[..], b"resp-header");
assert_eq!(&payload[..], b"resp-payload");
handle.streams.shutdown_state.teardown_token().cancel();
}
// --- Test: Drain accepts Event frames ---
#[tokio::test]
async fn test_uds_drain_accepts_events() {
let handle = TestTransportHandle::new_uds().await.unwrap();
let socket_path = get_socket_path(&handle);
handle.streams.shutdown_state.begin_drain();
sleep(Duration::from_millis(50)).await;
connect_and_send_frame(
&socket_path,
MessageType::Event,
b"evt-header",
b"evt-payload",
)
.await;
let (header, payload) = timeout(
Duration::from_secs(2),
handle.streams.event_stream.recv_async(),
)
.await
.expect("timeout")
.expect("recv");
assert_eq!(&header[..], b"evt-header");
assert_eq!(&payload[..], b"evt-payload");
handle.streams.shutdown_state.teardown_token().cancel();
}
// --- Test: New connection during drain still accepts responses ---
#[tokio::test]
async fn test_uds_new_connection_during_drain() {
let handle = TestTransportHandle::new_uds().await.unwrap();
let socket_path = get_socket_path(&handle);
// Begin drain BEFORE connecting
handle.streams.shutdown_state.begin_drain();
sleep(Duration::from_millis(50)).await;
// Establish a NEW connection after drain starts
connect_and_send_frame(
&socket_path,
MessageType::Response,
b"new-resp",
b"new-payload",
)
.await;
// Should arrive on the response stream
let (header, payload) = timeout(
Duration::from_secs(2),
handle.streams.response_stream.recv_async(),
)
.await
.expect("timeout")
.expect("recv");
assert_eq!(&header[..], b"new-resp");
assert_eq!(&payload[..], b"new-payload");
handle.streams.shutdown_state.teardown_token().cancel();
}
// --- Test: Full graceful shutdown lifecycle ---
#[tokio::test]
async fn test_uds_graceful_shutdown_lifecycle() {
let handle = TestTransportHandle::new_uds().await.unwrap();
let socket_path = get_socket_path(&handle);
// Verify normal operation: send a message, receive it
connect_and_send_frame(
&socket_path,
MessageType::Message,
b"normal-msg",
b"normal-pay",
)
.await;
let (header, _payload) = timeout(
Duration::from_secs(2),
handle.streams.message_stream.recv_async(),
)
.await
.expect("timeout")
.expect("recv");
assert_eq!(&header[..], b"normal-msg");
// Acquire an InFlightGuard (simulate in-progress request)
let guard = handle.streams.shutdown_state.acquire();
assert_eq!(handle.streams.shutdown_state.in_flight_count(), 1);
// Begin drain (Phase 1)
handle.streams.shutdown_state.begin_drain();
sleep(Duration::from_millis(50)).await;
// Verify new messages are rejected
let mut stream =
connect_and_send_frame(&socket_path, MessageType::Message, b"reject-me", b"").await;
let (msg_type, _, _) = read_one_frame(&mut stream).await;
assert_eq!(msg_type, MessageType::ShuttingDown);
// Verify responses still flow
connect_and_send_frame(&socket_path, MessageType::Response, b"still-ok", b"data").await;
let (header, _) = timeout(
Duration::from_secs(2),
handle.streams.response_stream.recv_async(),
)
.await
.expect("timeout")
.expect("recv");
assert_eq!(&header[..], b"still-ok");
// Spawn graceful_shutdown in background (will block on drain since guard is held)
let shutdown_state = handle.streams.shutdown_state.clone();
let shutdown_handle = tokio::spawn(async move {
// Phase 2: wait for drain
shutdown_state.wait_for_drain().await;
// Phase 3: teardown
shutdown_state.teardown_token().cancel();
});
// Verify shutdown hasn't completed yet (guard still held)
sleep(Duration::from_millis(100)).await;
assert!(!shutdown_handle.is_finished());
// Drop guard -> drain completes -> teardown fires
drop(guard);
timeout(Duration::from_secs(2), shutdown_handle)
.await
.expect("shutdown should complete")
.unwrap();
assert!(
handle
.streams
.shutdown_state
.teardown_token()
.is_cancelled()
);
}
// --- Test: Shutdown timeout forces teardown ---
#[tokio::test]
async fn test_uds_shutdown_timeout_forces_teardown() {
let handle = TestTransportHandle::new_uds().await.unwrap();
// Acquire guard and hold it
let _guard = handle.streams.shutdown_state.acquire();
let shutdown_state = handle.streams.shutdown_state.clone();
let shutdown_handle = tokio::spawn(async move {
shutdown_state.begin_drain();
// Phase 2: wait with short timeout
let _ =
tokio::time::timeout(Duration::from_millis(100), shutdown_state.wait_for_drain()).await;
// Phase 3: teardown (forced, guard still held)
shutdown_state.teardown_token().cancel();
});
timeout(Duration::from_secs(2), shutdown_handle)
.await
.expect("shutdown should complete via timeout")
.unwrap();
// Teardown should have fired even though guard is held
assert!(
handle
.streams
.shutdown_state
.teardown_token()
.is_cancelled()
);
// Guard is still held (not a problem — teardown was forced)
assert_eq!(handle.streams.shutdown_state.in_flight_count(), 1);
}
// --- Test: Outbound sends during drain ---
#[tokio::test]
async fn test_uds_outbound_sends_during_drain() {
// Create two transports and register them as peers
let handle_a = TestTransportHandle::new_uds().await.unwrap();
let handle_b = TestTransportHandle::new_uds().await.unwrap();
handle_a.register_peer(&handle_b).unwrap();
handle_b.register_peer(&handle_a).unwrap();
// Begin drain on transport A
handle_a.streams.shutdown_state.begin_drain();
sleep(Duration::from_millis(50)).await;
// Send a Response from A to B (outbound sends should work during drain)
handle_a.send(
handle_b.instance_id,
b"response-hdr".to_vec(),
b"response-pay".to_vec(),
MessageType::Response,
);
// B should receive the response
let (header, payload) = timeout(
Duration::from_secs(2),
handle_b.streams.response_stream.recv_async(),
)
.await
.expect("timeout")
.expect("recv");
assert_eq!(&header[..], b"response-hdr");
assert_eq!(&payload[..], b"response-pay");
handle_a.streams.shutdown_state.teardown_token().cancel();
handle_b.streams.shutdown_state.teardown_token().cancel();
}
// --- Test: Connection writer exits on teardown ---
#[tokio::test]
async fn test_uds_connection_writer_exits_on_teardown() {
let handle_a = TestTransportHandle::new_uds().await.unwrap();
let handle_b = TestTransportHandle::new_uds().await.unwrap();
handle_a.register_peer(&handle_b).unwrap();
// Send a message to establish the connection writer task
handle_a.send(
handle_b.instance_id,
b"setup".to_vec(),
b"data".to_vec(),
MessageType::Message,
);
// Wait for it to arrive
timeout(
Duration::from_secs(2),
handle_b.streams.message_stream.recv_async(),
)
.await
.expect("timeout")
.expect("recv");
// Shutdown transport A
handle_a.transport.shutdown();
// Give writer tasks time to exit
sleep(Duration::from_millis(200)).await;
// Sending after shutdown: the cancel_token is already cancelled, so any new
// writer task returns immediately without connecting. The error handler must
// be invoked and the message must not arrive on handle_b.
handle_a.error_handler.clear();
handle_a.send(
handle_b.instance_id,
b"should-fail".to_vec(),
b"data".to_vec(),
MessageType::Message,
);
// Give time for the async error path to complete
sleep(Duration::from_millis(100)).await;
// Error handler must have been invoked for the failed send
assert!(
handle_a.error_handler.error_count() >= 1,
"error handler should be invoked for post-shutdown send"
);
// The message must not have been delivered to handle_b
let not_delivered = timeout(
Duration::from_millis(100),
handle_b.streams.message_stream.recv_async(),
)
.await;
assert!(
not_delivered.is_err(),
"post-shutdown message must not arrive at handle_b"
);
handle_b.streams.shutdown_state.teardown_token().cancel();
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment