sq/crates/sq-sdk/src/capnp_producer.rs

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}")))
}