From 3725b633f6f7074d0c8870bafa3584cbc97577fd Mon Sep 17 00:00:00 2001
From: Ryan Bottriell <ryan@bottriell.ca>
Date: Thu, 14 Mar 2024 21:10:13 -0700
Subject: [PATCH] Use flatbuffers in grpc instead of protobuf fields

Signed-off-by: Ryan Bottriell <ryan@bottriell.ca>
---
 crates/spfs/build.rs                    | 2 +-
 crates/spfs/src/graph/object.rs         | 5 +++++
 crates/spfs/src/proto/conversions.rs    | 9 +++++++++
 crates/spfs/src/proto/defs/types.proto  | 2 ++
 crates/spfs/src/server/database.rs      | 2 +-
 crates/spfs/src/storage/rpc/database.rs | 2 +-
 6 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/crates/spfs/build.rs b/crates/spfs/build.rs
index 064739b74..fbfdb5cef 100644
--- a/crates/spfs/build.rs
+++ b/crates/spfs/build.rs
@@ -5,7 +5,7 @@
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     #[cfg(feature = "protobuf-src")]
     std::env::set_var("PROTOC", protobuf_src::protoc());
-    tonic_build::configure().compile(
+    tonic_build::configure().bytes(["buffer"]).compile(
         &[
             "src/proto/defs/database.proto",
             "src/proto/defs/repository.proto",
diff --git a/crates/spfs/src/graph/object.rs b/crates/spfs/src/graph/object.rs
index 58cd4afbc..6376328b5 100644
--- a/crates/spfs/src/graph/object.rs
+++ b/crates/spfs/src/graph/object.rs
@@ -467,6 +467,11 @@ impl<T: ObjectProto> FlatObject<T> {
             flatbuffers::root_unchecked::<'_, spfs_proto::AnyObject>(buf)
         }
     }
+
+    /// The inner bytes of this flat buffer
+    pub fn inner_bytes(&self) -> &bytes::Bytes {
+        &self.buf
+    }
 }
 
 impl<'buf, T> FlatObject<T>
diff --git a/crates/spfs/src/proto/conversions.rs b/crates/spfs/src/proto/conversions.rs
index c75eb8ae8..564e7236a 100644
--- a/crates/spfs/src/proto/conversions.rs
+++ b/crates/spfs/src/proto/conversions.rs
@@ -136,6 +136,14 @@ impl From<super::Error> for Error {
     }
 }
 
+impl<T: graph::ObjectProto> From<&graph::FlatObject<T>> for super::Object {
+    fn from(value: &graph::FlatObject<T>) -> Self {
+        Self {
+            kind: Some(super::object::Kind::Buffer(value.inner_bytes().clone())),
+        }
+    }
+}
+
 impl From<&graph::object::Enum> for super::Object {
     fn from(source: &graph::object::Enum) -> Self {
         use super::object::Kind;
@@ -172,6 +180,7 @@ impl TryFrom<super::Object> for graph::Object {
                 "Unexpected and unsupported object kind {:?}",
                 source.kind
             ))),
+            Some(Kind::Buffer(buf)) => graph::Object::new(buf),
             None => Err(Error::String(
                 "Expected non-empty object kind in rpc message".to_string(),
             )),
diff --git a/crates/spfs/src/proto/defs/types.proto b/crates/spfs/src/proto/defs/types.proto
index b7498f097..f7f13a206 100644
--- a/crates/spfs/src/proto/defs/types.proto
+++ b/crates/spfs/src/proto/defs/types.proto
@@ -21,6 +21,8 @@ message Object {
         Tree tree = 4 [deprecated = true];
         Blob blob = 5;
         bool mask = 6 [deprecated = true];
+        // A flatbuffer containing the object
+        bytes buffer = 7;
     }
 }
 
diff --git a/crates/spfs/src/server/database.rs b/crates/spfs/src/server/database.rs
index 46de0c89f..c4e2dffcf 100644
--- a/crates/spfs/src/server/database.rs
+++ b/crates/spfs/src/server/database.rs
@@ -48,7 +48,7 @@ impl proto::database_service_server::DatabaseService for DatabaseService {
         let request = request.into_inner();
         let digest = proto::handle_error!(convert_digest(request.digest));
         let object = { proto::handle_error!(self.repo.read_object(digest).await) };
-        let result = proto::ReadObjectResponse::ok((&object.into_enum()).into());
+        let result = proto::ReadObjectResponse::ok((&object).into());
         Ok(Response::new(result))
     }
 
diff --git a/crates/spfs/src/storage/rpc/database.rs b/crates/spfs/src/storage/rpc/database.rs
index 37b518878..0e77e623b 100644
--- a/crates/spfs/src/storage/rpc/database.rs
+++ b/crates/spfs/src/storage/rpc/database.rs
@@ -70,7 +70,7 @@ impl graph::DatabaseView for super::RpcRepository {
 impl graph::Database for super::RpcRepository {
     async fn write_object<T: ObjectProto>(&self, obj: &graph::FlatObject<T>) -> Result<()> {
         let request = proto::WriteObjectRequest {
-            object: Some((&obj.to_enum()).into()),
+            object: Some(obj.into()),
         };
         self.db_client
             .clone()