dicedb-rs 0.1.3

Rust SDK for DiceDb.
Documentation
use std::io::{Read, Write};

use crate::{
    commands::{Command, CommandExecutor, Value, WatchValue},
    errors::StreamError,
};

pub trait Stream {
    fn host(&self) -> &str;
    fn port(&self) -> u16;
    fn set_stream(&mut self, stream: std::net::TcpStream);
    fn tcp_stream(&mut self) -> &std::net::TcpStream;
    fn handshake(&mut self) -> Result<(), StreamError>;
}

pub trait Reconnectable {
    fn reconnect(&mut self, max_tries: u64) -> Result<(), StreamError>;
}

pub trait ValueReceiver {
    fn receive_value(&mut self) -> Result<Value, StreamError>;
}

pub trait WatchValueReceiver {
    fn recieve_watchvalue(&mut self) -> Result<WatchValue, StreamError>;
}

pub trait CommandSender {
    fn send_command(&mut self, command: Command) -> Result<(), StreamError>;
}

impl<T: Stream> Reconnectable for T {
    fn reconnect(&mut self, max_tries: u64) -> Result<(), StreamError> {
        let mut tries = 0;
        while tries < max_tries {
            tries += 1;
            let stream = std::net::TcpStream::connect(format!("{}:{}", self.host(), self.port()));
            match stream {
                Ok(stream) => {
                    self.set_stream(stream);
                    self.handshake()?;
                    return Ok(());
                }
                Err(_) => {
                    std::thread::sleep(std::time::Duration::from_secs(1));
                    continue;
                }
            }
        }
        Err(StreamError::IoError(std::io::Error::new(
            std::io::ErrorKind::Other,
            "Max attempts reached",
        )))
    }
}

const MAX_REQUEST_SIZE: usize = 32 * 1024 * 1024; // 32MB per session, meeh probably too much, fi.

impl<T: Stream> WatchValueReceiver for T {
    fn recieve_watchvalue(&mut self) -> Result<WatchValue, StreamError> {
        let mut buffer = vec![0; MAX_REQUEST_SIZE];
        let size = self.tcp_stream().read(&mut buffer)?;
        let reply_slice = &buffer[..size];
        let val = WatchValue::decode_watchvalue(reply_slice)?;
        Ok(val)
    }
}

impl<T: Stream> ValueReceiver for T {
    fn receive_value(&mut self) -> Result<Value, StreamError> {
        let mut buffer = vec![0; MAX_REQUEST_SIZE];
        let size = self.tcp_stream().read(&mut buffer)?;
        let reply_slice = &buffer[..size];
        let val = Value::decode_value(reply_slice)?;
        Ok(val)
    }
}

impl<T: Stream> CommandSender for T {
    fn send_command(&mut self, command: Command) -> Result<(), StreamError> {
        eprintln!("Sending command: {:?}", command);
        let serialized_command = command.encode();
        eprintln!("Sending command: {:?}", serialized_command);
        match self.tcp_stream().write_all(&serialized_command) {
            Ok(_) => Ok(()),
            Err(_) => {
                self.reconnect(10)?;
                self.tcp_stream().write_all(&serialized_command)?;
                Ok(())
            }
        }
    }
}

impl<T: Stream> CommandExecutor for T {
    fn execute_command(&mut self, command: Command) -> Result<Value, StreamError> {
        self.send_command(command)?;
        self.receive_value()
    }
}

#[cfg(test)]
mod tests {

    use crate::commandstream::CommandStream;

    use super::*;

    #[test]
    fn test_reconnect() {
        let mut command_client = CommandStream::new("localhost".to_string(), 7379).unwrap();
        let reconnect_result = command_client.reconnect(10);
        assert!(reconnect_result.is_ok());
    }
}