dynamo-runtime 1.4.0

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Event plane codec for serializing/deserializing envelopes and payloads.

use anyhow::{Result, anyhow, bail};
use bytes::Bytes;
use serde::{Deserialize, Serialize, de::DeserializeOwned};

use super::EventEnvelope;

/// Codec for serializing and deserializing event envelopes and payloads.
///
/// Currently only supports MessagePack for all transports.
#[derive(Debug, Clone, Copy)]
pub enum Codec {
    Msgpack(MsgpackCodec),
}

impl Default for Codec {
    fn default() -> Self {
        Codec::Msgpack(MsgpackCodec)
    }
}

impl Codec {
    /// Encode an EventEnvelope to wire bytes
    pub fn encode_envelope(&self, envelope: &EventEnvelope) -> Result<Bytes> {
        match self {
            Codec::Msgpack(c) => c.encode_envelope(envelope),
        }
    }

    /// Encode an event envelope while borrowing its immutable topic and payload.
    pub fn encode_envelope_parts(
        &self,
        publisher_id: u64,
        sequence: u64,
        published_at: u64,
        topic: &str,
        payload: &[u8],
    ) -> Result<Bytes> {
        match self {
            Codec::Msgpack(c) => {
                c.encode_envelope_parts(publisher_id, sequence, published_at, topic, payload)
            }
        }
    }

    /// Decode wire bytes to an EventEnvelope
    pub fn decode_envelope(&self, bytes: &Bytes) -> Result<EventEnvelope> {
        match self {
            Codec::Msgpack(c) => c.decode_envelope(bytes),
        }
    }

    /// Decode only the envelope identity without allocating its topic or payload.
    pub(crate) fn decode_envelope_identity(&self, bytes: &Bytes) -> Result<(u64, u64)> {
        match self {
            Codec::Msgpack(c) => c.decode_envelope_identity(bytes),
        }
    }

    /// Encode a typed payload to bytes (for embedding in envelope)
    pub fn encode_payload<T: Serialize>(&self, payload: &T) -> Result<Bytes> {
        match self {
            Codec::Msgpack(c) => c.encode_payload(payload),
        }
    }

    /// Decode payload bytes to a typed value
    pub fn decode_payload<T: DeserializeOwned>(&self, bytes: &Bytes) -> Result<T> {
        match self {
            Codec::Msgpack(c) => c.decode_payload(bytes),
        }
    }

    /// Codec name for debugging
    pub fn name(&self) -> &'static str {
        match self {
            Codec::Msgpack(c) => c.name(),
        }
    }
}

#[derive(Debug, Clone, Copy, Default)]
pub struct MsgpackCodec;

#[derive(Deserialize, Serialize)]
struct BorrowedEventEnvelope<'a> {
    publisher_id: u64,
    sequence: u64,
    published_at: u64,
    #[serde(borrow)]
    topic: &'a str,
    #[serde(borrow, serialize_with = "serialize_bytes")]
    payload: &'a [u8],
}

fn serialize_bytes<S>(bytes: &&[u8], serializer: S) -> std::result::Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    serializer.serialize_bytes(bytes)
}

impl MsgpackCodec {
    pub fn encode_envelope(&self, envelope: &EventEnvelope) -> Result<Bytes> {
        Ok(Bytes::from(rmp_serde::to_vec_named(envelope)?))
    }

    pub fn encode_envelope_parts(
        &self,
        publisher_id: u64,
        sequence: u64,
        published_at: u64,
        topic: &str,
        payload: &[u8],
    ) -> Result<Bytes> {
        let envelope = BorrowedEventEnvelope {
            publisher_id,
            sequence,
            published_at,
            topic,
            payload,
        };
        Ok(Bytes::from(rmp_serde::to_vec_named(&envelope)?))
    }

    pub fn decode_envelope(&self, bytes: &Bytes) -> Result<EventEnvelope> {
        let envelope: BorrowedEventEnvelope<'_> = rmp_serde::from_slice(bytes)?;
        let payload_start = (envelope.payload.as_ptr() as usize)
            .checked_sub(bytes.as_ptr() as usize)
            .ok_or_else(|| {
                anyhow!("MessagePack envelope payload is not borrowed from its frame")
            })?;
        let payload_end = payload_start
            .checked_add(envelope.payload.len())
            .ok_or_else(|| anyhow!("MessagePack envelope payload range overflowed"))?;

        if payload_end > bytes.len() {
            bail!("MessagePack envelope payload exceeds its frame");
        }

        Ok(EventEnvelope {
            publisher_id: envelope.publisher_id,
            sequence: envelope.sequence,
            published_at: envelope.published_at,
            topic: envelope.topic.to_owned(),
            payload: bytes.slice(payload_start..payload_end),
        })
    }

