astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-05-25 21:21:44 +03:00
Code Issues Packages Projects Releases Wiki Activity

Fix hot-path replay bounds and ME control allocations

Browse Source 
Signed-off-by: Alexey <247128645+axkurcom@users.noreply.github.com>
This commit is contained in:
Alexey
2026-05-20 14:05:22 +03:00
parent 70d02910b7
commit 8379b48f69
8 changed files with 91 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/transport/middle_proxy/codec.rs
View File

@@ -9,6 +9,7 @@ use crate::protocol::constants::*;
pub(crate) enum WriterCommand {
Data(Bytes),
DataAndFlush(Bytes),
ControlAndFlush([u8; 12]),
Close,
}
@@ -42,6 +43,13 @@ pub(crate) fn rpc_crc(mode: RpcChecksumMode, data: &[u8]) -> u32 {
}
}
pub(crate) fn build_control_payload(tag: u32, value: u64) -> [u8; 12] {
let mut payload = [0u8; 12];
payload[..4].copy_from_slice(&tag.to_le_bytes());
payload[4..].copy_from_slice(&value.to_le_bytes());
payload
}
pub(crate) fn build_rpc_frame(seq_no: i32, payload: &[u8], crc_mode: RpcChecksumMode) -> Vec<u8> {
let total_len = (4 + 4 + payload.len() + 4) as u32;
let mut frame = Vec::with_capacity(total_len as usize);

18
src/transport/middle_proxy/pool_writer.rs
View File

@@ -5,7 +5,6 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering};
use std::time::{Duration, Instant};
use bytes::Bytes;
use bytes::BytesMut;
use rand::RngExt;
use tokio::sync::mpsc;
@@ -17,7 +16,7 @@ use crate::crypto::SecureRandom;
use crate::error::{ProxyError, Result};
use crate::protocol::constants::{RPC_CLOSE_EXT_U32, RPC_PING_U32};
use super::codec::{RpcWriter, WriterCommand};
use super::codec::{RpcWriter, WriterCommand, build_control_payload};
use super::pool::{MePool, MeWriter, WriterContour};
use super::reader::reader_loop;
use super::wire::build_proxy_req_payload;
@@ -61,6 +60,9 @@ async fn writer_command_loop(
Some(WriterCommand::DataAndFlush(payload)) => {
rpc_writer.send_and_flush(&payload).await?;
}
Some(WriterCommand::ControlAndFlush(payload)) => {
rpc_writer.send_and_flush(&payload).await?;
}
Some(WriterCommand::Close) | None => return Ok(()),
}
}
@@ -130,9 +132,7 @@ async fn ping_loop(
_ = tokio::time::sleep(wait) => {}
}
let sent_id = ping_id;
let mut p = Vec::with_capacity(12);
p.extend_from_slice(&RPC_PING_U32.to_le_bytes());
p.extend_from_slice(&sent_id.to_le_bytes());
let payload = build_control_payload(RPC_PING_U32, sent_id as u64);
{
let mut tracker = ping_tracker_ping.lock().await;
cleanup_tick = cleanup_tick.wrapping_add(1);
@@ -149,7 +149,7 @@ async fn ping_loop(
ping_id = ping_id.wrapping_add(1);
stats_ping.increment_me_keepalive_sent();
if tx_ping
.send(WriterCommand::DataAndFlush(Bytes::from(p)))
.send(WriterCommand::ControlAndFlush(payload))
.await
.is_err()
{
@@ -253,12 +253,10 @@ async fn rpc_proxy_req_signal_loop(
stats_signal.increment_me_rpc_proxy_req_signal_response_total();
}
let mut close_payload = Vec::with_capacity(12);
close_payload.extend_from_slice(&RPC_CLOSE_EXT_U32.to_le_bytes());
close_payload.extend_from_slice(&conn_id.to_le_bytes());
let close_payload = build_control_payload(RPC_CLOSE_EXT_U32, conn_id);
if tx_signal
.send(WriterCommand::DataAndFlush(Bytes::from(close_payload)))
.send(WriterCommand::ControlAndFlush(close_payload))
.await
.is_err()
{

14
src/transport/middle_proxy/reader.rs
View File

@@ -19,7 +19,7 @@ use crate::error::{ProxyError, Result};
use crate::protocol::constants::*;
use crate::stats::Stats;
use super::codec::{RpcChecksumMode, WriterCommand, rpc_crc};
use super::codec::{RpcChecksumMode, WriterCommand, build_control_payload, rpc_crc};
use super::fairness::{
AdmissionDecision, DispatchAction, DispatchFeedback, PressureState, SchedulerDecision,
WorkerFairnessConfig, WorkerFairnessSnapshot, WorkerFairnessState,
@@ -464,10 +464,8 @@ pub(crate) async fn reader_loop(
} else if pt == RPC_PING_U32 && body.len() >= 8 {
let ping_id = i64::from_le_bytes(body[0..8].try_into().unwrap());
trace!(ping_id, "RPC_PING -> RPC_PONG");
let mut pong = Vec::with_capacity(12);
pong.extend_from_slice(&RPC_PONG_U32.to_le_bytes());
pong.extend_from_slice(&ping_id.to_le_bytes());
match tx.try_send(WriterCommand::DataAndFlush(Bytes::from(pong))) {
let pong = build_control_payload(RPC_PONG_U32, ping_id as u64);
match tx.try_send(WriterCommand::ControlAndFlush(pong)) {
Ok(()) => {}
Err(TrySendError::Full(_)) => {
debug!(ping_id, "PONG dropped: writer command channel is full");
@@ -667,10 +665,8 @@ mod tests {
}
async fn send_close_conn(tx: &mpsc::Sender<WriterCommand>, conn_id: u64) {
let mut p = Vec::with_capacity(12);
p.extend_from_slice(&RPC_CLOSE_CONN_U32.to_le_bytes());
p.extend_from_slice(&conn_id.to_le_bytes());
match tx.try_send(WriterCommand::DataAndFlush(Bytes::from(p))) {
let payload = build_control_payload(RPC_CLOSE_CONN_U32, conn_id);
match tx.try_send(WriterCommand::ControlAndFlush(payload)) {
Ok(()) => {}
Err(TrySendError::Full(_)) => {
debug!(

15
src/transport/middle_proxy/send.rs
View File

@@ -7,7 +7,6 @@ use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
use bytes::Bytes;
use tokio::sync::mpsc::error::TrySendError;
use tracing::{debug, warn};
@@ -17,7 +16,7 @@ use crate::network::IpFamily;
use crate::protocol::constants::{RPC_CLOSE_CONN_U32, RPC_CLOSE_EXT_U32};
use super::MePool;
use super::codec::WriterCommand;
use super::codec::{WriterCommand, build_control_payload};
use super::pool::WriterContour;
use super::registry::ConnMeta;
use super::wire::build_proxy_req_payload;
@@ -735,11 +734,9 @@ impl MePool {
pub async fn send_close(self: &Arc<Self>, conn_id: u64) -> Result<()> {
if let Some(w) = self.registry.get_writer(conn_id).await {
let mut p = Vec::with_capacity(12);
p.extend_from_slice(&RPC_CLOSE_EXT_U32.to_le_bytes());
p.extend_from_slice(&conn_id.to_le_bytes());
let payload = build_control_payload(RPC_CLOSE_EXT_U32, conn_id);
if w.tx
.send(WriterCommand::DataAndFlush(Bytes::from(p)))
.send(WriterCommand::ControlAndFlush(payload))
.await
.is_err()
{
@@ -756,10 +753,8 @@ impl MePool {
pub async fn send_close_conn(self: &Arc<Self>, conn_id: u64) -> Result<()> {
if let Some(w) = self.registry.get_writer(conn_id).await {
let mut p = Vec::with_capacity(12);
p.extend_from_slice(&RPC_CLOSE_CONN_U32.to_le_bytes());
p.extend_from_slice(&conn_id.to_le_bytes());
match w.tx.try_send(WriterCommand::DataAndFlush(Bytes::from(p))) {
let payload = build_control_payload(RPC_CLOSE_CONN_U32, conn_id);
match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) {
Ok(()) => {}
Err(TrySendError::Full(cmd)) => {
let _ = tokio::time::timeout(Duration::from_millis(50), w.tx.send(cmd)).await;

1
src/transport/middle_proxy/tests/send_adversarial_tests.rs
View File

@@ -165,6 +165,7 @@ async fn recv_data_count(rx: &mut mpsc::Receiver<WriterCommand>, budget: Duratio
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
Ok(Some(WriterCommand::Data(_))) => data_count += 1,
Ok(Some(WriterCommand::DataAndFlush(_))) => data_count += 1,
Ok(Some(WriterCommand::ControlAndFlush(_))) => data_count += 1,
Ok(Some(WriterCommand::Close)) => {}
Ok(None) => break,
Err(_) => break,