feat(cache): admit gateway-fetched blocks locally (#514)

Author: beinan

crates/talon-cache-client/src/worker_client.rs

@@ -17,12 +17,12 @@ use std::path::Path;
 use std::sync::Arc;
 use std::time::Duration;
 
-use talon_core::{ObjectId, RequestId, Version};
+use talon_core::{BlockId, ObjectId, RequestId, Version};
 use talon_transport::frame::{FrameHeader, MsgType, HEADER_LEN};
 use talon_transport::{
-    decode_error_payload, encode_cached_request, encode_delete, encode_put_header, encode_request,
-    CachedRangeRequest, DataPlaneError, DeleteRequest, Flags, PutRequest, RangeRequest,
-    MAX_CONTROL_PAYLOAD_LEN,
+    decode_error_payload, encode_cached_block_put_header, encode_cached_request, encode_delete,
+    encode_put_header, encode_request, CachedBlockPutRequest, CachedRangeRequest, DataPlaneError,
+    DeleteRequest, Flags, PutRequest, RangeRequest, MAX_CONTROL_PAYLOAD_LEN,
 };
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpStream;
@@ -197,6 +197,74 @@ impl WorkerClient {
         }
     }
 
+    /// Admit one complete, versioned block without asking the worker to access
+    /// its configured backend.
+    pub async fn admit_cached_block(
+        &self,
+        block: &BlockId,
+        object_len: u64,
+        body: &[u8],
+    ) -> Result<(), WorkerError> {
+        let request_id = RequestId::next();
+        let header = encode_cached_block_put_header(
+            request_id.0,
+            &CachedBlockPutRequest {
+                block: block.clone(),
+                object_len,
+                body_len: body.len() as u64,
+            },
+        )?;
+        match self.admit_exchange(&header, body).await {
+            Ok(()) => Ok(()),
+            Err((true, _)) => {
+                let mut stream = self.pool.fresh(&self.addr).await?;
+                self.pool
+                    .with_deadline(
+                        "worker admit_cached_block retry",
+                        streamed_put_timeout(body.len() as u64),
+                        async {
+                            stream.write_all(&header).await?;
+                            stream.write_all(body).await?;
+                            stream.flush().await?;
+                            read_range_reply(&mut stream, 0).await.map(|_| ())
+                        },
+                    )
+                    .await?;
+                self.pool.release(&self.addr, stream);
+                Ok(())
+            }
+            Err((false, error)) => Err(error),
+        }
+    }
+
+    async fn admit_exchange(&self, header: &[u8], body: &[u8]) -> Result<(), (bool, WorkerError)> {
+        let (mut stream, reused) = self
+            .pool
+            .checkout(&self.addr)
+            .await
+            .map_err(|error| (false, WorkerError::from(error)))?;
+        let result = self
+            .pool
+            .with_deadline(
+                "worker admit_cached_block",
+                streamed_put_timeout(body.len() as u64),
+                async {
+                    stream.write_all(header).await?;
+                    stream.write_all(body).await?;
+                    stream.flush().await?;
+                    read_range_reply(&mut stream, 0).await.map(|_| ())
+                },
+            )
+            .await;
+        match result {
+            Ok(()) => {
+                self.pool.release(&self.addr, stream);
+                Ok(())
+            }
+            Err(error) => Err((reused, error)),
+        }
+    }
+
     /// One request/response over a pooled-or-fresh connection.
     ///
     /// On success releases the connection for reuse. On error returns
