tokio-tutorial-mini-redis/src/bin/server.rs

   1 use std::collections::HashMap;
   2 use std::sync::{Arc, Mutex};
   3
   4 use bytes::Bytes;
   5 use mini_redis::{Connection, Frame};
   6 use mini_redis::Command::{self, Get, Set};
   7 use tokio::net::{TcpListener, TcpStream};
   8
   9 type Db = Arc<Mutex<HashMap<String, Bytes>>>;
  10
  11 const LISTEN: &str = "127.0.0.1:6379";
  12
  13 #[tokio::main]
  14 async fn main() {
  15     env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("debug")).init();
  16
  17     let listener = TcpListener::bind(LISTEN).await.unwrap();
  18     log::info!("Listening on {}", LISTEN);
  19     let db: Db = Arc::new(Mutex::new(HashMap::new()));
  20
  21     loop {
  22         // The second item contains the IP and port of the new connection
  23         let (socket, addr) = listener.accept().await.unwrap();
  24         log::debug!("got connection from {:?}", addr);
  25         let db_i = db.clone();
  26         tokio::spawn(async move { process(socket, db_i).await });
  27     }
  28 }
  29
  30 async fn process(socket: TcpStream, db: Db) {
  31     let mut connection = Connection::new(socket);
  32
  33     while let Some(frame) = connection.read_frame().await.unwrap() {
  34         let response = match Command::from_frame(frame).unwrap() {
  35             Set(cmd) => {
  36                 // The value is stored as `Vec<u8>`
  37                 db.lock().unwrap().insert(cmd.key().to_string(), cmd.value().clone());
  38                 log::debug!("Set {} → {:?}", &cmd.key(), &cmd.value());
  39                 Frame::Simple("OK".to_string())
  40             }
  41             Get(cmd) => {
  42                 if let Some(value) = db.lock().unwrap().get(cmd.key()) {
  43                     log::debug!("Get {} → {:?}", &cmd.key(), &value);
  44                     Frame::Bulk(value.clone())
  45                 } else {
  46                     log::debug!("Get {} unknown key", &cmd.key());
  47                     Frame::Null
  48                 }
  49             }
  50             cmd => panic!("unimplemented {:?}", cmd),
  51         };
  52
  53         // Write the response to the client
  54         connection.write_frame(&response).await.unwrap();
  55     }
  56 }