145 lines
3.3 KiB
Rust
145 lines
3.3 KiB
Rust
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)
|
|
}
|
|
}
|
|
}
|