@@ -575,14 +643,50 @@ mod tests {
     use std::time::Duration;
     use talon_core::Backend;
     use talon_transport::{
-        decode_cached_request, decode_request, encode_error, response_header_ok,
+        decode_cached_block_put_header, decode_cached_request, decode_request, encode_error,
+        response_header_ok,
     };
     use tokio::net::TcpListener;
 
     fn object() -> ObjectId {
         ObjectId::new(Backend::Azure, "container", "path/to/blob.bin")
     }
 
+    #[tokio::test]
+    async fn cache_admission_sends_exact_versioned_block_and_body() {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap().to_string();
+        let block = BlockId::new(object(), 8, 8, Version::new("etag-v2"));
+        let expected = block.clone();
+        tokio::spawn(async move {
+            let (mut socket, _) = listener.accept().await.unwrap();
+            let mut header_bytes = [0_u8; HEADER_LEN];
+            socket.read_exact(&mut header_bytes).await.unwrap();
+            let header = FrameHeader::decode(&header_bytes).unwrap();
+            assert_eq!(header.msg_type, MsgType::AdmitCachedBlock);
+            let mut payload = vec![0_u8; header.length as usize];
+            socket.read_exact(&mut payload).await.unwrap();
+            let mut frame = header_bytes.to_vec();
+            frame.extend_from_slice(&payload);
+            let (_, request) = decode_cached_block_put_header(&frame).unwrap();
+            assert_eq!(request.block, expected);
+            assert_eq!(request.object_len, 13);
+            assert_eq!(request.body_len, 5);
+            let mut body = vec![0_u8; 5];
+            socket.read_exact(&mut body).await.unwrap();
+            assert_eq!(body, b"tail!");
+            socket
+                .write_all(&response_header_ok(header.request_id, 0))
+                .await
+                .unwrap();
+        });
+
+        WorkerClient::new(addr)
+            .admit_cached_block(&block, 13, b"tail!")
+            .await
+            .unwrap();
+    }
+
     /// What a mock write-worker records: the object and body it received.
     type RecordedPut = Arc<std::sync::Mutex<Option<(ObjectId, Vec<u8>)>>>;
 

crates/talon-transport/src/data.rs

@@ -12,7 +12,7 @@
 //! module only handles the request encode/decode and the response header shape.
 
 use serde::{Deserialize, Serialize};
-use talon_core::{ObjectId, Version};
+use talon_core::{BlockId, ObjectId, Version};
 
 use crate::frame::{Flags, FrameError, FrameHeader, MsgType, HEADER_LEN};
 
@@ -102,6 +102,21 @@ pub struct PutRequest {
     pub body_len: u64,
 }
 
+/// A request to admit one complete origin-fetched block into the local cache.
+///
+/// The raw block bytes (`body_len` of them) follow this framed header. The
+/// worker validates the block geometry against `object_len` and never accesses
+/// its backend while handling this operation.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct CachedBlockPutRequest {
+    /// Exact versioned cache identity of the block.
+    pub block: BlockId,
+    /// Authoritative total object length observed by the gateway.
+    pub object_len: u64,
+    /// Number of raw block bytes following this header.
+    pub body_len: u64,
+}
+
 /// A client→worker request to delete an object (#226). No body follows.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct DeleteRequest {
@@ -121,6 +136,9 @@ pub enum DataError {
     /// A non-`Put` frame was handed to the put codec.
     #[error("expected a Put frame, got {0:?}")]
     NotPut(MsgType),
+    /// A non-`AdmitCachedBlock` frame was handed to the admission codec.
+    #[error("expected an AdmitCachedBlock frame, got {0:?}")]
+    NotCachedBlockPut(MsgType),
     /// A non-`Delete` frame was handed to the delete codec.
     #[error("expected a Delete frame, got {0:?}")]
     NotDelete(MsgType),
@@ -234,6 +252,39 @@ pub fn decode_put_header(buf: &[u8]) -> Result<(FrameHeader, PutRequest), DataEr
     Ok((header, req))
 }
 
