nebkor-maelstrom/src/kv.rs
2024-06-09 15:17:36 -07:00

81 lines
2.5 KiB
Rust

use serde_json::Value;
use crate::{check_err, mk_payload, Body, RpcResult, Runner};
/// A convenient way to do RPC with Maelstrom's KV services.
#[derive(Debug, Default, Clone)]
pub struct Kv {
service: &'static str,
}
impl Kv {
/// Construct a proxy to the `seq-kv` service.
pub fn seq() -> Self {
Kv { service: "seq-kv" }
}
/// Construct a proxy to the `lin-kv` service.
pub fn lin() -> Self {
Kv { service: "lin-kv" }
}
/// Construct a proxy to the `lww-kv` service.
pub fn lww() -> Self {
Kv { service: "lww-kv" }
}
/// Returns the `Value` from a remote call to the KV service for key `key`.
///
/// # Examples
///
/// ```no_run
/// # use nebkor_maelstrom::{kv::Kv, Node, Runner, Message};
/// # struct Foo;
/// # impl Node for Foo {fn handle(&mut self, runner: &Runner, msg: Message) {}}
/// # let runner = Runner::new(Foo);
/// // get a proxy to the `seq-kv` service
/// let kv = Kv::seq();
/// let result = kv.read(&runner, "MY_KEY");
/// // if "MY_KEY" had previously been written into the store, then the Result will be `Ok`
/// // and the body will be `Some(Value)`.
/// assert!(result.is_ok() && result.unwrap().is_some());
/// ```
pub fn read(&self, runner: &Runner, key: &str) -> RpcResult {
let payload = mk_payload(&[("key", key.into())]);
let body = Body::from_type("read").with_payload(payload);
let rx = runner.rpc(self.service, body);
let msg = rx.recv().unwrap();
check_err(&msg)?;
Ok(Some(msg.body.payload.get("value").unwrap().to_owned()))
}
/// The success value is `None`.
pub fn write(&self, runner: &Runner, key: &str, val: Value) -> RpcResult {
let payload = mk_payload(&[("key", key.into()), ("value", val)]);
let body = Body::from_type("write").with_payload(payload);
let msg = runner.rpc(self.service, body).recv().unwrap();
check_err(&msg)?;
Ok(None)
}
/// The success value is `None`.
pub fn cas(
&self,
runner: &Runner,
key: &str,
from: Value,
to: Value,
create: bool,
) -> RpcResult {
let payload = mk_payload(&[
("key", key.into()),
("from", from),
("to", to),
("create_if_not_exists", create.into()),
]);
let body = Body::from_type("cas").with_payload(payload);
let msg = runner.rpc(self.service, body).recv().unwrap();
check_err(&msg)?;
Ok(None)
}
}