exercises

Log | Files | Refs | README

server.rs (2696B)

      1 use bytes::Bytes;
      2 use mini_redis::{Connection, Frame};
      3 use std::collections::HashMap;
      4 use std::sync::{Arc, Mutex};
      5 use tokio::net::{TcpListener, TcpStream};
      6 
      7 // Using Arc allows the HashMap to be referenced concurrently from many tasks,
      8 // potentially running on many threads. Throughout Tokio, the term handle is
      9 // used to reference a value that provides access to some shared state.
     10 //
     11 // Use std::sync::Mutex instead of tokio::sync::Mutex
     12 type Db = Arc<Mutex<HashMap<String, Bytes>>>;
     13 
     14 #[tokio::main]
     15 async fn main() {
     16     // Bind the listener to the address
     17     let listener = TcpListener::bind("127.0.0.1:6379").await.unwrap();
     18 
     19     println!("Listening");
     20 
     21     // The HashMap will be shared across many tasks and potentially many threads.
     22     // To support this, it is wrapped in Arc<Mutex<_>>
     23     let db = Arc::new(Mutex::new(HashMap::new()));
     24 
     25     loop {
     26         // The second item contains the IP and port of the new connection.
     27         let (socket, _) = listener.accept().await.unwrap();
     28 
     29         // Clone the handle to the hash map
     30         let db = db.clone();
     31 
     32         println!("Accepted");
     33 
     34         // A new task is spawned for each inbound socket. The socket is
     35         // moved to the new task and processed there.
     36         tokio::spawn(async move {
     37             process(socket, db).await;
     38         });
     39     }
     40 }
     41 
     42 /// Process function handles incoming commands. It uses a HashMap to store values.
     43 /// SET commands will insert into the HashMap and GET values will load them.
     44 /// Additionally, we will use a loop to accept more than one command per connection.
     45 async fn process(socket: TcpStream, db: Db) {
     46     use mini_redis::Command::{self, Get, Set};
     47 
     48     // Connection, provided by `mini-redis`, handles parsing frames from the socket
     49     let mut connection = Connection::new(socket);
     50 
     51     // Use `read_frame` to receive a command from the connection.
     52     while let Some(frame) = connection.read_frame().await.unwrap() {
     53         let response = match Command::from_frame(frame).unwrap() {
     54             Set(cmd) => {
     55                 let mut db = db.lock().unwrap();
     56                 // The value is stored as Vec<u8>
     57                 db.insert(cmd.key().to_string(), cmd.value().clone());
     58                 Frame::Simple("OK".to_string())
     59             }
     60             Get(cmd) => {
     61                 let db = db.lock().unwrap();
     62                 if let Some(value) = db.get(cmd.key()) {
     63                     Frame::Bulk(value.clone())
     64                 } else {
     65                     Frame::Null
     66                 }
     67             }
     68             cmd => panic!("unimplemented {cmd:?}"),
     69         };
     70 
     71         // Write the response to the client
     72         connection.write_frame(&response).await.unwrap();
     73     }
     74 }