refactor content encoder

src/client/writer.rs

@@ -21,7 +21,8 @@ use tokio_io::AsyncWrite;
 
 use body::{Binary, Body};
 use header::ContentEncoding;
-use server::encoding::{ContentEncoder, TransferEncoding};
+use server::encoding::{ContentEncoder, Output, TransferEncoding};
+use server::shared::SharedBytes;
 use server::WriterState;
 
 use client::ClientRequest;
@@ -41,21 +42,18 @@ pub(crate) struct HttpClientWriter {
     flags: Flags,
     written: u64,
     headers_size: u32,
-    buffer: Box<BytesMut>,
+    buffer: Output,
     buffer_capacity: usize,
-    encoder: ContentEncoder,
 }
 
 impl HttpClientWriter {
     pub fn new() -> HttpClientWriter {
-        let encoder = ContentEncoder::Identity(TransferEncoding::eof());
         HttpClientWriter {
             flags: Flags::empty(),
             written: 0,
             headers_size: 0,
             buffer_capacity: 0,
-            buffer: Box::new(BytesMut::new()),
-            encoder,
+            buffer: Output::Buffer(SharedBytes::empty()),
         }
     }
 
@@ -75,7 +73,7 @@ impl HttpClientWriter {
         &mut self, stream: &mut T,
     ) -> io::Result<WriterState> {
         while !self.buffer.is_empty() {
-            match stream.write(self.buffer.as_ref()) {
+            match stream.write(self.buffer.as_ref().as_ref()) {
                 Ok(0) => {
                     self.disconnected();
                     return Ok(WriterState::Done);
@@ -113,16 +111,18 @@ impl HttpClientWriter {
     pub fn start(&mut self, msg: &mut ClientRequest) -> io::Result<()> {
         // prepare task
         self.flags.insert(Flags::STARTED);
-        self.encoder = content_encoder(self.buffer.as_mut(), msg);
         if msg.upgrade() {
             self.flags.insert(Flags::UPGRADE);
         }
 
         // render message
         {
+            // output buffer
+            let buffer = self.buffer.get_mut();
+
             // status line
             writeln!(
-                Writer(&mut self.buffer),
+                Writer(buffer),
                 "{} {} {:?}\r",
                 msg.method(),
                 msg.uri()
@@ -134,53 +134,49 @@ impl HttpClientWriter {
 
             // write headers
             if let Body::Binary(ref bytes) = *msg.body() {
-                self.buffer
-                    .reserve(msg.headers().len() * AVERAGE_HEADER_SIZE + bytes.len());
+                buffer.reserve(msg.headers().len() * AVERAGE_HEADER_SIZE + bytes.len());
             } else {
-                self.buffer
-                    .reserve(msg.headers().len() * AVERAGE_HEADER_SIZE);
+                buffer.reserve(msg.headers().len() * AVERAGE_HEADER_SIZE);
             }
 
             for (key, value) in msg.headers() {
                 let v = value.as_ref();
                 let k = key.as_str().as_bytes();
-                self.buffer.reserve(k.len() + v.len() + 4);
-                self.buffer.put_slice(k);
-                self.buffer.put_slice(b": ");
-                self.buffer.put_slice(v);
-                self.buffer.put_slice(b"\r\n");
+                buffer.reserve(k.len() + v.len() + 4);
+                buffer.put_slice(k);
+                buffer.put_slice(b": ");
+                buffer.put_slice(v);
+                buffer.put_slice(b"\r\n");
             }
 
             // set date header
             if !msg.headers().contains_key(DATE) {
-                self.buffer.extend_from_slice(b"date: ");
-                set_date(&mut self.buffer);
-                self.buffer.extend_from_slice(b"\r\n\r\n");
+                buffer.extend_from_slice(b"date: ");
+                set_date(buffer);
+                buffer.extend_from_slice(b"\r\n\r\n");
             } else {
-                self.buffer.extend_from_slice(b"\r\n");
+                buffer.extend_from_slice(b"\r\n");
+            }
         }
         self.headers_size = self.buffer.len() as u32;
 
+        self.buffer = content_encoder(self.buffer.take(), msg);
+
         if msg.body().is_binary() {
             if let Body::Binary(bytes) = msg.replace_body(Body::Empty) {
                 self.written += bytes.len() as u64;
-                    self.encoder.write(bytes.as_ref())?;
+                self.buffer.write(bytes.as_ref())?;
             }
         } else {
             self.buffer_capacity = msg.write_buffer_capacity();
         }
-        }
         Ok(())
     }
 
     pub fn write(&mut self, payload: &[u8]) -> io::Result<WriterState> {
         self.written += payload.len() as u64;
         if !self.flags.contains(Flags::DISCONNECTED) {
-            if self.flags.contains(Flags::UPGRADE) {
-                self.buffer.extend(payload);
-            } else {
-                self.encoder.write(payload)?;
-            }
+            self.buffer.write(payload)?;
         }
 
         if self.buffer.len() > self.buffer_capacity {
@@ -191,9 +187,7 @@ impl HttpClientWriter {
     }
 
     pub fn write_eof(&mut self) -> io::Result<()> {
-        self.encoder.write_eof()?;
-
-        if self.encoder.is_eof() {
+        if self.buffer.write_eof()? {
             Ok(())
         } else {
             Err(io::Error::new(
@@ -221,21 +215,20 @@ impl HttpClientWriter {
     }
 }
 
-fn content_encoder(buf: &mut BytesMut, req: &mut ClientRequest) -> ContentEncoder {
+fn content_encoder(buf: SharedBytes, req: &mut ClientRequest) -> Output {
     let version = req.version();
     let mut body = req.replace_body(Body::Empty);
     let mut encoding = req.content_encoding();
 
-    let mut transfer = match body {
+    let transfer = match body {
         Body::Empty => {
             req.headers_mut().remove(CONTENT_LENGTH);
-            TransferEncoding::length(0)
+            TransferEncoding::length(0, buf)
         }
         Body::Binary(ref mut bytes) => {
             if encoding.is_compression() {
-                let mut tmp = BytesMut::new();
-                let mut transfer = TransferEncoding::eof();
-                transfer.set_buffer(&mut tmp);
+                let mut tmp = SharedBytes::empty();
+                let mut transfer = TransferEncoding::eof(tmp);
                 let mut enc = match encoding {
                     #[cfg(feature = "flate2")]
                     ContentEncoding::Deflate => ContentEncoder::Deflate(
@@ -256,7 +249,7 @@ fn content_encoder(buf: &mut BytesMut, req: &mut ClientRequest) -> ContentEncode
                 // TODO return error!
                 let _ = enc.write(bytes.as_ref());
                 let _ = enc.write_eof();
-                *bytes = Binary::from(tmp.take());
+                *bytes = Binary::from(enc.buf_mut().take());
 
                 req.headers_mut().insert(
                     CONTENT_ENCODING,
@@ -268,7 +261,7 @@ fn content_encoder(buf: &mut BytesMut, req: &mut ClientRequest) -> ContentEncode
             let _ = write!(b, "{}", bytes.len());
             req.headers_mut()
                 .insert(CONTENT_LENGTH, HeaderValue::try_from(b.freeze()).unwrap());
-            TransferEncoding::eof()
+            TransferEncoding::eof(buf)
         }
         Body::Streaming(_) | Body::Actor(_) => {
             if req.upgrade() {
@@ -282,9 +275,9 @@ fn content_encoder(buf: &mut BytesMut, req: &mut ClientRequest) -> ContentEncode
                     encoding = ContentEncoding::Identity;
                     req.headers_mut().remove(CONTENT_ENCODING);
                 }
-                TransferEncoding::eof()
+                TransferEncoding::eof(buf)
             } else {
-                streaming_encoding(version, req)
+                streaming_encoding(buf, version, req)
             }
         }
     };
@@ -295,10 +288,9 @@ fn content_encoder(buf: &mut BytesMut, req: &mut ClientRequest) -> ContentEncode
             HeaderValue::from_static(encoding.as_str()),
         );
     }
-    transfer.set_buffer(buf);
 
     req.replace_body(body);
-    match encoding {
+    let enc = match encoding {
         #[cfg(feature = "flate2")]
         ContentEncoding::Deflate => ContentEncoder::Deflate(DeflateEncoder::new(
             transfer,
@@ -310,23 +302,24 @@ fn content_encoder(buf: &mut BytesMut, req: &mut ClientRequest) -> ContentEncode
         }
         #[cfg(feature = "brotli")]
         ContentEncoding::Br => ContentEncoder::Br(BrotliEncoder::new(transfer, 5)),
-        ContentEncoding::Identity | ContentEncoding::Auto => {
-            ContentEncoder::Identity(transfer)
-        }
-    }
+        ContentEncoding::Identity | ContentEncoding::Auto => return Output::TE(transfer),
+    };
+    Output::Encoder(enc)
 }
 
-fn streaming_encoding(version: Version, req: &mut ClientRequest) -> TransferEncoding {
+fn streaming_encoding(
+    buf: SharedBytes, version: Version, req: &mut ClientRequest,
+) -> TransferEncoding {
     if req.chunked() {
         // Enable transfer encoding
         req.headers_mut().remove(CONTENT_LENGTH);
         if version == Version::HTTP_2 {
             req.headers_mut().remove(TRANSFER_ENCODING);
-            TransferEncoding::eof()
+            TransferEncoding::eof(buf)
         } else {
             req.headers_mut()
                 .insert(TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
-            TransferEncoding::chunked()
+            TransferEncoding::chunked(buf)
         }
     } else {
         // if Content-Length is specified, then use it as length hint
@@ -349,9 +342,9 @@ fn streaming_encoding(version: Version, req: &mut ClientRequest) -> TransferEnco
 
         if !chunked {
             if let Some(len) = len {
-                TransferEncoding::length(len)
+                TransferEncoding::length(len, buf)
             } else {
-                TransferEncoding::eof()
+                TransferEncoding::eof(buf)
             }
         } else {
             // Enable transfer encoding
@@ -359,11 +352,11 @@ fn streaming_encoding(version: Version, req: &mut ClientRequest) -> TransferEnco
                 Version::HTTP_11 => {
                     req.headers_mut()
                         .insert(TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
-                    TransferEncoding::chunked()
+                    TransferEncoding::chunked(buf)
                 }
                 _ => {
                     req.headers_mut().remove(TRANSFER_ENCODING);
-                    TransferEncoding::eof()
+                    TransferEncoding::eof(buf)
                 }
             }
         }

src/fs.rs

@@ -1127,6 +1127,7 @@ mod tests {
 
         let response = srv.execute(request.send()).unwrap();
 
+        println!("RESP: {:?}", response);
         let te = response
             .headers()
             .get(header::TRANSFER_ENCODING)

src/server/encoding.rs

@@ -1,7 +1,7 @@
 use std::fmt::Write as FmtWrite;
 use std::io::{Read, Write};
 use std::str::FromStr;
-use std::{cmp, io, mem, ptr};
+use std::{cmp, io, mem};
 
 #[cfg(feature = "brotli")]
 use brotli2::write::{BrotliDecoder, BrotliEncoder};
@@ -25,6 +25,8 @@ use httprequest::HttpInnerMessage;
 use httpresponse::HttpResponse;
 use payload::{PayloadSender, PayloadStatus, PayloadWriter};
 
+use super::shared::SharedBytes;
+
 pub(crate) enum PayloadType {
     Sender(PayloadSender),
     Encoding(Box<EncodedPayload>),
@@ -368,6 +370,83 @@ impl PayloadStream {
     }
 }
 
+pub(crate) enum Output {
+    Buffer(SharedBytes),
+    Encoder(ContentEncoder),
+    TE(TransferEncoding),
+    Empty,
+}
+
+impl Output {
+    pub fn take(&mut self) -> SharedBytes {
+        match mem::replace(self, Output::Empty) {
+            Output::Buffer(bytes) => bytes,
+            _ => panic!(),
+        }
+    }
+    pub fn as_ref(&mut self) -> &SharedBytes {
+        match self {
+            Output::Buffer(ref mut bytes) => bytes,
+            Output::Encoder(ref mut enc) => enc.buf_ref(),
+            Output::TE(ref mut te) => te.buf_ref(),
+            Output::Empty => panic!(),
+        }
+    }
+    pub fn get_mut(&mut self) -> &mut BytesMut {
+        match self {
+            Output::Buffer(ref mut bytes) => bytes.get_mut(),
+            _ => panic!(),
+        }
+    }
+    pub fn split_to(&mut self, cap: usize) -> BytesMut {
+        match self {
+            Output::Buffer(ref mut bytes) => bytes.split_to(cap),
+            Output::Encoder(ref mut enc) => enc.buf_mut().split_to(cap),
+            Output::TE(ref mut te) => te.buf_mut().split_to(cap),
+            Output::Empty => BytesMut::new(),
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        match self {
+            Output::Buffer(ref bytes) => bytes.len(),
+            Output::Encoder(ref enc) => enc.len(),
+            Output::TE(ref te) => te.len(),
+            Output::Empty => 0,
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        match self {
+            Output::Buffer(ref bytes) => bytes.is_empty(),
+            Output::Encoder(ref enc) => enc.is_empty(),
+            Output::TE(ref te) => te.is_empty(),
+            Output::Empty => true,
+        }
+    }
+
+    pub fn write(&mut self, data: &[u8]) -> Result<(), io::Error> {
+        match self {
+            Output::Buffer(ref mut bytes) => {
+                bytes.extend_from_slice(data);
+                Ok(())
+            }
+            Output::Encoder(ref mut enc) => enc.write(data),
+            Output::TE(ref mut te) => te.encode(data).map(|_| ()),
+            Output::Empty => Ok(()),
+        }
+    }
+
+    pub fn write_eof(&mut self) -> Result<bool, io::Error> {
+        match self {
+            Output::Buffer(_) => Ok(true),
+            Output::Encoder(ref mut enc) => enc.write_eof(),
+            Output::TE(ref mut te) => Ok(te.encode_eof()),
+            Output::Empty => Ok(true),
+        }
+    }
+}
+
 pub(crate) enum ContentEncoder {
     #[cfg(feature = "flate2")]
     Deflate(DeflateEncoder<TransferEncoding>),
@@ -379,14 +458,10 @@ pub(crate) enum ContentEncoder {
 }
 
 impl ContentEncoder {
-    pub fn empty() -> ContentEncoder {
-        ContentEncoder::Identity(TransferEncoding::eof())
-    }
-
     pub fn for_server(
-        buf: &mut BytesMut, req: &HttpInnerMessage, resp: &mut HttpResponse,
+        buf: SharedBytes, req: &HttpInnerMessage, resp: &mut HttpResponse,
         response_encoding: ContentEncoding,
-    ) -> ContentEncoder {
+    ) -> Output {
         let version = resp.version().unwrap_or_else(|| req.version);
         let is_head = req.method == Method::HEAD;
         let mut len = 0;
@@ -439,7 +514,7 @@ impl ContentEncoder {
                 if req.method != Method::HEAD {
                     resp.headers_mut().remove(CONTENT_LENGTH);
                 }
-                TransferEncoding::length(0)
+                TransferEncoding::length(0, buf)
             }
             &Body::Binary(_) => {
                 #[cfg(any(feature = "brotli", feature = "flate2"))]
@@ -447,9 +522,8 @@ impl ContentEncoder {
                     if !(encoding == ContentEncoding::Identity
                         || encoding == ContentEncoding::Auto)
                     {
-                        let mut tmp = BytesMut::default();
-                        let mut transfer = TransferEncoding::eof();
-                        transfer.set_buffer(&mut tmp);
+                        let mut tmp = SharedBytes::empty();
+                        let mut transfer = TransferEncoding::eof(tmp);
                         let mut enc = match encoding {
                             #[cfg(feature = "flate2")]
                             ContentEncoding::Deflate => ContentEncoder::Deflate(
@@ -473,7 +547,7 @@ impl ContentEncoder {
                         // TODO return error!
                         let _ = enc.write(bin.as_ref());
                         let _ = enc.write_eof();
-                        let body = tmp.take();
+                        let body = enc.buf_mut().take();
                         len = body.len();
 
                         encoding = ContentEncoding::Identity;
@@ -491,7 +565,7 @@ impl ContentEncoder {
                 } else {
                     // resp.headers_mut().remove(CONTENT_LENGTH);
                 }
-                TransferEncoding::eof()
+                TransferEncoding::eof(buf)
             }
             &Body::Streaming(_) | &Body::Actor(_) => {
                 if resp.upgrade() {
@@ -502,14 +576,14 @@ impl ContentEncoder {
                         encoding = ContentEncoding::Identity;
                         resp.headers_mut().remove(CONTENT_ENCODING);
                     }
-                    TransferEncoding::eof()
+                    TransferEncoding::eof(buf)
                 } else {
                     if !(encoding == ContentEncoding::Identity
                         || encoding == ContentEncoding::Auto)
                     {
                         resp.headers_mut().remove(CONTENT_LENGTH);
                     }
-                    ContentEncoder::streaming_encoding(version, resp)
+                    ContentEncoder::streaming_encoding(buf, version, resp)
                 }
             }
         };
@@ -518,9 +592,8 @@ impl ContentEncoder {
             resp.set_body(Body::Empty);
             transfer.kind = TransferEncodingKind::Length(0);
         }
-        transfer.set_buffer(buf);
 
-        match encoding {
+        let enc = match encoding {
             #[cfg(feature = "flate2")]
             ContentEncoding::Deflate => ContentEncoder::Deflate(DeflateEncoder::new(
                 transfer,
@@ -533,13 +606,14 @@ impl ContentEncoder {
             #[cfg(feature = "brotli")]
             ContentEncoding::Br => ContentEncoder::Br(BrotliEncoder::new(transfer, 3)),
             ContentEncoding::Identity | ContentEncoding::Auto => {
-                ContentEncoder::Identity(transfer)
-            }
+                return Output::TE(transfer)
             }
+        };
+        Output::Encoder(enc)
     }
 
     fn streaming_encoding(
-        version: Version, resp: &mut HttpResponse,
+        buf: SharedBytes, version: Version, resp: &mut HttpResponse,
     ) -> TransferEncoding {
         match resp.chunked() {
             Some(true) => {
@@ -547,14 +621,14 @@ impl ContentEncoder {
                 resp.headers_mut().remove(CONTENT_LENGTH);
                 if version == Version::HTTP_2 {
                     resp.headers_mut().remove(TRANSFER_ENCODING);
-                    TransferEncoding::eof()
+                    TransferEncoding::eof(buf)
                 } else {
                     resp.headers_mut()
                         .insert(TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
-                    TransferEncoding::chunked()
+                    TransferEncoding::chunked(buf)
                 }
             }
-            Some(false) => TransferEncoding::eof(),
+            Some(false) => TransferEncoding::eof(buf),
             None => {
                 // if Content-Length is specified, then use it as length hint
                 let (len, chunked) =
@@ -577,9 +651,9 @@ impl ContentEncoder {
 
                 if !chunked {
                     if let Some(len) = len {
-                        TransferEncoding::length(len)
+                        TransferEncoding::length(len, buf)
                     } else {
-                        TransferEncoding::eof()
+                        TransferEncoding::eof(buf)
                     }
                 } else {
                     // Enable transfer encoding
@@ -589,11 +663,11 @@ impl ContentEncoder {
                                 TRANSFER_ENCODING,
                                 HeaderValue::from_static("chunked"),
                             );
-                            TransferEncoding::chunked()
+                            TransferEncoding::chunked(buf)
                         }
                         _ => {
                             resp.headers_mut().remove(TRANSFER_ENCODING);
-                            TransferEncoding::eof()
+                            TransferEncoding::eof(buf)
                         }
                     }
                 }
@@ -604,15 +678,54 @@ impl ContentEncoder {
 
 impl ContentEncoder {
     #[inline]
-    pub fn is_eof(&self) -> bool {
+    pub fn len(&self) -> usize {
         match *self {
             #[cfg(feature = "brotli")]
-            ContentEncoder::Br(ref encoder) => encoder.get_ref().is_eof(),
+            ContentEncoder::Br(ref encoder) => encoder.get_ref().len(),
             #[cfg(feature = "flate2")]
-            ContentEncoder::Deflate(ref encoder) => encoder.get_ref().is_eof(),
+            ContentEncoder::Deflate(ref encoder) => encoder.get_ref().len(),
             #[cfg(feature = "flate2")]
-            ContentEncoder::Gzip(ref encoder) => encoder.get_ref().is_eof(),
-            ContentEncoder::Identity(ref encoder) => encoder.is_eof(),
+            ContentEncoder::Gzip(ref encoder) => encoder.get_ref().len(),
+            ContentEncoder::Identity(ref encoder) => encoder.len(),
+        }
+    }
+
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        match *self {
+            #[cfg(feature = "brotli")]
+            ContentEncoder::Br(ref encoder) => encoder.get_ref().is_empty(),
+            #[cfg(feature = "flate2")]
+            ContentEncoder::Deflate(ref encoder) => encoder.get_ref().is_empty(),
+            #[cfg(feature = "flate2")]
+            ContentEncoder::Gzip(ref encoder) => encoder.get_ref().is_empty(),
+            ContentEncoder::Identity(ref encoder) => encoder.is_empty(),
+        }
+    }
+
+    #[inline]
+    pub(crate) fn buf_mut(&mut self) -> &mut BytesMut {
+        match *self {
+            #[cfg(feature = "brotli")]
+            ContentEncoder::Br(ref mut encoder) => encoder.get_mut().buf_mut(),
+            #[cfg(feature = "flate2")]
+            ContentEncoder::Deflate(ref mut encoder) => encoder.get_mut().buf_mut(),
+            #[cfg(feature = "flate2")]
+            ContentEncoder::Gzip(ref mut encoder) => encoder.get_mut().buf_mut(),
+            ContentEncoder::Identity(ref mut encoder) => encoder.buf_mut(),
+        }
+    }
+
+    #[inline]
+    pub(crate) fn buf_ref(&mut self) -> &SharedBytes {
+        match *self {
+            #[cfg(feature = "brotli")]
+            ContentEncoder::Br(ref mut encoder) => encoder.get_mut().buf_ref(),
+            #[cfg(feature = "flate2")]
+            ContentEncoder::Deflate(ref mut encoder) => encoder.get_mut().buf_ref(),
+            #[cfg(feature = "flate2")]
+            ContentEncoder::Gzip(ref mut encoder) => encoder.get_mut().buf_ref(),
+            ContentEncoder::Identity(ref mut encoder) => encoder.buf_ref(),
         }
     }
 
@@ -620,7 +733,7 @@ impl ContentEncoder {
     #[inline(always)]
     pub fn write_eof(&mut self) -> Result<bool, io::Error> {
         let encoder =
-            mem::replace(self, ContentEncoder::Identity(TransferEncoding::eof()));
+            mem::replace(self, ContentEncoder::Identity(TransferEncoding::empty()));
 
         match encoder {
             #[cfg(feature = "brotli")]
@@ -695,9 +808,9 @@ impl ContentEncoder {
 }
 
 /// Encoders to handle different Transfer-Encodings.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub(crate) struct TransferEncoding {
-    buf: *mut BytesMut,
+    buf: Option<SharedBytes>,
     kind: TransferEncodingKind,
 }
 
@@ -716,40 +829,55 @@ enum TransferEncodingKind {
 }
 
 impl TransferEncoding {
-    pub(crate) fn set_buffer(&mut self, buf: *mut BytesMut) {
-        self.buf = buf;
+    fn take(self) -> SharedBytes {
+        self.buf.unwrap()
+    }
+
+    fn buf_ref(&mut self) -> &SharedBytes {
+        self.buf.as_ref().unwrap()
+    }
+
+    fn len(&self) -> usize {
+        self.buf.as_ref().unwrap().len()
+    }
+
+    fn is_empty(&self) -> bool {
+        self.buf.as_ref().unwrap().is_empty()
+    }
+
+    fn buf_mut(&mut self) -> &mut BytesMut {
+        self.buf.as_mut().unwrap().get_mut()
     }
 
     #[inline]
-    pub fn eof() -> TransferEncoding {
+    pub fn empty() -> TransferEncoding {
         TransferEncoding {
+            buf: None,
             kind: TransferEncodingKind::Eof,
-            buf: ptr::null_mut(),
         }
     }
 
     #[inline]
-    pub fn chunked() -> TransferEncoding {
+    pub fn eof(buf: SharedBytes) -> TransferEncoding {
         TransferEncoding {
+            buf: Some(buf),
+            kind: TransferEncodingKind::Eof,
+        }
+    }
+
+    #[inline]
+    pub fn chunked(buf: SharedBytes) -> TransferEncoding {
+        TransferEncoding {
+            buf: Some(buf),
             kind: TransferEncodingKind::Chunked(false),
-            buf: ptr::null_mut(),
         }
     }
 
     #[inline]
-    pub fn length(len: u64) -> TransferEncoding {
+    pub fn length(len: u64, buf: SharedBytes) -> TransferEncoding {
         TransferEncoding {
+            buf: Some(buf),
             kind: TransferEncodingKind::Length(len),
-            buf: ptr::null_mut(),
-        }
-    }
-
-    #[inline]
-    pub fn is_eof(&self) -> bool {
-        match self.kind {
-            TransferEncodingKind::Eof => true,
-            TransferEncodingKind::Chunked(ref eof) => *eof,
-            TransferEncodingKind::Length(ref remaining) => *remaining == 0,
         }
     }
 
@@ -759,9 +887,7 @@ impl TransferEncoding {
         match self.kind {
             TransferEncodingKind::Eof => {
                 let eof = msg.is_empty();
-                debug_assert!(!self.buf.is_null());
-                let buf = unsafe { &mut *self.buf };
-                buf.extend(msg);
+                self.buf.as_mut().unwrap().extend_from_slice(msg);
                 Ok(eof)
             }
             TransferEncodingKind::Chunked(ref mut eof) => {
@@ -771,16 +897,13 @@ impl TransferEncoding {
 
                 if msg.is_empty() {
                     *eof = true;
-                    debug_assert!(!self.buf.is_null());
-                    let buf = unsafe { &mut *self.buf };
-                    buf.extend_from_slice(b"0\r\n\r\n");
+                    self.buf.as_mut().unwrap().extend_from_slice(b"0\r\n\r\n");
                 } else {
                     let mut buf = BytesMut::new();
                     writeln!(&mut buf, "{:X}\r", msg.len())
                         .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
 
-                    debug_assert!(!self.buf.is_null());
-                    let b = unsafe { &mut *self.buf };
+                    let b = self.buf.as_mut().unwrap();
                     b.reserve(buf.len() + msg.len() + 2);
                     b.extend_from_slice(buf.as_ref());
                     b.extend_from_slice(msg);
@@ -795,8 +918,10 @@ impl TransferEncoding {
                     }
                     let len = cmp::min(*remaining, msg.len() as u64);
 
-                    debug_assert!(!self.buf.is_null());
-                    unsafe { &mut *self.buf }.extend(&msg[..len as usize]);
+                    self.buf
+                        .as_mut()
+                        .unwrap()
+                        .extend_from_slice(&msg[..len as usize]);
 
                     *remaining -= len as u64;
                     Ok(*remaining == 0)
@@ -816,10 +941,7 @@ impl TransferEncoding {
             TransferEncodingKind::Chunked(ref mut eof) => {
                 if !*eof {
                     *eof = true;
-
-                    debug_assert!(!self.buf.is_null());
-                    let buf = unsafe { &mut *self.buf };
-                    buf.extend_from_slice(b"0\r\n\r\n");
+                    self.buf.as_mut().unwrap().extend_from_slice(b"0\r\n\r\n");
                 }
                 true
             }
@@ -912,15 +1034,14 @@ mod tests {
 
     #[test]
     fn test_chunked_te() {
-        let mut bytes = BytesMut::new();
-        let mut enc = TransferEncoding::chunked();
+        let bytes = SharedBytes::empty();
+        let mut enc = TransferEncoding::chunked(bytes);
         {
-            enc.set_buffer(&mut bytes);
             assert!(!enc.encode(b"test").ok().unwrap());
             assert!(enc.encode(b"").ok().unwrap());
         }
         assert_eq!(
-            bytes.take().freeze(),
+            enc.buf_mut().take().freeze(),
             Bytes::from_static(b"4\r\ntest\r\n0\r\n\r\n")
         );
     }

src/server/h1writer.rs

@@ -6,7 +6,7 @@ use std::io;
 use std::rc::Rc;
 use tokio_io::AsyncWrite;
 
-use super::encoding::ContentEncoder;
+use super::encoding::{ContentEncoder, Output};
 use super::helpers;
 use super::settings::WorkerSettings;
 use super::shared::SharedBytes;
@@ -32,10 +32,9 @@ bitflags! {
 pub(crate) struct H1Writer<T: AsyncWrite, H: 'static> {
     flags: Flags,
     stream: T,
-    encoder: ContentEncoder,
     written: u64,
     headers_size: u32,
-    buffer: SharedBytes,
+    buffer: Output,
     buffer_capacity: usize,
     settings: Rc<WorkerSettings<H>>,
 }
@@ -46,10 +45,9 @@ impl<T: AsyncWrite, H: 'static> H1Writer<T, H> {
     ) -> H1Writer<T, H> {
         H1Writer {
             flags: Flags::KEEPALIVE,
-            encoder: ContentEncoder::empty(),
             written: 0,
             headers_size: 0,
-            buffer: buf,
+            buffer: Output::Buffer(buf),
             buffer_capacity: 0,
             stream,
             settings,
@@ -66,7 +64,7 @@ impl<T: AsyncWrite, H: 'static> H1Writer<T, H> {
     }
 
     pub fn disconnected(&mut self) {
-        self.buffer.take();
+        self.buffer = Output::Empty;
     }
 
     pub fn keepalive(&self) -> bool {
@@ -106,7 +104,8 @@ impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
 
     #[inline]
     fn buffer(&mut self) -> &mut BytesMut {
-        self.buffer.get_mut()
+        //self.buffer.get_mut()
+        unimplemented!()
     }
 
     fn start(
@@ -114,8 +113,6 @@ impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
         encoding: ContentEncoding,
     ) -> io::Result<WriterState> {
         // prepare task
-        self.encoder =
-            ContentEncoder::for_server(self.buffer.get_mut(), req, msg, encoding);
         if msg.keep_alive().unwrap_or_else(|| req.keep_alive()) {
             self.flags = Flags::STARTED | Flags::KEEPALIVE;
         } else {
@@ -143,7 +140,9 @@ impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
 
         // render message
         {
+            // output buffer
             let mut buffer = self.buffer.get_mut();
+
             let reason = msg.reason().as_bytes();
             let mut is_bin = if let Body::Binary(ref bytes) = body {
                 buffer.reserve(
@@ -222,9 +221,12 @@ impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
             self.headers_size = buffer.len() as u32;
         }
 
+        // output encoding
+        self.buffer = ContentEncoder::for_server(self.buffer.take(), req, msg, encoding);
+
         if let Body::Binary(bytes) = body {
             self.written = bytes.len() as u64;
-            self.encoder.write(bytes.as_ref())?;
+            self.buffer.write(bytes.as_ref())?;
         } else {
             // capacity, makes sense only for streaming or actor
             self.buffer_capacity = msg.write_buffer_capacity();
@@ -253,19 +255,19 @@ impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
                             Ok(val) => val,
                         };
                         if n < pl.len() {
-                            self.buffer.extend_from_slice(&pl[n..]);
+                            self.buffer.write(&pl[n..]);
                             return Ok(WriterState::Done);
                         }
                     } else {
-                        self.buffer.extend(payload);
+                        self.buffer.write(payload.as_ref());
                     }
                 } else {
                     // TODO: add warning, write after EOF
-                    self.encoder.write(payload.as_ref())?;
+                    self.buffer.write(payload.as_ref())?;
                 }
             } else {
                 // could be response to EXCEPT header
-                self.buffer.extend_from_slice(payload.as_ref())
+                self.buffer.write(payload.as_ref());
             }
         }
 
@@ -277,7 +279,7 @@ impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
     }
 
     fn write_eof(&mut self) -> io::Result<WriterState> {
-        if !self.encoder.write_eof()? {
+        if !self.buffer.write_eof()? {
             Err(io::Error::new(
                 io::ErrorKind::Other,
                 "Last payload item, but eof is not reached",
@@ -293,7 +295,7 @@ impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
     fn poll_completed(&mut self, shutdown: bool) -> Poll<(), io::Error> {
         if !self.buffer.is_empty() {
             let written = {
-                match Self::write_data(&mut self.stream, self.buffer.as_ref()) {
+                match Self::write_data(&mut self.stream, self.buffer.as_ref().as_ref()) {
                     Err(err) => {
                         if err.kind() == io::ErrorKind::WriteZero {
                             self.disconnected();

src/server/h2writer.rs

@@ -11,7 +11,7 @@ use std::{cmp, io};
 use http::header::{HeaderValue, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING};
 use http::{HttpTryFrom, Version};
 
-use super::encoding::ContentEncoder;
+use super::encoding::{ContentEncoder, Output};
 use super::helpers;
 use super::settings::WorkerSettings;
 use super::shared::SharedBytes;
@@ -35,10 +35,9 @@ bitflags! {
 pub(crate) struct H2Writer<H: 'static> {
     respond: SendResponse<Bytes>,
     stream: Option<SendStream<Bytes>>,
-    encoder: ContentEncoder,
     flags: Flags,
     written: u64,
-    buffer: SharedBytes,
+    buffer: Output,
     buffer_capacity: usize,
     settings: Rc<WorkerSettings<H>>,
 }
@@ -51,10 +50,9 @@ impl<H: 'static> H2Writer<H> {
             respond,
             settings,
             stream: None,
-            encoder: ContentEncoder::empty(),
             flags: Flags::empty(),
             written: 0,
-            buffer: buf,
+            buffer: Output::Buffer(buf),
             buffer_capacity: 0,
         }
     }
@@ -87,8 +85,7 @@ impl<H: 'static> Writer for H2Writer<H> {
     ) -> io::Result<WriterState> {
         // prepare response
         self.flags.insert(Flags::STARTED);
-        self.encoder =
-            ContentEncoder::for_server(self.buffer.get_mut(), req, msg, encoding);
+        self.buffer = ContentEncoder::for_server(self.buffer.take(), req, msg, encoding);
 
         // http2 specific
         msg.headers_mut().remove(CONNECTION);
@@ -150,7 +147,7 @@ impl<H: 'static> Writer for H2Writer<H> {
             } else {
                 self.flags.insert(Flags::EOF);
                 self.written = bytes.len() as u64;
-                self.encoder.write(bytes.as_ref())?;
+                self.buffer.write(bytes.as_ref())?;
                 if let Some(ref mut stream) = self.stream {
                     self.flags.insert(Flags::RESERVED);
                     stream.reserve_capacity(cmp::min(self.buffer.len(), CHUNK_SIZE));
@@ -170,10 +167,10 @@ impl<H: 'static> Writer for H2Writer<H> {
         if !self.flags.contains(Flags::DISCONNECTED) {
             if self.flags.contains(Flags::STARTED) {
                 // TODO: add warning, write after EOF
-                self.encoder.write(payload.as_ref())?;
+                self.buffer.write(payload.as_ref())?;
             } else {
                 // might be response for EXCEPT
-                self.buffer.extend_from_slice(payload.as_ref())
+                error!("Not supported");
             }
         }
 
@@ -186,7 +183,7 @@ impl<H: 'static> Writer for H2Writer<H> {
 
     fn write_eof(&mut self) -> io::Result<WriterState> {
         self.flags.insert(Flags::EOF);
-        if !self.encoder.write_eof()? {
+        if !self.buffer.write_eof()? {
             Err(io::Error::new(
                 io::ErrorKind::Other,
                 "Last payload item, but eof is not reached",

src/server/shared.rs

@@ -5,8 +5,6 @@ use std::rc::Rc;
 
 use bytes::BytesMut;
 
-use body::Binary;
-
 #[derive(Debug)]
 pub(crate) struct SharedBytesPool(RefCell<VecDeque<BytesMut>>);
 
@@ -50,6 +48,10 @@ impl SharedBytes {
         SharedBytes(Some(bytes), Some(pool))
     }
 
+    pub fn empty() -> SharedBytes {
+        SharedBytes(Some(BytesMut::new()), None)
+    }
+
     #[inline]
     pub(crate) fn get_mut(&mut self) -> &mut BytesMut {
         self.0.as_mut().unwrap()
@@ -79,9 +81,8 @@ impl SharedBytes {
     }
 
     #[inline]
-    pub fn extend(&mut self, data: &Binary) {
-        let buf = self.get_mut();
-        buf.extend_from_slice(data.as_ref());
+    pub fn reserve(&mut self, cap: usize) {
+        self.get_mut().reserve(cap);
     }
 
     #[inline]