exercises

Log | Files | Refs | README

connection.rs (4347B)

      1 use bytes::{Buf, BytesMut};
      2 use mini_redis::{Frame, Result, frame::Error::Incomplete};
      3 use std::io::Cursor;
      4 use tokio::{
      5     io::{self, AsyncReadExt, AsyncWriteExt, BufWriter},
      6     net::TcpStream,
      7 };
      8 
      9 pub struct Connection {
     10     stream: BufWriter<TcpStream>,
     11     buffer: BytesMut,
     12 }
     13 
     14 impl Connection {
     15     pub fn new(stream: TcpStream) -> Connection {
     16         Connection {
     17             stream: BufWriter::new(stream),
     18             // Allocate the buffer with 4kb of capacity.
     19             buffer: BytesMut::with_capacity(4096),
     20         }
     21     }
     22 
     23     /// Read a frame from the connection
     24     ///
     25     /// Returns `None` if EOF is reached
     26     pub async fn read_frame(&mut self) -> Result<Option<Frame>> {
     27         loop {
     28             // Attempt to parse a frame from the buffered data. If
     29             // enough data has been buffered, the frame is
     30             // returned.
     31             if let Some(frame) = self.parse_frame()? {
     32                 return Ok(Some(frame));
     33             }
     34 
     35             // There is not enough buffered data to read a frame.
     36             // Attempt to read more data from the socket.
     37             //
     38             // On success, the number of bytes is returned. `0`
     39             // indicates "end of stream".
     40             if 0 == self.stream.read_buf(&mut self.buffer).await? {
     41                 // The remote closed the connection. For this to be
     42                 // a clean shutdown, there should be no data in the
     43                 // read buffer. If there is, this means that the
     44                 // peer closed the socket while sending a frame.
     45                 if self.buffer.is_empty() {
     46                     return Ok(None);
     47                 } else {
     48                     return Err("connection reset by peer".into());
     49                 }
     50             }
     51         }
     52     }
     53 
     54     fn parse_frame(&mut self) -> Result<Option<Frame>> {
     55         // Create the `T: Buf` type
     56         let mut buf = Cursor::new(&self.buffer[..]);
     57 
     58         // Check whether a full frame is available
     59         match Frame::check(&mut buf) {
     60             Ok(_) => {
     61                 // Get the bytes length of the frame
     62                 let len = buf.position() as usize;
     63 
     64                 // Reset the internal cursor for the call to `parse`.
     65                 buf.set_position(0);
     66 
     67                 // Parse the frame
     68                 let frame = Frame::parse(&mut buf)?;
     69 
     70                 // Discard the frame from the buffer
     71                 self.buffer.advance(len);
     72 
     73                 // Return the frame to the caller
     74                 Ok(Some(frame))
     75             }
     76             Err(Incomplete) => Ok(None),
     77             Err(e) => Err(e.into()),
     78         }
     79     }
     80 
     81     /// Write a frame to the connection
     82     pub async fn write_frame(&mut self, frame: &Frame) -> io::Result<()> {
     83         match frame {
     84             Frame::Simple(val) => {
     85                 self.stream.write_u8(b'+').await?;
     86                 self.stream.write_all(val.as_bytes()).await?;
     87                 self.stream.write_all(b"\r\n").await?;
     88             }
     89             Frame::Error(val) => {
     90                 self.stream.write_u8(b'-').await?;
     91                 self.stream.write_all(val.as_bytes()).await?;
     92                 self.stream.write_all(b"\r\n").await?;
     93             }
     94             Frame::Integer(val) => {
     95                 self.stream.write_u8(b':').await?;
     96                 self.write_decimal(*val).await?;
     97             }
     98             Frame::Bulk(val) => {
     99                 let len = val.len();
    100 
    101                 self.stream.write_u8(b'$').await?;
    102                 self.write_decimal(len as u64).await?;
    103                 self.stream.write_all(val).await?;
    104                 self.stream.write_all(b"\r\n").await?;
    105             }
    106             Frame::Null => {
    107                 self.stream.write_all(b"$-1\r\n").await?;
    108             }
    109             Frame::Array(frames) => unimplemented!(),
    110         }
    111 
    112         self.stream.flush().await;
    113 
    114         Ok(())
    115     }
    116 
    117     /// Write a decimal frame to the stream
    118     async fn write_decimal(&mut self, val: u64) -> io::Result<()> {
    119         use std::io::Write;
    120 
    121         // Convert the value to a string
    122         let mut buf = [0u8; 12];
    123         let mut buf = Cursor::new(&mut buf[..]);
    124         write!(&mut buf, "{}", val)?;
    125 
    126         let pos = buf.position() as usize;
    127         self.stream.write_all(&buf.get_ref()[..pos]).await?;
    128         self.stream.write_all(b"\r\n").await?;
    129 
    130         Ok(())
    131     }
    132 }