feat: add capnp

Signed-off-by: kjuulh <contact@kjuulh.io>

crates/sq-capnp-interface/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "sq-capnp-interface"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+capnp = { workspace = true }
+bytes = { workspace = true }
+tokio-util = { workspace = true, features = ["codec"] }
+
+[build-dependencies]
+capnpc = { workspace = true }

crates/sq-capnp-interface/build.rs

@@ -0,0 +1,6 @@
+fn main() {
+    capnpc::CompilerCommand::new()
+        .file("schema/data_plane.capnp")
+        .run()
+        .expect("capnp schema compilation failed");
+}

crates/sq-capnp-interface/schema/data_plane.capnp

@@ -0,0 +1,65 @@
+@0xb8f6c1e2a3d4e5f6;
+
+struct MessageHeader {
+  key @0 :Text;
+  value @1 :Data;
+}
+
+struct PublishMessage {
+  topic @0 :Text;
+  key @1 :Data;
+  value @2 :Data;
+  headers @3 :List(MessageHeader);
+}
+
+struct PublishRequest {
+  messages @0 :List(PublishMessage);
+  ackMode @1 :UInt8;
+  producerId @2 :Text;
+}
+
+struct PublishResult {
+  topic @0 :Text;
+  partition @1 :UInt32;
+  offset @2 :UInt64;
+}
+
+struct PublishResponse {
+  results @0 :List(PublishResult);
+}
+
+struct SubscribeRequest {
+  topic @0 :Text;
+  partition @1 :UInt32;
+  consumerGroup @2 :Text;
+  startOffset @3 :UInt64;
+  hasStartOffset @4 :Bool;
+  maxBatchSize @5 :UInt32;
+}
+
+struct ConsumedMessage {
+  offset @0 :UInt64;
+  topic @1 :Text;
+  partition @2 :UInt32;
+  key @3 :Data;
+  value @4 :Data;
+  headers @5 :List(MessageHeader);
+  timestampMs @6 :UInt64;
+}
+
+struct SubscribeResponse {
+  messages @0 :List(ConsumedMessage);
+}
+
+struct AckRequest {
+  consumerGroup @0 :Text;
+  topic @1 :Text;
+  partition @2 :UInt32;
+  offset @3 :UInt64;
+}
+
+struct AckResponse {}
+
+struct ErrorResponse {
+  message @0 :Text;
+}

crates/sq-capnp-interface/src/codec.rs

