trending/OpenShell
1
0
Fork 0
mirror of https://github.com/NVIDIA/OpenShell.git synced 2026-11-01 07:53:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): accept the full reconnect log buffer (#5049)

Browse Source 
* fix(supervisor): batch buffered log replay after reconnect

Signed-off-by: divesh <dgude@nvidia.com>

* fix(gateway): accept the full reconnect log buffer

Signed-off-by: divesh <dgude@nvidia.com>

* fix(logging): batch reconnect replay and warn on truncation

Signed-off-by: divesh <dgude@nvidia.com>

---------

Signed-off-by: divesh <dgude@nvidia.com>
This commit is contained in:
Divesh
2026-10-30 05:11:21 +00:00
committed by GitHub
parent cc9995438f
commit 3a5dfcd2f1
2 changed files with 266 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/openshell-server/src/grpc/policy.rs
+156 -1
View File
@@ -5412,6 +5412,8 @@ pub(super) async fn handle_push_sandbox_logs(
state: &Arc<ServerState>,
request: Request<tonic::Streaming<PushSandboxLogsRequest>>,
) -> Result<Response<PushSandboxLogsResponse>, Status> {
const MAX_BATCH_LINES: usize = 100;
let principal = request
.extensions()
.get::<Principal>()
@@ -5437,9 +5439,20 @@ pub(super) async fn handle_push_sandbox_logs(
)
.await?;
let logs: Vec<_> = batch.logs.into_iter().take(100).collect();
let received_lines = batch.logs.len();
let logs: Vec<_> = batch.logs.into_iter().take(MAX_BATCH_LINES).collect();
validate_sandbox_log_timestamps(&logs)?;
if received_lines > logs.len() {
warn!(
sandbox_id = %batch.sandbox_id,
received_lines,
accepted_lines = logs.len(),
dropped_lines = received_lines - logs.len(),
"Truncated sandbox log batch"
);
}
for log in logs {
let mut log = log;
log.source = "sandbox".to_string();
@@ -8265,6 +8278,148 @@ mod tests {
request
}
#[tokio::test]
async fn push_sandbox_logs_accepts_batched_replay_and_warns_on_truncation() {
use openshell_core::proto::open_shell_client::OpenShellClient;
use openshell_core::proto::open_shell_server::OpenShellServer;
use tokio_stream::wrappers::TcpListenerStream;
let state = test_server_state().await;
let sandbox_id = "sb-log-replay";
state
.store
.put_message(&test_sandbox(
sandbox_id,
"log-replay",
ProtoSandboxPolicy::default(),
vec![],
))
.await
.unwrap();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let principal = crate::auth::principal::test_support::sandbox(sandbox_id);
let service = OpenShellServer::with_interceptor(
crate::grpc::OpenShellService::new(state.clone()),
move |mut request: Request<()>| {
request.extensions_mut().insert(principal.clone());
Ok(request)
},
);
let server = tokio::spawn(
tonic::transport::Server::builder()
.add_service(service)
.serve_with_incoming(TcpListenerStream::new(listener)),
);
let mut client = OpenShellClient::connect(format!("http://{address}"))
.await
.unwrap();
let diagnostics = tempfile::NamedTempFile::new().unwrap();
let subscriber = tracing_subscriber::fmt()
.json()
.with_env_filter("off,openshell_server::grpc::policy=warn")
.with_writer(std::sync::Mutex::new(diagnostics.reopen().unwrap()))
.finish();
// This test uses a current-thread runtime, including the gateway tasks.
let _capture = tracing::subscriber::set_default(subscriber);
let mut seen_warnings = 0;
for (batch_size, count, accepted_count) in [
(200, 200, 100),
(1, 0, 0),
(1, 1, 1),
(50, 50, 50),
(100, 100, 100),
(101, 101, 100),
(199, 199, 100),
(201, 201, 100),
(50, 200, 200),
(50, 201, 201),
] {
state.tracing_log_bus.remove(sandbox_id);
let logs: Vec<_> = (0..count)
.map(|index| SandboxLogLine {
sandbox_id: "untrusted-id".to_string(),
event_time: openshell_core::time::timestamp_from_millis(index).ok(),
level: "INFO".to_string(),
target: "reconnect-test".to_string(),
message: format!("line {index}"),
source: "gateway".to_string(),
fields: HashMap::from([("index".to_string(), index.to_string())]),
})
.collect();
let expected: Vec<_> = logs
.iter()
.take(accepted_count)
.cloned()
.map(|mut log| {
log.sandbox_id = sandbox_id.to_string();
log.source = "sandbox".to_string();
log
})
.collect();
let received_lines = logs.len();
let requests: Vec<_> = logs
.chunks(batch_size)
.map(|chunk| PushSandboxLogsRequest {
sandbox_id: sandbox_id.to_string(),
logs: chunk.to_vec(),
})
.collect();
client
.push_sandbox_logs(tokio_stream::iter(requests))
.await
.unwrap();
let received: Vec<_> = state
.tracing_log_bus
.tail(sandbox_id, 201)
.into_iter()
.map(|entry| {
let Some(openshell_core::proto::sandbox_stream_event::Payload::Log(log)) =
entry.event.payload
else {
panic!("expected a log event");
};
log
})
.collect();
assert_eq!(
received.len(),
expected.len(),
"gateway ingestion of {count} log lines in batches of {batch_size}"
);
assert_eq!(
received, expected,
"{count} log lines, batches of {batch_size}"
);
let warnings: Vec<serde_json::Value> = std::fs::read_to_string(diagnostics.path())
.unwrap()
.lines()
.map(|line| serde_json::from_str(line).unwrap())
.collect();
let emitted = &warnings[seen_warnings..];
assert_eq!(emitted.len(), usize::from(received_lines > accepted_count));
if let Some(warning) = emitted.first() {
assert_eq!(warning["level"], "WARN");
assert_eq!(
warning["fields"],
serde_json::json!({
"message": "Truncated sandbox log batch",
"sandbox_id": sandbox_id,
"received_lines": received_lines,
"accepted_lines": accepted_count,
"dropped_lines": received_lines - accepted_count,
})
);
}
seen_warnings = warnings.len();
}
server.abort();
}
#[test]
fn sandbox_log_timestamp_validation_rejects_malformed_batches_atomically() {
let logs = vec![
crates/openshell-supervisor-process/src/log_push.rs
+110 -16
View File
@@ -180,13 +180,15 @@ pub fn start_log_push(
const MAX_BACKOFF: tokio::time::Duration = tokio::time::Duration::from_secs(30);
/// Initial backoff delay after a connection failure.
const INITIAL_BACKOFF: tokio::time::Duration = tokio::time::Duration::from_secs(1);
/// Lines per regular or reconnect-replay batch, below the gateway's ingest cap.
const LOG_PUSH_BATCH_LINES: usize = 50;
async fn run_push_loop(
gateway: GatewayEndpoint,
sandbox_id: String,
mut rx: mpsc::Receiver<SandboxLogLine>,
) {
let mut batch = Vec::with_capacity(50);
let mut batch = Vec::with_capacity(LOG_PUSH_BATCH_LINES);
let mut backoff = INITIAL_BACKOFF;
let mut attempt: u64 = 0;
@@ -236,20 +238,13 @@ async fn run_push_loop(
});
// --- Flush any lines buffered during reconnect ---
if !batch.is_empty() {
let lines = std::mem::take(&mut batch);
if push_tx
.send(PushSandboxLogsRequest {
sandbox_id: sandbox_id.clone(),
logs: lines,
})
.await
.is_err()
{
// RPC died immediately — go back to reconnect.
backoff = INITIAL_BACKOFF;
continue;
}
if flush_reconnect_buffer(&push_tx, &sandbox_id, &mut batch)
.await
.is_err()
{
// RPC died immediately — go back to reconnect.
backoff = INITIAL_BACKOFF;
continue;
}
// --- Batch and send loop (runs until stream breaks) ---
@@ -274,7 +269,7 @@ async fn run_push_loop(
return;
};
batch.push(line);
if batch.len() >= 50 {
if batch.len() >= LOG_PUSH_BATCH_LINES {
let lines = std::mem::take(&mut batch);
if push_tx.send(PushSandboxLogsRequest {
sandbox_id: sandbox_id.clone(),
@@ -316,6 +311,24 @@ async fn run_push_loop(
}
}
async fn flush_reconnect_buffer(
tx: &mpsc::Sender<PushSandboxLogsRequest>,
sandbox_id: &str,
batch: &mut Vec<SandboxLogLine>,
) -> Result<(), mpsc::error::SendError<PushSandboxLogsRequest>> {
// The reconnect buffer can hold 200 lines, but the gateway only ingests
// the first 100 of each request. Replay using the regular batch size.
while !batch.is_empty() {
let count = batch.len().min(LOG_PUSH_BATCH_LINES);
tx.send(PushSandboxLogsRequest {
sandbox_id: sandbox_id.to_string(),
logs: batch.drain(..count).collect(),
})
.await?;
}
Ok(())
}
/// Drain incoming log lines during a backoff delay so the tracing layer's
/// `try_send` doesn't fill up. Lines received during backoff are kept in `batch`
/// (up to a limit) so they can be sent after reconnecting.
@@ -524,6 +537,87 @@ mod tests {
}
}
#[tokio::test]
async fn reconnect_replay_retains_all_buffered_lines_at_the_gateway_ingest_cap() {
let expected: Vec<_> = (0..200).map(|i| test_line(&format!("line {i}"))).collect();
let (tx, mut rx) = mpsc::channel(LOG_PUSH_QUEUE_LINES);
for line in &expected {
tx.try_send(line.clone()).unwrap();
}
drop(tx);
let mut batch = Vec::new();
drain_during_backoff(&mut rx, &mut batch, INITIAL_BACKOFF).await;
assert_eq!(batch, expected);
let (push_tx, mut push_rx) = mpsc::channel(32);
flush_reconnect_buffer(&push_tx, "sb-test", &mut batch)
.await
.unwrap();
drop(push_tx);
let mut received = Vec::new();
let mut batch_sizes = Vec::new();
while let Some(request) = push_rx.recv().await {
assert_eq!(request.sandbox_id, "sb-test");
batch_sizes.push(request.logs.len());
// Match handle_push_sandbox_logs in openshell-server/src/grpc/policy.rs.
received.extend(request.logs.into_iter().take(100));
}
eprintln!(
"reconnect replay: batches={batch_sizes:?}, retained={}/{}",
received.len(),
expected.len()
);
assert_eq!(
received.len(),
expected.len(),
"the gateway ingest cap must not truncate reconnect-buffered logs"
);
assert_eq!(received, expected, "replay must preserve order and fields");
assert!(batch.is_empty());
}
#[tokio::test]
async fn reconnect_replay_preserves_empty_and_partial_batches() {
for count in [0, 1, 49, 50, 51, 100, 101, 199] {
let expected: Vec<_> = (0..count)
.map(|i| test_line(&format!("line {i}")))
.collect();
let mut batch = expected.clone();
let (tx, mut rx) = mpsc::channel(32);
flush_reconnect_buffer(&tx, "sb-test", &mut batch)
.await
.unwrap();
drop(tx);
let mut received = Vec::new();
while let Some(request) = rx.recv().await {
assert_eq!(request.sandbox_id, "sb-test");
assert!(!request.logs.is_empty());
assert!(request.logs.len() <= LOG_PUSH_BATCH_LINES);
received.extend(request.logs);
}
assert_eq!(received, expected, "replay of {count} lines");
assert!(batch.is_empty());
}
}
#[tokio::test]
async fn reconnect_replay_returns_an_error_if_the_stream_closes() {
let mut batch: Vec<_> = (0..200).map(|i| test_line(&format!("line {i}"))).collect();
let (tx, rx) = mpsc::channel(1);
drop(rx);
let error = flush_reconnect_buffer(&tx, "sb-test", &mut batch)
.await
.expect_err("a closed stream must return to the reconnect loop");
assert_eq!(error.0.logs.len(), LOG_PUSH_BATCH_LINES);
assert_eq!(batch.len(), 200 - LOG_PUSH_BATCH_LINES);
assert_eq!(batch[0].message, "line 50");
}
#[tokio::test]
async fn drain_during_backoff_buffers_up_to_the_cap_and_drops_the_rest() {
let (tx, mut rx) = mpsc::channel::<SandboxLogLine>(1024);