fix: convert runner protocol to protobuf

Author: MasterPtato

packages/edge/infra/client/actor-kv/Cargo.toml

@@ -11,7 +11,8 @@ fdb-util.workspace = true
 foundationdb.workspace = true
 futures-util = { version = "0.3" }
 indexmap = { version = "2.0" }
-prost = "0.13.3"
+pegboard-config.workspace = true
+prost = "0.14"
 rivet-util-id.workspace = true
 serde = { version = "1.0.195", features = ["derive"] }
 serde_json = "1.0.111"

packages/edge/infra/client/actor-kv/src/entry.rs

@@ -1,13 +1,27 @@
 use anyhow::*;
 use foundationdb as fdb;
+use pegboard_config::runner_protocol::proto::kv;
 use prost::Message;
-use serde::{Deserialize, Serialize};
 
-use crate::{key::Key, metadata::Metadata};
+use crate::key::Key;
+
+/// Represents a Rivet KV value.
+#[derive(Clone, Debug)]
+pub struct Entry {
+	inner: kv::Entry,
+}
+
+impl std::ops::Deref for Entry {
+	type Target = kv::Entry;
+
+	fn deref(&self) -> &Self::Target {
+		&self.inner
+	}
+}
 
 #[derive(Default)]
 pub(crate) struct EntryBuilder {
-	metadata: Option<Metadata>,
+	metadata: Option<kv::Metadata>,
 	value: Vec<u8>,
 	next_idx: usize,
 }
@@ -19,7 +33,7 @@ impl EntryBuilder {
 				// We ignore setting the metadata again because it means the same key was given twice in the
 				// input keys for `ActorKv::get`. We don't perform automatic deduplication.
 				if self.metadata.is_none() {
-					self.metadata = Some(Metadata::decode(value.value())?);
+					self.metadata = Some(kv::Metadata::decode(value.value())?);
 				}
 			}
 			SubKey::Chunk(idx, value) => {
@@ -40,21 +54,17 @@ impl EntryBuilder {
 		ensure!(!self.value.is_empty(), "empty value at key {key:?}");
 
 		Ok(Entry {
-			metadata: self
-				.metadata
-				.with_context(|| format!("no metadata for key {key:?}"))?,
-			value: self.value,
+			inner: kv::Entry {
+				metadata: Some(
+					self.metadata
+						.with_context(|| format!("no metadata for key {key:?}"))?,
+				),
+				value: self.value,
+			},
 		})
 	}
 }
 