+/// Encode a cache-admission header. The caller writes exactly `body_len` raw
+/// bytes after the returned frame.
+pub fn encode_cached_block_put_header(
+    request_id: u32,
+    req: &CachedBlockPutRequest,
+) -> Result<Vec<u8>, DataError> {
+    let body = bincode::serialize(req)?;
+    let header = FrameHeader::new(MsgType::AdmitCachedBlock, request_id, body.len() as u32);
+    let mut buf = Vec::with_capacity(HEADER_LEN + body.len());
+    buf.extend_from_slice(&header.encode());
+    buf.extend_from_slice(&body);
+    Ok(buf)
+}
+
+/// Decode a cache-admission header (without its trailing raw body).
+pub fn decode_cached_block_put_header(
+    buf: &[u8],
+) -> Result<(FrameHeader, CachedBlockPutRequest), DataError> {
+    let header = FrameHeader::decode(buf)?;
+    if header.msg_type != MsgType::AdmitCachedBlock {
+        return Err(DataError::NotCachedBlockPut(header.msg_type));
+    }
+    let declared = header.length as usize;
+    let body = &buf[HEADER_LEN..];
+    if body.len() != declared {
+        return Err(DataError::LengthMismatch {
+            declared,
+            actual: body.len(),
+        });
+    }
+    Ok((header, bincode::deserialize(body)?))
+}
+
 /// Encode a [`DeleteRequest`] into `header || bincode(req)` (no body follows).
 pub fn encode_delete(request_id: u32, req: &DeleteRequest) -> Result<Vec<u8>, DataError> {
     let body = bincode::serialize(req)?;
@@ -457,6 +508,21 @@ mod tests {
         assert!((header.length as usize) < 1024);
     }
 
+    #[test]
+    fn cached_block_put_header_round_trips() {
+        let request = CachedBlockPutRequest {
+            block: BlockId::new(obj(), 8, 8, Version::new("etag-v3")),
+            object_len: 13,
+            body_len: 5,
+        };
+        let encoded = encode_cached_block_put_header(43, &request).unwrap();
+        let (header, decoded) = decode_cached_block_put_header(&encoded).unwrap();
+        assert_eq!(header.msg_type, MsgType::AdmitCachedBlock);
+        assert_eq!(header.request_id, 43);
+        assert_eq!(decoded, request);
+        assert!((header.length as usize) < 1024);
+    }
+
     #[test]
     fn delete_round_trips() {
         let req = DeleteRequest { object: obj() };

crates/talon-transport/src/frame.rs

@@ -61,6 +61,9 @@ pub enum MsgType {
     /// Read a versioned range only when every requested byte is resident.
     /// This operation must never access the worker backend.
     GetCachedRange = 6,
+    /// Admit one complete, versioned block into the worker cache without
+    /// accessing the configured backend.
+    AdmitCachedBlock = 7,
 }
 
 impl MsgType {
@@ -74,6 +77,7 @@ impl MsgType {
             4 => MsgType::Ping,
             5 => MsgType::Delete,
             6 => MsgType::GetCachedRange,
+            7 => MsgType::AdmitCachedBlock,
             other => return Err(FrameError::UnknownMsgType(other)),
         })
     }
@@ -215,6 +219,7 @@ mod tests {
             MsgType::Ping,
             MsgType::Delete,
             MsgType::GetCachedRange,
+            MsgType::AdmitCachedBlock,
         ]
         .into_iter()
         .enumerate()

crates/talon-transport/src/lib.rs

@@ -23,9 +23,10 @@ pub use codec::{
     CONTROL_SCHEMA_VERSION, MIN_CONTROL_SCHEMA_VERSION,
 };
 pub use data::{
-    decode_cached_request, decode_delete, decode_error_payload, decode_put_header, decode_request,
-    encode_cached_request, encode_delete, encode_error, encode_put_header, encode_request,
-    encode_typed_error, response_header_ok, CachedRangeRequest, DataError, DataErrorCode,
+    decode_cached_block_put_header, decode_cached_request, decode_delete, decode_error_payload,
+    decode_put_header, decode_request, encode_cached_block_put_header, encode_cached_request,
+    encode_delete, encode_error, encode_put_header, encode_request, encode_typed_error,
+    response_header_ok, CachedBlockPutRequest, CachedRangeRequest, DataError, DataErrorCode,
     DataPlaneError, DeleteRequest, PutRequest, RangeRequest,
 };
 pub use frame::{Flags, FrameError, FrameHeader, MsgType, HEADER_LEN, MAGIC, PROTOCOL_VERSION};

crates/talon-transport/src/limits.rs

@@ -48,7 +48,9 @@ pub fn max_payload_for(msg_type: MsgType) -> u32 {
     match msg_type {
         MsgType::Control => MAX_CONTROL_PAYLOAD_LEN,
         MsgType::Ping => MAX_PING_PAYLOAD_LEN,
-        MsgType::Put | MsgType::Delete | MsgType::GetCachedRange => MAX_CONTROL_PAYLOAD_LEN,
+        MsgType::Put | MsgType::Delete | MsgType::GetCachedRange | MsgType::AdmitCachedBlock => {
+            MAX_CONTROL_PAYLOAD_LEN
+        }
         MsgType::Get | MsgType::GetRange => MAX_PAYLOAD_LEN,
     }
 }

crates/talon-worker/src/main.rs

@@ -1367,6 +1367,97 @@ mod tests {
 
         drop(client);
         server.await.unwrap();
+
+        std::fs::remove_dir_all(root).ok();
+    }
+
+    #[tokio::test]
+    async fn tokio_conn_admits_a_complete_block_without_origin_access() {
+        use talon_core::BlockId;
+        use talon_transport::data::{
+            encode_cached_block_put_header, CachedBlockPutRequest, CachedRangeRequest,
+        };
+
+        let (worker, observability, _, root) = test_worker();
+        observability.readiness().set_backend_ready(true);
+        observability.readiness().set_store_ready(true);
+        observability.readiness().set_control_registered(true);
+        let object = ObjectId::new(talon_core::Backend::Azure, "c", "admitted");
+        let block = BlockId::new(object.clone(), 0, 8, Version::new("origin-v2"));
+        let request = CachedBlockPutRequest {
+            block,
+            object_len: 8,
+            body_len: 8,
+        };
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server_worker = Arc::clone(&worker);
+        let server_observability = Arc::clone(&observability);
+        let server = tokio::spawn(async move {
+            let (stream, _) = listener.accept().await.unwrap();
+            handle_conn(stream, server_worker, server_observability)
+                .await
+                .unwrap();
+        });
+
+        let mut client = TcpStream::connect(address).await.unwrap();
+        client
+            .write_all(&encode_cached_block_put_header(7, &request).unwrap())
+            .await
+            .unwrap();
+        client.write_all(b"gateway!").await.unwrap();
+        let mut reply = [0_u8; HEADER_LEN];
+        client.read_exact(&mut reply).await.unwrap();
+        let reply = FrameHeader::decode(&reply).unwrap();
+        assert_eq!(reply.length, 0);
+        assert!(!reply.flags.contains(talon_transport::Flags::ERROR));
+        assert_eq!(
+            worker
+                .serve_cached(&CachedRangeRequest {
+                    object: object.clone(),
+                    version: Version::new("origin-v2"),
+                    offset: 0,
+                    len: 8,
+                })
+                .await
+                .unwrap(),
+            Bytes::from_static(b"gateway!")
+        );
+        drop(client);
+        server.await.unwrap();
+
+        let truncated = CachedBlockPutRequest {
+            block: BlockId::new(object.clone(), 0, 8, Version::new("truncated")),
+            object_len: 8,
+            body_len: 8,
+        };
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server_worker = Arc::clone(&worker);
+        let server_observability = Arc::clone(&observability);
+        let server = tokio::spawn(async move {
+            let (stream, _) = listener.accept().await.unwrap();
+            handle_conn(stream, server_worker, server_observability)
+                .await
+                .is_err()
+        });
+        let mut client = TcpStream::connect(address).await.unwrap();
+        client
+            .write_all(&encode_cached_block_put_header(8, &truncated).unwrap())
+            .await
+            .unwrap();
+        client.write_all(b"short").await.unwrap();
+        client.shutdown().await.unwrap();
+        assert!(server.await.unwrap());
+        assert!(worker
+            .serve_cached(&CachedRangeRequest {
+                object,
+                version: Version::new("truncated"),
+                offset: 0,
+                len: 1,
+            })
+            .await
+            .is_err());
         std::fs::remove_dir_all(root).ok();
     }