    pub(crate) fn decode_envelope_identity(&self, bytes: &Bytes) -> Result<(u64, u64)> {
        let envelope: BorrowedEventEnvelope<'_> = rmp_serde::from_slice(bytes)?;
        Ok((envelope.publisher_id, envelope.sequence))
    }

    pub fn encode_payload<T: Serialize>(&self, payload: &T) -> Result<Bytes> {
        Ok(Bytes::from(rmp_serde::to_vec_named(payload)?))
    }

    pub fn decode_payload<T: DeserializeOwned>(&self, bytes: &Bytes) -> Result<T> {
        Ok(rmp_serde::from_slice(bytes)?)
    }

    pub fn name(&self) -> &'static str {
        "msgpack"
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)]
    struct TestEvent {
        worker_id: u64,
        message: String,
    }

    #[test]
    fn test_msgpack_codec_envelope_roundtrip() {
        let codec = MsgpackCodec;

        let envelope = EventEnvelope {
            publisher_id: 12345,
            sequence: 42,
            published_at: 1700000000000,
            topic: "test-topic".to_string(),
            payload: Bytes::from("test payload"),
        };

        let encoded = codec.encode_envelope(&envelope).unwrap();
        let decoded = codec.decode_envelope(&encoded).unwrap();

        assert_eq!(decoded.publisher_id, envelope.publisher_id);
        assert_eq!(decoded.sequence, envelope.sequence);
        assert_eq!(decoded.published_at, envelope.published_at);
        assert_eq!(decoded.topic, envelope.topic);
        assert_eq!(decoded.payload, envelope.payload);

        let frame_start = encoded.as_ptr() as usize;
        let frame_end = frame_start + encoded.len();
        let payload_start = decoded.payload.as_ptr() as usize;
        let payload_end = payload_start + decoded.payload.len();
        assert!(
            (frame_start..=frame_end).contains(&payload_start) && payload_end <= frame_end,
            "decoded payload must share the encoded frame's backing storage"
        );
    }

    #[test]
    fn test_msgpack_codec_borrowed_envelope_roundtrip() {
        let codec = MsgpackCodec;
        let payload = b"borrowed payload";

        let borrowed = codec
            .encode_envelope_parts(12345, 42, 1700000000000, "test-topic", payload)
            .unwrap();
        let owned = codec
            .encode_envelope(&EventEnvelope {
                publisher_id: 12345,
                sequence: 42,
                published_at: 1700000000000,
                topic: "test-topic".to_string(),
                payload: Bytes::from_static(payload),
            })
            .unwrap();
        assert_eq!(borrowed, owned);

        let decoded = codec.decode_envelope(&borrowed).unwrap();

        assert_eq!(decoded.publisher_id, 12345);
        assert_eq!(decoded.sequence, 42);
        assert_eq!(decoded.published_at, 1700000000000);
        assert_eq!(decoded.topic, "test-topic");
        assert_eq!(decoded.payload.as_ref(), payload);
    }

    #[test]
    fn test_msgpack_codec_envelope_identity_decode() {
        let codec = MsgpackCodec;
        let encoded = codec
            .encode_envelope_parts(12345, 42, 1700000000000, "test-topic", b"payload")
            .unwrap();

        assert_eq!(
            codec.decode_envelope_identity(&encoded).unwrap(),
            (12345, 42)
        );
    }

    #[test]
    fn test_msgpack_codec_rejects_malformed_envelope() {
        let codec = MsgpackCodec;

        assert!(codec.decode_envelope(&Bytes::from_static(&[0xc1])).is_err());
        assert!(
            codec
                .decode_envelope_identity(&Bytes::from_static(&[0xc1]))
                .is_err()
        );
    }

    #[test]
    fn test_msgpack_codec_payload_roundtrip() {
        let codec = MsgpackCodec;

        let event = TestEvent {
            worker_id: 123,
            message: "hello world".to_string(),
        };

        let encoded = codec.encode_payload(&event).unwrap();
        let decoded: TestEvent = codec.decode_payload(&encoded).unwrap();

        assert_eq!(decoded, event);
    }
}