-/// Represents a Rivet KV value.
-#[derive(Clone, Debug, Serialize, Deserialize)]
-pub struct Entry {
-	pub metadata: Metadata,
-	pub value: Vec<u8>,
-}
-
 /// Represents FDB keys within a Rivet KV key.
 pub(crate) enum SubKey {
 	Metadata(fdb::future::FdbValue),

packages/edge/infra/client/actor-kv/src/key.rs

@@ -1,11 +1,14 @@
 use foundationdb::tuple::{
 	Bytes, PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset,
 };
-use serde::{Deserialize, Serialize};
+use pegboard_config::runner_protocol::proto::kv;
+use prost::Message;
 
 // TODO: Custom deser impl that uses arrays instead of objects?
-#[derive(Clone, Serialize, Deserialize)]
-pub struct Key(Vec<Vec<u8>>);
+#[derive(Clone)]
+pub struct Key {
+	inner: kv::Key,
+}
 
 impl std::fmt::Debug for Key {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -15,24 +18,23 @@ impl std::fmt::Debug for Key {
 
 impl PartialEq for Key {
 	fn eq(&self, other: &Self) -> bool {
-		self.0 == other.0
+		self.inner == other.inner
 	}
 }
 
 impl Eq for Key {}
 
 impl std::hash::Hash for Key {
 	fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-		for buffer in &self.0 {
+		for buffer in &self.inner.segments {
 			state.write(buffer);
 		}
 	}
 }
 
 impl Key {
 	pub fn len(&self) -> usize {
-		// Arbitrary 4 accounting for nesting overhead
-		self.0.iter().fold(0, |acc, x| acc + x.len()) + 4 * self.0.len()
+		self.inner.encoded_len()
 	}
 }
 
@@ -47,7 +49,7 @@ impl TuplePack for Key {
 		w.write_all(&[fdb_util::codes::NESTED])?;
 		offset += 1;
 
-		for v in self.0.iter() {
+		for v in self.inner.segments.iter() {
 			offset += v.pack(w, tuple_depth.increment())?;
 		}
 
@@ -62,22 +64,29 @@ impl<'de> TupleUnpack<'de> for Key {
 	fn unpack(mut input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
 		input = fdb_util::parse_code(input, fdb_util::codes::NESTED)?;
 
-		let mut vec = Vec::new();
+		let mut segments = Vec::new();
 		while !is_end_of_tuple(input, true) {
 			let (rem, v) = Bytes::unpack(input, tuple_depth.increment())?;
 			input = rem;
-			vec.push(v.into_owned());
+			segments.push(v.into_owned());
 		}
 
 		input = fdb_util::parse_code(input, fdb_util::codes::NIL)?;
 
-		Ok((input, Key(vec)))
+		Ok((
+			input,
+			Key {
+				inner: kv::Key { segments },
+			},
+		))
 	}
 }
 
 /// Same as Key: except when packing, it leaves off the NIL byte to allow for an open range.
-#[derive(Clone, Serialize, Deserialize)]
-pub struct ListKey(Vec<Vec<u8>>);
+#[derive(Clone)]
+pub struct ListKey {
+	inner: kv::Key,
+}
 
 impl std::fmt::Debug for ListKey {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -96,7 +105,7 @@ impl TuplePack for ListKey {
 		w.write_all(&[fdb_util::codes::NESTED])?;
 		offset += 1;
 
-		for v in &self.0 {
+		for v in &self.inner.segments {
 			offset += v.pack(w, tuple_depth.increment())?;
 		}
 
@@ -108,8 +117,7 @@ impl TuplePack for ListKey {
 
 impl ListKey {
 	pub fn len(&self) -> usize {
-		// Arbitrary 4 accounting for nesting overhead
-		self.0.iter().fold(0, |acc, x| acc + x.len()) + 4 * self.0.len()
+		self.inner.encoded_len()
 	}
 }
 

packages/edge/infra/client/actor-kv/src/lib.rs

@@ -14,15 +14,14 @@ use indexmap::IndexMap;
 pub use key::Key;
 use list_query::ListLimitReached;
 pub use list_query::ListQuery;
-pub use metadata::Metadata;
+use pegboard_config::runner_protocol::proto::kv;
 use prost::Message;
 use tokio::sync::Mutex;
 use utils::{validate_entries, validate_keys, TransactionExt};
 
 mod entry;
 pub mod key;
 mod list_query;
-mod metadata;
 mod utils;
 
 const MAX_KEY_SIZE: usize = 2 * 1024;
@@ -325,8 +324,8 @@ impl ActorKv {
 								// Clear previous before setting
 								tx.clear_subspace_range(&key_subspace);
 
-								let metadata = Metadata {
-									kv_version: self.version.as_bytes().to_vec(),
+								let metadata = kv::Metadata {
+									version: self.version.as_bytes().to_vec(),
 									create_ts: utils::now(),
 								};
 								let mut buf = Vec::new();

packages/edge/infra/client/actor-kv/src/list_query.rs

@@ -1,34 +1,39 @@
 use anyhow::*;
 use foundationdb::tuple::Subspace;
 use indexmap::IndexMap;
-use serde::{Deserialize, Serialize};
 
 use crate::{
 	entry::EntryBuilder,
 	key::{Key, ListKey},
 	MAX_KEY_SIZE,
 };
 
-#[derive(Clone, Debug, Serialize, Deserialize)]
-#[serde(rename_all = "snake_case")]
+#[derive(Clone, Debug)]
 pub enum ListQuery {
 	All,
-	RangeInclusive(ListKey, Key),
-	RangeExclusive(ListKey, Key),
+	Range {
+		start: ListKey,
+		end: Key,
+		exclusive: bool,
+	},
 	Prefix(ListKey),
 }
 
 impl ListQuery {
 	pub(crate) fn range(&self, subspace: &Subspace) -> (Vec<u8>, Vec<u8>) {
 		match self {
 			ListQuery::All => subspace.range(),
-			ListQuery::RangeInclusive(start, end) => (
+			ListQuery::Range {
+				start,
+				end,
+				exclusive,
+			} => (
 				subspace.subspace(&start).range().0,
-				subspace.subspace(&end).range().1,
-			),
-			ListQuery::RangeExclusive(start, end) => (
-				subspace.subspace(&start).range().0,
-				subspace.subspace(&end).range().1,
+				if *exclusive {
+					subspace.subspace(&end).range().0
+				} else {
+					subspace.subspace(&end).range().1
+				},
 			),
 			ListQuery::Prefix(prefix) => subspace.subspace(&prefix).range(),
 		}
@@ -37,7 +42,7 @@ impl ListQuery {
 	pub(crate) fn validate(&self) -> Result<()> {
 		match self {
 			ListQuery::All => {}
-			ListQuery::RangeInclusive(start, end) => {
+			ListQuery::Range { start, end, .. } => {
 				ensure!(
 					start.len() <= MAX_KEY_SIZE,
 					"start key is too long (max 2048 bytes)"
@@ -47,16 +52,6 @@ impl ListQuery {
 					"end key is too long (max 2048 bytes)"
 				);
 			}
-			ListQuery::RangeExclusive(start, end) => {
-				ensure!(
-					start.len() <= MAX_KEY_SIZE,
-					"startAfter key is too long (max 2048 bytes)"
-				);
-				ensure!(
-					end.len() <= MAX_KEY_SIZE,
-					"end key is too long (max 2048 bytes)"
-				);
-			}
 			ListQuery::Prefix(prefix) => {
 				ensure!(
 					prefix.len() <= MAX_KEY_SIZE,

packages/edge/infra/client/actor-kv/src/metadata.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, result::Result::Ok};
+use std::result::Result::Ok;
 
 use anyhow::*;
 use foundationdb as fdb;

packages/edge/infra/client/actor-kv/src/utils.rs

@@ -10,11 +10,14 @@ anyhow = "1.0"
 indexmap = { version = "2.0" }
 ipnet = { version = "2.10.1", features = ["serde"] }
 pegboard.workspace = true
-pegboard-actor-kv.workspace = true
+prost = "0.14"
 rivet-util-id.workspace = true
 schemars = { version = "0.8.21", features = ["url", "uuid1"] }
 serde = { version = "1.0.195", features = ["derive"] }
 serde_json = "1.0"
 tokio-util = { version = "0.7", features = ["codec"] }
 url = "2.2.2"
-uuid = { version = "1.6.1", features = ["v4"] }
+uuid = { version = "1.6.1", features = ["v4"] }
+
+[build-dependencies]
+prost-build = "0.14"

packages/edge/infra/client/config/Cargo.toml

@@ -0,0 +1,13 @@
+use std::io::Result;
+
+fn main() -> Result<()> {
+	prost_build::compile_protos(
+		&[
+			"resources/proto/kv.proto",
+			"resources/proto/runner_protocol.proto",
+		],
+		&["resources/proto/"],
+	)?;
+
+	Ok(())
+}