@@ -0,0 +1,185 @@
+use bytes::{Buf, BufMut, Bytes, BytesMut};
+use tokio_util::codec::{Decoder, Encoder, LengthDelimitedCodec};
+
+// Opcodes
+pub const OP_PUBLISH_REQ: u8 = 0x01;
+pub const OP_PUBLISH_RES: u8 = 0x81;
+pub const OP_SUBSCRIBE_REQ: u8 = 0x02;
+pub const OP_SUBSCRIBE_RES: u8 = 0x82;
+pub const OP_ACK_REQ: u8 = 0x03;
+pub const OP_ACK_RES: u8 = 0x83;
+pub const OP_SUBSCRIBE_END: u8 = 0x84;
+pub const OP_ERROR: u8 = 0xFE;
+
+/// A decoded frame: opcode + capnp payload bytes.
+pub struct Frame {
+    pub opcode: u8,
+    pub payload: Bytes,
+}
+
+/// Codec that wraps `LengthDelimitedCodec` and prepends a 1-byte opcode.
+///
+/// Wire format: `[4-byte big-endian frame length][1-byte opcode][capnp payload]`
+pub struct SqCodec {
+    inner: LengthDelimitedCodec,
+}
+
+impl SqCodec {
+    pub fn new() -> Self {
+        Self {
+            inner: LengthDelimitedCodec::builder()
+                .max_frame_length(16 * 1024 * 1024) // 16 MB
+                .new_codec(),
+        }
+    }
+}
+
+impl Default for SqCodec {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Decoder for SqCodec {
+    type Item = Frame;
+    type Error = std::io::Error;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        match self.inner.decode(src)? {
+            Some(mut buf) => {
+                if buf.is_empty() {
+                    return Err(std::io::Error::new(
+                        std::io::ErrorKind::InvalidData,
+                        "empty frame",
+                    ));
+                }
+                let opcode = buf.get_u8();
+                let payload = buf.freeze();
+                Ok(Some(Frame { opcode, payload }))
+            }
+            None => Ok(None),
+        }
+    }
+}
+
+impl Encoder<Frame> for SqCodec {
+    type Error = std::io::Error;
+
+    fn encode(&mut self, item: Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
+        let mut buf = BytesMut::with_capacity(1 + item.payload.len());
+        buf.put_u8(item.opcode);
+        buf.extend_from_slice(&item.payload);
+        self.inner.encode(buf.freeze(), dst)
+    }
+}
+
+/// Serialize a capnp message builder into bytes.
+pub fn serialize_capnp(builder: &capnp::message::Builder<capnp::message::HeapAllocator>) -> Bytes {
+    let mut buf = Vec::new();
+    capnp::serialize::write_message(&mut buf, builder).expect("capnp serialize failed");
+    Bytes::from(buf)
+}
+
+/// Build a Frame from an opcode and a capnp message builder.
+pub fn build_frame(
+    opcode: u8,
+    builder: &capnp::message::Builder<capnp::message::HeapAllocator>,
+) -> Frame {
+    Frame {
+        opcode,
+        payload: serialize_capnp(builder),
+    }
+}
+
+/// Build an error frame with a text message.
+pub fn error_frame(msg: &str) -> Frame {
+    let mut builder = capnp::message::Builder::new_default();
+    {
+        let mut err = builder.init_root::<crate::data_plane_capnp::error_response::Builder>();
+        err.set_message(msg);
+    }
+    build_frame(OP_ERROR, &builder)
+}
+
+/// Deserialize a capnp message from a byte slice.
+pub fn read_capnp(payload: &[u8]) -> capnp::Result<capnp::message::Reader<capnp::serialize::OwnedSegments>> {
+    let mut cursor = std::io::Cursor::new(payload);
+    capnp::serialize::read_message(
+        &mut cursor,
+        capnp::message::ReaderOptions::new(),
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio_util::codec::{Decoder, Encoder};
+
+    #[test]
+    fn roundtrip_frame() {
+        let mut codec = SqCodec::new();
+        let original = Frame {
+            opcode: OP_PUBLISH_REQ,
+            payload: Bytes::from_static(b"hello"),
+        };
+
+        let mut buf = BytesMut::new();
+        codec
+            .encode(
+                Frame {
+                    opcode: original.opcode,
+                    payload: original.payload.clone(),
+                },
+                &mut buf,
+            )
+            .unwrap();
+
+        let decoded = codec.decode(&mut buf).unwrap().unwrap();
+        assert_eq!(decoded.opcode, OP_PUBLISH_REQ);
+        assert_eq!(decoded.payload, Bytes::from_static(b"hello"));
+    }
+
+    #[test]
+    fn capnp_publish_roundtrip() {
+        // Build a PublishRequest
+        let mut builder = capnp::message::Builder::new_default();
+        {
+            let mut req = builder.init_root::<crate::data_plane_capnp::publish_request::Builder>();
+            req.set_ack_mode(1);
+            req.set_producer_id("test");
+            let mut msgs = req.init_messages(1);
+            let mut msg = msgs.reborrow().get(0);
+            msg.set_topic("orders");
+            msg.set_key(b"key1");
+            msg.set_value(b"value1");
+        }
+
+        let frame = build_frame(OP_PUBLISH_REQ, &builder);
+        assert_eq!(frame.opcode, OP_PUBLISH_REQ);
+
+        // Decode
+        let reader = read_capnp(&frame.payload).unwrap();
+        let req = reader
+            .get_root::<crate::data_plane_capnp::publish_request::Reader>()
+            .unwrap();
+        assert_eq!(req.get_ack_mode(), 1);
+        assert_eq!(req.get_producer_id().unwrap(), "test");
+        let msgs = req.get_messages().unwrap();
+        assert_eq!(msgs.len(), 1);
+        assert_eq!(msgs.get(0).get_topic().unwrap(), "orders");
+        assert_eq!(msgs.get(0).get_key().unwrap(), b"key1");
+        assert_eq!(msgs.get(0).get_value().unwrap(), b"value1");
+    }
+
+    #[test]
+    fn error_frame_roundtrip() {
+        let frame = error_frame("something went wrong");
+        assert_eq!(frame.opcode, OP_ERROR);
+
+        let reader = read_capnp(&frame.payload).unwrap();
+        let err = reader
+            .get_root::<crate::data_plane_capnp::error_response::Reader>()
+            .unwrap();
+        assert_eq!(err.get_message().unwrap(), "something went wrong");
+    }
+}

crates/sq-capnp-interface/src/lib.rs

@@ -0,0 +1,6 @@
+pub mod codec;
+
+#[allow(dead_code)]
+pub mod data_plane_capnp {
+    include!(concat!(env!("OUT_DIR"), "/schema/data_plane_capnp.rs"));
+}