Use flatbuffers in grpc instead of protobuf fields

Signed-off-by: Ryan Bottriell <ryan@bottriell.ca>
spkenv · Mar 15, 2024 · 9cd9f93 · 9cd9f93
1 parent 3da6292
commit 9cd9f93
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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(),
             )),
@@ -460,9 +469,9 @@ impl TryFrom<super::DigestSearchCriteria> for graph::DigestSearchCriteria {
                 super::digest_search_criteria::Criteria::All(_) => {
                     Ok(graph::DigestSearchCriteria::All)
                 }
-                super::digest_search_criteria::Criteria::StartsWith(bytes) => {
-                    Ok(graph::DigestSearchCriteria::StartsWith(bytes.bytes.into()))
-                }
+                super::digest_search_criteria::Criteria::StartsWith(bytes) => Ok(
+                    graph::DigestSearchCriteria::StartsWith(bytes.bytes.to_vec().into()),
+                ),
             },
             None => Err("Unknown criteria kind".into()),
         }

diff --git a/crates/spfs/src/proto/defs/types.proto b/crates/spfs/src/proto/defs/types.proto
@@ -15,12 +15,14 @@ message DateTime {
 
 message Object {
     oneof kind {
-        Platform platform = 1;
-        Layer layer = 2;
-        Manifest manifest = 3;
+        Platform platform = 1 ;
+        Layer layer = 2 ;
+        Manifest manifest = 3 ;
         Tree tree = 4 [deprecated = true];
-        Blob blob = 5;
+        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
@@ -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
@@ -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()