roto/roto-tonic/src/lib.rs

use bytes::{Buf, BufMut, Bytes, BytesMut};
use http_body::Body;
use roto_runtime::RotoMessage;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::Mutex;
use std::task::{Context, Poll};
use tonic::codec::{Codec, DecodeBuf, Decoder, EncodeBuf, Encoder};

pub mod generated {
    pub mod helloworld;
    pub mod interop;
}

pub struct RotoCodec<T, U> {
    _phantom: PhantomData<(T, U)>,
}

impl<T, U> Default for RotoCodec<T, U> {
    fn default() -> Self {
        Self {
            _phantom: PhantomData,
        }
    }
}

impl<T, U> Codec for RotoCodec<T, U>
where
    T: RotoMessage + Send + 'static,
    U: RotoMessage + Send + 'static,
{
    type Encode = U;
    type Decode = T;
    type Encoder = RotoEncoder<U>;
    type Decoder = RotoDecoder<T>;

    fn encoder(&mut self) -> Self::Encoder {
        RotoEncoder(PhantomData)
    }

    fn decoder(&mut self) -> Self::Decoder {
        RotoDecoder(PhantomData)
    }
}

pub struct RotoEncoder<U>(PhantomData<U>);

impl<U> Encoder for RotoEncoder<U>
where
    U: RotoMessage,
{
    type Item = U;
    type Error = tonic::Status;

    fn encode(&mut self, message: Self::Item, buf: &mut EncodeBuf<'_>) -> Result<(), Self::Error> {
        buf.put_slice(&message.bytes());
        Ok(())
    }
}

pub struct RotoDecoder<T>(PhantomData<T>);

impl<T> Decoder for RotoDecoder<T>
where
    T: RotoMessage,
{
    type Item = T;
    type Error = tonic::Status;

    fn decode(&mut self, buf: &mut DecodeBuf<'_>) -> Result<Option<Self::Item>, Self::Error> {
        if buf.remaining() == 0 {
            return Ok(None);
        }

        let bytes = buf.copy_to_bytes(buf.remaining());
        match T::decode(bytes) {
            Ok(msg) => Ok(Some(msg)),
            Err(e) => Err(tonic::Status::internal(format!("Roto decode error: {}", e))),
        }
    }
}

pub struct BufferPool {
    pool: Mutex<Vec<BytesMut>>,
    default_capacity: usize,
}

impl BufferPool {
    pub fn new(default_capacity: usize) -> Self {
        Self {
            pool: Mutex::new(Vec::new()),
            default_capacity,
        }
    }

    pub fn get(&self) -> BytesMut {
        self.pool
            .lock()
            .unwrap()
            .pop()
            .unwrap_or_else(|| BytesMut::with_capacity(self.default_capacity))
    }

    pub fn put(&self, mut buf: BytesMut) {
        buf.clear();
        if buf.capacity() >= self.default_capacity {
            self.pool.lock().unwrap().push(buf);
        }
    }
}

pub struct StatusBody {
    pub data: Option<Bytes>,
    pub trailers: Option<http::HeaderMap>,
}

impl StatusBody {
    pub fn new(data: Option<Bytes>, status: u8) -> Self {
        let mut trailers = http::HeaderMap::new();
        trailers.insert("grpc-status", status.to_string().parse().unwrap());
        Self {
            data,
            trailers: Some(trailers),
        }
    }
}

impl Body for StatusBody {
    type Data = Bytes;
    type Error = tonic::Status;

    fn poll_frame(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
        if let Some(data) = self.data.take() {
            Poll::Ready(Some(Ok(http_body::Frame::data(data))))
        } else if let Some(trailers) = self.trailers.take() {
            Poll::Ready(Some(Ok(http_body::Frame::trailers(trailers))))
        } else {
            Poll::Ready(None)
        }
    }
}