use std::io;
use std::marker::Unpin;
use byteorder::{BigEndian, ByteOrder};
use bytes::Bytes;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use Operation::*;
#[derive(Debug, PartialEq, Eq)]
pub enum Operation {
Ping,
Read { key: Bytes, offset: u64, limit: u64 },
Evict(Bytes),
}
impl Operation {
pub async fn decode<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Self, io::Error> {
let opcode = reader.read_u8().await?;
match opcode {
1 => Ok(Ping),
2 => {
let key_length = reader.read_u16().await? as usize;
let mut key = vec![0; key_length + 16];
reader.read_exact(&mut key).await?;
let offset = BigEndian::read_u64(&key[key_length..]);
let limit = BigEndian::read_u64(&key[key_length + 8..]);
key.truncate(key_length);
Ok(Read {
key: key.into(),
offset,
limit,
})
}
3 => {
let key_length = reader.read_u16().await? as usize;
let mut key = vec![0; key_length];
reader.read_exact(&mut key).await?;
Ok(Evict(key.into()))
}
_ => Err(io::Error::new(
io::ErrorKind::Other,
format!("unknown opcode({})", opcode),
)),
}
}
pub async fn encode<W: AsyncWrite + Unpin>(&self, writer: &mut W) -> Result<(), io::Error> {
writer.write_u8(self.opcode()).await?;
match self {
Ping => Ok(()),
Read { key, offset, limit } => {
writer.write_u16(key.len() as u16).await?;
writer.write_all(&key).await?;
writer.write_u64(*offset).await?;
writer.write_u64(*limit).await
}
Evict(key) => {
writer.write_u16(key.len() as u16).await?;
writer.write_all(&key).await
}
}
}
fn size(&self) -> usize {
match self {
Ping => 1,
Read { key, offset, limit } => 2 + key.len() + 16,
Evict(key) => 2 + key.len(),
}
}
fn opcode(&self) -> u8 {
match self {
Ping => 1,
Read {
key: _,
offset: _,
limit: _,
} => 2,
Evict(_) => 3,
}
}
}
#[cfg(test)]
mod tests {
use std::io;
use super::Operation;
#[tokio::test]
async fn test_operation() -> Result<(), io::Error> {
for o in [
Operation::Ping,
Operation::Read {
key: "aaa".into(),
offset: 0,
limit: 0,
},
Operation::Read {
key: "bbb".into(),
offset: 1 << 10,
limit: 1 << 10,
},
Operation::Evict("ccc".into()),
] {
let mut buffer = Vec::new();
o.encode(&mut buffer).await?;
let mut s = &buffer[..];
let op = Operation::decode(&mut s).await?;
assert_eq!(op, o);
}
Ok(())
}
}