1use 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}