455 lines
15 KiB
Rust
455 lines
15 KiB
Rust
//! TCP streaming protocol for inter-agent exchange.
|
|
//!
|
|
//! Uses a simplified length-prefixed binary framing:
|
|
//!
|
|
//! ```text
|
|
//! [4 bytes: payload length (big-endian)] [1 byte: msg_type] [3 bytes: msg_id] [payload]
|
|
//! ```
|
|
//!
|
|
//! Message types follow the spec in 10-operations-api.md section 6.2.
|
|
|
|
use std::sync::Arc;
|
|
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
|
|
use rvf_runtime::QueryOptions;
|
|
|
|
use crate::http::SharedStore;
|
|
|
|
/// TCP message types (client -> server).
|
|
const MSG_QUERY: u8 = 0x01;
|
|
const MSG_INGEST: u8 = 0x02;
|
|
const MSG_DELETE: u8 = 0x03;
|
|
const MSG_STATUS: u8 = 0x04;
|
|
|
|
/// TCP message types (server -> client).
|
|
const MSG_QUERY_RESULT: u8 = 0x81;
|
|
const MSG_INGEST_ACK: u8 = 0x82;
|
|
const MSG_DELETE_ACK: u8 = 0x83;
|
|
const MSG_STATUS_RESP: u8 = 0x84;
|
|
const MSG_ERROR: u8 = 0xFF;
|
|
|
|
/// Maximum frame payload: 16 MB.
|
|
const MAX_FRAME_SIZE: u32 = 16 * 1024 * 1024;
|
|
|
|
/// Start the TCP listener on the given address.
|
|
pub async fn serve_tcp(addr: &str, store: SharedStore) -> std::io::Result<()> {
|
|
let listener = TcpListener::bind(addr).await?;
|
|
|
|
loop {
|
|
let (stream, _peer) = listener.accept().await?;
|
|
let store = Arc::clone(&store);
|
|
tokio::spawn(async move {
|
|
if let Err(e) = handle_connection(stream, store).await {
|
|
eprintln!("tcp connection error: {e}");
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
async fn handle_connection(mut stream: TcpStream, store: SharedStore) -> std::io::Result<()> {
|
|
loop {
|
|
// Read frame header: 4 bytes length + 1 byte msg_type + 3 bytes msg_id = 8 bytes
|
|
let mut header = [0u8; 8];
|
|
match stream.read_exact(&mut header).await {
|
|
Ok(_) => {}
|
|
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(()),
|
|
Err(e) => return Err(e),
|
|
}
|
|
|
|
let payload_len = u32::from_be_bytes([header[0], header[1], header[2], header[3]]);
|
|
let msg_type = header[4];
|
|
let msg_id = [header[5], header[6], header[7]];
|
|
|
|
if payload_len > MAX_FRAME_SIZE {
|
|
send_error(&mut stream, &msg_id, 0x0104, "frame too large").await?;
|
|
return Ok(());
|
|
}
|
|
|
|
let mut payload = vec![0u8; payload_len as usize];
|
|
stream.read_exact(&mut payload).await?;
|
|
|
|
let response = match msg_type {
|
|
MSG_QUERY => handle_query(&payload, &store).await,
|
|
MSG_INGEST => handle_ingest(&payload, &store).await,
|
|
MSG_DELETE => handle_delete(&payload, &store).await,
|
|
MSG_STATUS => handle_status(&store).await,
|
|
_ => Err(TcpError {
|
|
code: 0x0107,
|
|
message: "unknown message type".into(),
|
|
}),
|
|
};
|
|
|
|
match response {
|
|
Ok((resp_type, resp_payload)) => {
|
|
send_frame(&mut stream, resp_type, &msg_id, &resp_payload).await?;
|
|
}
|
|
Err(e) => {
|
|
send_error(&mut stream, &msg_id, e.code, &e.message).await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
struct TcpError {
|
|
code: u16,
|
|
message: String,
|
|
}
|
|
|
|
async fn send_frame(
|
|
stream: &mut TcpStream,
|
|
msg_type: u8,
|
|
msg_id: &[u8; 3],
|
|
payload: &[u8],
|
|
) -> std::io::Result<()> {
|
|
let len = payload.len() as u32;
|
|
let mut frame = Vec::with_capacity(8 + payload.len());
|
|
frame.extend_from_slice(&len.to_be_bytes());
|
|
frame.push(msg_type);
|
|
frame.extend_from_slice(msg_id);
|
|
frame.extend_from_slice(payload);
|
|
stream.write_all(&frame).await
|
|
}
|
|
|
|
async fn send_error(
|
|
stream: &mut TcpStream,
|
|
msg_id: &[u8; 3],
|
|
code: u16,
|
|
description: &str,
|
|
) -> std::io::Result<()> {
|
|
let desc_bytes = description.as_bytes();
|
|
let mut payload = Vec::with_capacity(4 + desc_bytes.len());
|
|
payload.extend_from_slice(&code.to_le_bytes());
|
|
payload.extend_from_slice(&(desc_bytes.len() as u16).to_le_bytes());
|
|
payload.extend_from_slice(desc_bytes);
|
|
send_frame(stream, MSG_ERROR, msg_id, &payload).await
|
|
}
|
|
|
|
/// Handle a QUERY message. Payload is a simplified JSON-encoded query
|
|
/// for ease of inter-agent use (vector, k as little-endian).
|
|
async fn handle_query(payload: &[u8], store: &SharedStore) -> Result<(u8, Vec<u8>), TcpError> {
|
|
// Simplified binary protocol:
|
|
// [4 bytes: k (LE)] [4 bytes: dim (LE)] [dim * 4 bytes: vector f32s (LE)]
|
|
if payload.len() < 8 {
|
|
return Err(TcpError {
|
|
code: 0x0200,
|
|
message: "payload too short for query".into(),
|
|
});
|
|
}
|
|
|
|
let k = u32::from_le_bytes([payload[0], payload[1], payload[2], payload[3]]) as usize;
|
|
let dim = u32::from_le_bytes([payload[4], payload[5], payload[6], payload[7]]) as usize;
|
|
|
|
let expected = 8 + dim * 4;
|
|
if payload.len() < expected {
|
|
return Err(TcpError {
|
|
code: 0x0200,
|
|
message: "payload too short for vector data".into(),
|
|
});
|
|
}
|
|
|
|
let mut vector = Vec::with_capacity(dim);
|
|
for i in 0..dim {
|
|
let offset = 8 + i * 4;
|
|
let val = f32::from_le_bytes([
|
|
payload[offset],
|
|
payload[offset + 1],
|
|
payload[offset + 2],
|
|
payload[offset + 3],
|
|
]);
|
|
vector.push(val);
|
|
}
|
|
|
|
let results = {
|
|
let s = store.lock().await;
|
|
s.query(&vector, k, &QueryOptions::default())
|
|
.map_err(|e| TcpError {
|
|
code: 0x0200,
|
|
message: format!("{e:?}"),
|
|
})?
|
|
};
|
|
|
|
// Response: [4 bytes: result_count (LE)] [per result: 8 bytes id (LE) + 4 bytes dist (LE)]
|
|
let mut resp = Vec::with_capacity(4 + results.len() * 12);
|
|
resp.extend_from_slice(&(results.len() as u32).to_le_bytes());
|
|
for r in &results {
|
|
resp.extend_from_slice(&r.id.to_le_bytes());
|
|
resp.extend_from_slice(&r.distance.to_le_bytes());
|
|
}
|
|
|
|
Ok((MSG_QUERY_RESULT, resp))
|
|
}
|
|
|
|
/// Handle an INGEST message.
|
|
/// Binary payload: [4 bytes: count (LE)] [2 bytes: dim (LE)] [per vector: 8 bytes id (LE) + dim*4 bytes data (LE)]
|
|
async fn handle_ingest(payload: &[u8], store: &SharedStore) -> Result<(u8, Vec<u8>), TcpError> {
|
|
if payload.len() < 6 {
|
|
return Err(TcpError {
|
|
code: 0x0300,
|
|
message: "payload too short for ingest header".into(),
|
|
});
|
|
}
|
|
|
|
let count = u32::from_le_bytes([payload[0], payload[1], payload[2], payload[3]]) as usize;
|
|
let dim = u16::from_le_bytes([payload[4], payload[5]]) as usize;
|
|
|
|
let entry_size = 8 + dim * 4;
|
|
let expected = 6 + count * entry_size;
|
|
if payload.len() < expected {
|
|
return Err(TcpError {
|
|
code: 0x0300,
|
|
message: "payload too short for ingest data".into(),
|
|
});
|
|
}
|
|
|
|
let mut ids = Vec::with_capacity(count);
|
|
let mut vectors: Vec<Vec<f32>> = Vec::with_capacity(count);
|
|
|
|
let mut offset = 6;
|
|
for _ in 0..count {
|
|
let id = u64::from_le_bytes([
|
|
payload[offset],
|
|
payload[offset + 1],
|
|
payload[offset + 2],
|
|
payload[offset + 3],
|
|
payload[offset + 4],
|
|
payload[offset + 5],
|
|
payload[offset + 6],
|
|
payload[offset + 7],
|
|
]);
|
|
offset += 8;
|
|
|
|
let mut vec_data = Vec::with_capacity(dim);
|
|
for _ in 0..dim {
|
|
let val = f32::from_le_bytes([
|
|
payload[offset],
|
|
payload[offset + 1],
|
|
payload[offset + 2],
|
|
payload[offset + 3],
|
|
]);
|
|
vec_data.push(val);
|
|
offset += 4;
|
|
}
|
|
|
|
ids.push(id);
|
|
vectors.push(vec_data);
|
|
}
|
|
|
|
let vec_refs: Vec<&[f32]> = vectors.iter().map(|v| v.as_slice()).collect();
|
|
|
|
let result = {
|
|
let mut s = store.lock().await;
|
|
s.ingest_batch(&vec_refs, &ids, None)
|
|
.map_err(|e| TcpError {
|
|
code: 0x0300,
|
|
message: format!("{e:?}"),
|
|
})?
|
|
};
|
|
|
|
// Response: [8 bytes: accepted (LE)] [8 bytes: rejected (LE)] [4 bytes: epoch (LE)]
|
|
let mut resp = Vec::with_capacity(20);
|
|
resp.extend_from_slice(&result.accepted.to_le_bytes());
|
|
resp.extend_from_slice(&result.rejected.to_le_bytes());
|
|
resp.extend_from_slice(&result.epoch.to_le_bytes());
|
|
|
|
Ok((MSG_INGEST_ACK, resp))
|
|
}
|
|
|
|
/// Handle a DELETE message.
|
|
/// Binary payload: [4 bytes: count (LE)] [per id: 8 bytes (LE)]
|
|
async fn handle_delete(payload: &[u8], store: &SharedStore) -> Result<(u8, Vec<u8>), TcpError> {
|
|
if payload.len() < 4 {
|
|
return Err(TcpError {
|
|
code: 0x0300,
|
|
message: "payload too short for delete".into(),
|
|
});
|
|
}
|
|
|
|
let count = u32::from_le_bytes([payload[0], payload[1], payload[2], payload[3]]) as usize;
|
|
|
|
if payload.len() < 4 + count * 8 {
|
|
return Err(TcpError {
|
|
code: 0x0300,
|
|
message: "payload too short for delete ids".into(),
|
|
});
|
|
}
|
|
|
|
let mut ids = Vec::with_capacity(count);
|
|
let mut offset = 4;
|
|
for _ in 0..count {
|
|
let id = u64::from_le_bytes([
|
|
payload[offset],
|
|
payload[offset + 1],
|
|
payload[offset + 2],
|
|
payload[offset + 3],
|
|
payload[offset + 4],
|
|
payload[offset + 5],
|
|
payload[offset + 6],
|
|
payload[offset + 7],
|
|
]);
|
|
ids.push(id);
|
|
offset += 8;
|
|
}
|
|
|
|
let result = {
|
|
let mut s = store.lock().await;
|
|
s.delete(&ids).map_err(|e| TcpError {
|
|
code: 0x0300,
|
|
message: format!("{e:?}"),
|
|
})?
|
|
};
|
|
|
|
// Response: [8 bytes: deleted (LE)] [4 bytes: epoch (LE)]
|
|
let mut resp = Vec::with_capacity(12);
|
|
resp.extend_from_slice(&result.deleted.to_le_bytes());
|
|
resp.extend_from_slice(&result.epoch.to_le_bytes());
|
|
|
|
Ok((MSG_DELETE_ACK, resp))
|
|
}
|
|
|
|
/// Handle a STATUS message (no payload needed).
|
|
async fn handle_status(store: &SharedStore) -> Result<(u8, Vec<u8>), TcpError> {
|
|
let s = store.lock().await;
|
|
let st = s.status();
|
|
|
|
// Simplified STATUS_RESP: epoch(4) + total_vectors(8) + total_segments(4) +
|
|
// file_size(8) + profile_id(1) + read_only(1)
|
|
let mut resp = Vec::with_capacity(26);
|
|
resp.extend_from_slice(&st.current_epoch.to_le_bytes());
|
|
resp.extend_from_slice(&st.total_vectors.to_le_bytes());
|
|
resp.extend_from_slice(&st.total_segments.to_le_bytes());
|
|
resp.extend_from_slice(&st.file_size.to_le_bytes());
|
|
resp.push(st.profile_id);
|
|
resp.push(if st.read_only { 1 } else { 0 });
|
|
|
|
Ok((MSG_STATUS_RESP, resp))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rvf_runtime::{RvfOptions, RvfStore};
|
|
use tempfile::TempDir;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::TcpListener;
|
|
use tokio::sync::Mutex;
|
|
|
|
fn create_test_store() -> (TempDir, SharedStore) {
|
|
let dir = TempDir::new().unwrap();
|
|
let path = dir.path().join("test.rvf");
|
|
let options = RvfOptions {
|
|
dimension: 4,
|
|
..Default::default()
|
|
};
|
|
let store = RvfStore::create(&path, options).unwrap();
|
|
(dir, Arc::new(Mutex::new(store)))
|
|
}
|
|
|
|
fn build_frame(msg_type: u8, msg_id: [u8; 3], payload: &[u8]) -> Vec<u8> {
|
|
let len = payload.len() as u32;
|
|
let mut frame = Vec::with_capacity(8 + payload.len());
|
|
frame.extend_from_slice(&len.to_be_bytes());
|
|
frame.push(msg_type);
|
|
frame.extend_from_slice(&msg_id);
|
|
frame.extend_from_slice(payload);
|
|
frame
|
|
}
|
|
|
|
async fn read_frame(stream: &mut TcpStream) -> (u8, [u8; 3], Vec<u8>) {
|
|
let mut header = [0u8; 8];
|
|
stream.read_exact(&mut header).await.unwrap();
|
|
let payload_len = u32::from_be_bytes([header[0], header[1], header[2], header[3]]);
|
|
let msg_type = header[4];
|
|
let msg_id = [header[5], header[6], header[7]];
|
|
let mut payload = vec![0u8; payload_len as usize];
|
|
stream.read_exact(&mut payload).await.unwrap();
|
|
(msg_type, msg_id, payload)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tcp_status() {
|
|
let (_dir, store) = create_test_store();
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let addr = listener.local_addr().unwrap();
|
|
|
|
let store_clone = Arc::clone(&store);
|
|
tokio::spawn(async move {
|
|
let (stream, _) = listener.accept().await.unwrap();
|
|
handle_connection(stream, store_clone).await.unwrap();
|
|
});
|
|
|
|
let mut client = TcpStream::connect(addr).await.unwrap();
|
|
|
|
// Send STATUS request (no payload)
|
|
let frame = build_frame(MSG_STATUS, [0, 0, 1], &[]);
|
|
client.write_all(&frame).await.unwrap();
|
|
|
|
let (msg_type, msg_id, payload) = read_frame(&mut client).await;
|
|
assert_eq!(msg_type, MSG_STATUS_RESP);
|
|
assert_eq!(msg_id, [0, 0, 1]);
|
|
assert_eq!(payload.len(), 26);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tcp_ingest_and_query() {
|
|
let (_dir, store) = create_test_store();
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let addr = listener.local_addr().unwrap();
|
|
|
|
let store_clone = Arc::clone(&store);
|
|
tokio::spawn(async move {
|
|
let (stream, _) = listener.accept().await.unwrap();
|
|
handle_connection(stream, store_clone).await.unwrap();
|
|
});
|
|
|
|
let mut client = TcpStream::connect(addr).await.unwrap();
|
|
|
|
// Ingest 2 vectors of dim=4
|
|
let mut ingest_payload = Vec::new();
|
|
ingest_payload.extend_from_slice(&2u32.to_le_bytes()); // count
|
|
ingest_payload.extend_from_slice(&4u16.to_le_bytes()); // dim
|
|
// Vector 1: id=1, data=[1,0,0,0]
|
|
ingest_payload.extend_from_slice(&1u64.to_le_bytes());
|
|
ingest_payload.extend_from_slice(&1.0f32.to_le_bytes());
|
|
ingest_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
ingest_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
ingest_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
// Vector 2: id=2, data=[0,1,0,0]
|
|
ingest_payload.extend_from_slice(&2u64.to_le_bytes());
|
|
ingest_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
ingest_payload.extend_from_slice(&1.0f32.to_le_bytes());
|
|
ingest_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
ingest_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
|
|
let frame = build_frame(MSG_INGEST, [0, 0, 2], &ingest_payload);
|
|
client.write_all(&frame).await.unwrap();
|
|
|
|
let (msg_type, _, payload) = read_frame(&mut client).await;
|
|
assert_eq!(msg_type, MSG_INGEST_ACK);
|
|
let accepted = u64::from_le_bytes(payload[0..8].try_into().unwrap());
|
|
assert_eq!(accepted, 2);
|
|
|
|
// Query for nearest to [1,0,0,0], k=1
|
|
let mut query_payload = Vec::new();
|
|
query_payload.extend_from_slice(&1u32.to_le_bytes()); // k
|
|
query_payload.extend_from_slice(&4u32.to_le_bytes()); // dim
|
|
query_payload.extend_from_slice(&1.0f32.to_le_bytes());
|
|
query_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
query_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
query_payload.extend_from_slice(&0.0f32.to_le_bytes());
|
|
|
|
let frame = build_frame(MSG_QUERY, [0, 0, 3], &query_payload);
|
|
client.write_all(&frame).await.unwrap();
|
|
|
|
let (msg_type, _, payload) = read_frame(&mut client).await;
|
|
assert_eq!(msg_type, MSG_QUERY_RESULT);
|
|
let count = u32::from_le_bytes(payload[0..4].try_into().unwrap());
|
|
assert_eq!(count, 1);
|
|
let id = u64::from_le_bytes(payload[4..12].try_into().unwrap());
|
|
assert_eq!(id, 1);
|
|
}
|
|
}
|