Fix stuff

roto-tonic/src/lib.rs

@@ -1,7 +1,12 @@
 use std::marker::PhantomData;
 use tonic::codec::{Codec, Decoder, Encoder, DecodeBuf, EncodeBuf};
-use bytes::{Buf, BufMut};
+use bytes::{Buf, BufMut, Bytes, BytesMut};
 use roto_runtime::RotoMessage;
+use std::sync::{Arc, Mutex};
+use std::pin::Pin;
+use std::future::Future;
+use std::task::{Context, Poll};
+use http_body::Body;
 
 pub struct RotoCodec<T, U> {
     _phantom: PhantomData<(T, U)>,
@@ -70,3 +75,46 @@ where
         }
     }
 }
+
+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(crate) Option<Bytes>);
+
+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.0.take() {
+            Poll::Ready(Some(Ok(http_body::Frame::data(data))))
+        } else {
+            Poll::Ready(None)
+        }
+    }
+}