Merge pull request #97 from http-rs/more-logs

Server encoder state machine

Author: yoshuawuyts

src/chunked/encoder.rs

@@ -72,19 +72,37 @@ impl ChunkedEncoder {
         buf: &mut [u8],
     ) -> Poll<io::Result<usize>> {
         self.bytes_written = 0;
+        let res = self.run(res, cx, buf);
+        log::trace!("ChunkedEncoder {} bytes written", self.bytes_written);
+        res
+    }
+
+    /// Execute the right method for the current state.
+    fn run(
+        &mut self,
+        res: &mut Response,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
         match self.state {
-            State::Start => self.init(res, cx, buf),
+            State::Start => self.dispatch(State::EncodeChunks, res, cx, buf),
             State::EncodeChunks => self.encode_chunks(res, cx, buf),
             State::EndOfChunks => self.encode_chunks_eos(res, cx, buf),
-            State::ReceiveTrailers => self.encode_trailers(res, cx, buf),
+            State::ReceiveTrailers => self.receive_trailers(res, cx, buf),
             State::EncodeTrailers => self.encode_trailers(res, cx, buf),
-            State::EndOfStream => self.encode_eos(cx, buf),
-            State::End => Poll::Ready(Ok(0)),
+            State::EndOfStream => self.encode_eos(res, cx, buf),
+            State::End => Poll::Ready(Ok(self.bytes_written)),
         }
     }
 
     /// Switch the internal state to a new state.
-    fn set_state(&mut self, state: State) {
+    fn dispatch(
+        &mut self,
+        state: State,
+        res: &mut Response,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
         use State::*;
         log::trace!("ChunkedEncoder state: {:?} -> {:?}", self.state, state);
 
@@ -96,21 +114,11 @@ impl ChunkedEncoder {
             ReceiveTrailers => assert!(matches!(state, EncodeTrailers | EndOfStream)),
             EncodeTrailers => assert!(matches!(state, EndOfStream)),
             EndOfStream => assert!(matches!(state, End)),
-            End => panic!("No state transitions allowed after the stream has ended"),
+            End => panic!("No state transitions allowed after the ChunkedEncoder has ended"),
         }
 
         self.state = state;
-    }
-
-    /// Init encoding.
-    fn init(
-        &mut self,
-        res: &mut Response,
-        cx: &mut Context<'_>,
-        buf: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
-        self.set_state(State::EncodeChunks);
-        self.encode_chunks(res, cx, buf)
+        self.run(res, cx, buf)
     }
 
     /// Stream out data using chunked encoding.
@@ -134,8 +142,7 @@ impl ChunkedEncoder {
         // If the stream doesn't have any more bytes left to read we're done
         // sending chunks and it's time to move on.
         if src.len() == 0 {
-            self.set_state(State::EndOfChunks);
-            return self.encode_chunks_eos(res, cx, buf);
+            return self.dispatch(State::EndOfChunks, res, cx, buf);
         }
 
         // Each chunk is prefixed with the length of the data in hex, then a
@@ -180,7 +187,6 @@ impl ChunkedEncoder {
         self.bytes_written += CRLF_LEN;
 
         // Finally return how many bytes we've written to the buffer.
-        log::trace!("sending {} bytes", self.bytes_written);
         Poll::Ready(Ok(self.bytes_written))
     }
 
@@ -203,8 +209,7 @@ impl ChunkedEncoder {
         buf[idx + 2] = LF;
         self.bytes_written += 1 + CRLF_LEN;
 
-        self.set_state(State::ReceiveTrailers);
-        return self.receive_trailers(res, cx, buf);
+        self.dispatch(State::ReceiveTrailers, res, cx, buf)
     }
 
     /// Receive trailers sent to the response, and store them in an internal
@@ -216,31 +221,32 @@ impl ChunkedEncoder {
         buf: &mut [u8],
     ) -> Poll<io::Result<usize>> {
         // TODO: actually wait for trailers to be received.
-        self.set_state(State::EncodeTrailers);
-        self.encode_trailers(res, cx, buf)
+        self.dispatch(State::EncodeTrailers, res, cx, buf)
     }
 
     /// Send trailers to the buffer.
     fn encode_trailers(
         &mut self,
-        _res: &mut Response,
+        res: &mut Response,
         cx: &mut Context<'_>,
         buf: &mut [u8],
     ) -> Poll<io::Result<usize>> {
         // TODO: actually encode trailers here.
-        self.set_state(State::EndOfStream);
-        self.encode_eos(cx, buf)
+        self.dispatch(State::EndOfStream, res, cx, buf)
     }
 
     /// Encode the end of the stream.
-    fn encode_eos(&mut self, _cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+    fn encode_eos(
+        &mut self,
+        res: &mut Response,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
         let idx = self.bytes_written;
         // Write the final CRLF
         buf[idx] = CR;
         buf[idx + 1] = LF;
         self.bytes_written += CRLF_LEN;
-
-        self.set_state(State::End);
-        return Poll::Ready(Ok(self.bytes_written));
+        self.dispatch(State::End, res, cx, buf)
     }
 }