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 #[tokio::main]
  12 async fn main() {
  13     let listener = TcpListener::bind("127.0.0.1:6379").await.unwrap();
  14     let db: Db = Arc::new(Mutex::new(HashMap::new()));
  15
  16     loop {
  17         // The second item contains the IP and port of the new connection
  18         let (socket, _) = listener.accept().await.unwrap();
  19         let db_i = db.clone();
  20         tokio::spawn(async move { process(socket, db_i).await });
  21     }
  22 }
  23
  24 async fn process(socket: TcpStream, db: Db) {
  25     let mut connection = Connection::new(socket);
  26
  27     while let Some(frame) = connection.read_frame().await.unwrap() {
  28         let response = match Command::from_frame(frame).unwrap() {
  29             Set(cmd) => {
  30                 // The value is stored as `Vec<u8>`
  31                 db.lock().unwrap().insert(cmd.key().to_string(), cmd.value().clone());
  32                 Frame::Simple("OK".to_string())
  33             }
  34             Get(cmd) => {
  35                 if let Some(value) = db.lock().unwrap().get(cmd.key()) {
  36                     Frame::Bulk(value.clone())
  37                 } else {
  38                     Frame::Null
  39                 }
  40             }
  41             cmd => panic!("unimplemented {:?}", cmd),
  42         };
  43
  44         // Write the response to the client
  45         connection.write_frame(&response).await.unwrap();
  46     }
  47 }