lin_kv/
lin_kv.rs

1/// ```bash
2/// $ cargo build --examples
3/// $ RUST_LOG=debug ~/Projects/maelstrom/maelstrom test -w lin-kv --bin ./target/debug/examples/lin_kv --node-count 4 --concurrency 2n --time-limit 20 --rate 100 --log-stderr
4/// ````
5use async_trait::async_trait;
6use maelstrom::kv::{lin_kv, Storage, KV};
7use maelstrom::protocol::Message;
8use maelstrom::{done, Node, Result, Runtime};
9use serde::{Deserialize, Serialize};
10use std::sync::Arc;
11use tokio_context::context::Context;
12
13pub(crate) fn main() -> Result<()> {
14    Runtime::init(try_main())
15}
16
17async fn try_main() -> Result<()> {
18    let runtime = Runtime::new();
19    let handler = Arc::new(handler(runtime.clone()));
20    runtime.with_handler(handler).run().await
21}
22
23#[derive(Clone)]
24struct Handler {
25    s: Storage,
26}
27
28#[async_trait]
29impl Node for Handler {
30    async fn process(&self, runtime: Runtime, req: Message) -> Result<()> {
31        let (ctx, _handler) = Context::new();
32        let msg: Result<Request> = req.body.as_obj();
33        match msg {
34            Ok(Request::Read { key }) => {
35                let value = self.s.get(ctx, key.to_string()).await?;
36                return runtime.reply(req, Request::ReadOk { value }).await;
37            }
38            Ok(Request::Write { key, value }) => {
39                self.s.put(ctx, key.to_string(), value).await?;
40                return runtime.reply(req, Request::WriteOk {}).await;
41            }
42            Ok(Request::Cas { key, from, to, put }) => {
43                self.s.cas(ctx, key.to_string(), from, to, put).await?;
44                return runtime.reply(req, Request::CasOk {}).await;
45            }
46            _ => done(runtime, req),
47        }
48    }
49}
50
51fn handler(runtime: Runtime) -> Handler {
52    Handler { s: lin_kv(runtime) }
53}
54
55#[derive(Serialize, Deserialize)]
56#[serde(rename_all = "snake_case", tag = "type")]
57enum Request {
58    Read {
59        key: u64,
60    },
61    ReadOk {
62        value: i64,
63    },
64    Write {
65        key: u64,
66        value: i64,
67    },
68    WriteOk {},
69    Cas {
70        key: u64,
71        from: i64,
72        to: i64,
73        #[serde(default, rename = "create_if_not_exists")]
74        put: bool,
75    },
76    CasOk {},
77}