feat: add capnp

Signed-off-by: kjuulh <contact@kjuulh.io>
2026-02-27 12:15:35 +01:00
parent 3162971c89
commit 749ae245c7
115 changed files with 16596 additions and 31 deletions
--- a/crates/sq-sdk/src/capnp_producer.rs
+++ b/crates/sq-sdk/src/capnp_producer.rs
@@ -0,0 +1,145 @@
+use sq_capnp_interface::codec::{self, OP_ERROR, OP_PUBLISH_REQ, OP_PUBLISH_RES};
+use sq_capnp_interface::data_plane_capnp;
+
+use crate::capnp_connection::Connection;
+use crate::error::SqError;
+use crate::producer::{ProducerMessage, SendResult};
+use crate::types::AckMode;
+
+/// Configuration for an SQ producer (Cap'n Proto transport).
+pub struct ProducerConfig {
+    /// Server address (e.g., "127.0.0.1:6064").
+    pub address: String,
+    /// Acknowledgment mode.
+    pub ack_mode: AckMode,
+    /// Producer identifier.
+    pub producer_id: String,
+}
+
+impl Default for ProducerConfig {
+    fn default() -> Self {
+        Self {
+            address: "127.0.0.1:6064".to_string(),
+            ack_mode: AckMode::All,
+            producer_id: "default".to_string(),
+        }
+    }
+}
+
+/// SQ producer using Cap'n Proto over TCP.
+pub struct Producer {
+    conn: Connection,
+    config: ProducerConfig,
+}
+
+impl Producer {
+    /// Connect to an SQ server and create a new producer.
+    pub async fn connect(config: ProducerConfig) -> Result<Self, SqError> {
+        let conn = Connection::connect(&config.address).await?;
+        Ok(Self { conn, config })
+    }
+
+    /// Send a single message.
+    pub async fn send(
+        &mut self,
+        topic: &str,
+        key: Option<&[u8]>,
+        value: &[u8],
+    ) -> Result<SendResult, SqError> {
+        let results = self
+            .send_batch(vec![ProducerMessage {
+                topic: topic.to_string(),
+                key: key.map(|k| k.to_vec()),
+                value: value.to_vec(),
+                headers: Vec::new(),
+            }])
+            .await?;
+        Ok(results.into_iter().next().unwrap())
+    }
+
+    /// Send a batch of messages.
+    pub async fn send_batch(
+        &mut self,
+        messages: Vec<ProducerMessage>,
+    ) -> Result<Vec<SendResult>, SqError> {
+        // Build capnp request.
+        let mut builder = capnp::message::Builder::new_default();
+        {
+            let mut req = builder.init_root::<data_plane_capnp::publish_request::Builder>();
+            req.set_ack_mode(self.config.ack_mode.to_capnp_u8());
+            req.set_producer_id(&self.config.producer_id[..]);
+
+            let mut msg_list = req.init_messages(messages.len() as u32);
+            for (i, m) in messages.iter().enumerate() {
+                let mut entry = msg_list.reborrow().get(i as u32);
+                entry.set_topic(&m.topic[..]);
+                entry.set_key(m.key.as_deref().unwrap_or(&[]));
+                entry.set_value(&m.value);
+
+                let mut headers = entry.init_headers(m.headers.len() as u32);
+                for (j, (k, v)) in m.headers.iter().enumerate() {
+                    let mut h = headers.reborrow().get(j as u32);
+                    h.set_key(&k[..]);
+                    h.set_value(v);
+                }
+            }
+        }
+
+        let frame = codec::build_frame(OP_PUBLISH_REQ, &builder);
+        self.conn.send_frame(frame).await?;
+
+        // Read response.
+        let resp_frame = self.conn.recv_frame().await?;
+
+        if resp_frame.opcode == OP_ERROR {
+            let msg = decode_error(&resp_frame.payload)?;
+            return Err(SqError::Server(msg));
+        }
+
+        if resp_frame.opcode != OP_PUBLISH_RES {
+            return Err(SqError::Server(format!(
+                "unexpected opcode: 0x{:02x}",
+                resp_frame.opcode
+            )));
+        }
+
+        // Decode response.
+        let reader = codec::read_capnp(&resp_frame.payload)
+            .map_err(|e| SqError::Server(format!("decode error: {e}")))?;
+        let resp = reader
+            .get_root::<data_plane_capnp::publish_response::Reader>()
+            .map_err(|e| SqError::Server(format!("schema error: {e}")))?;
+
+        let results = resp
+            .get_results()
+            .map_err(|e| SqError::Server(format!("schema error: {e}")))?;
+
+        let mut send_results = Vec::with_capacity(results.len() as usize);
+        for i in 0..results.len() {
+            let r = results.get(i);
+            send_results.push(SendResult {
+                topic: r
+                    .get_topic()
+                    .map_err(|e| SqError::Server(format!("schema error: {e}")))?
+                    .to_string()
+                    .map_err(|e| SqError::Server(format!("utf8 error: {e}")))?,
+                partition: r.get_partition(),
+                offset: r.get_offset(),
+            });
+        }
+
+        Ok(send_results)
+    }
+}
+
+fn decode_error(payload: &[u8]) -> Result<String, SqError> {
+    let reader = codec::read_capnp(payload)
+        .map_err(|e| SqError::Server(format!("decode error: {e}")))?;
+    let err = reader
+        .get_root::<data_plane_capnp::error_response::Reader>()
+        .map_err(|e| SqError::Server(format!("schema error: {e}")))?;
+    err.get_message()
+        .map_err(|e| SqError::Server(format!("schema error: {e}")))?
+        .to_string()
+        .map_err(|e| SqError::Server(format!("utf8 error: {e}")))
+}