Add streaming database API support to database service

This introduces the streaming variants of Read/WriteMetadataBlob, along
with implementations in the test service, and some tests verifying that
the two of them work together.

No sealed memory production code is changed in this commit, that will be
follow-up work.

Bug: b/476476987
Change-Id: I8bb0d9e4d3540551eb0b354a929fd10d6a6a6964

Author: jblebrun

oak_private_memory/BUILD

@@ -254,15 +254,14 @@ rust_test(
 )
 
 rust_test(
-    name = "search_test",
-    srcs = ["test/search_test.rs"],
+    name = "external_db_client_test",
+    srcs = ["test/external_db_client_test.rs"],
     deps = [
-        "//database",
+        ":external_db_client",
+        ":private_memory_test_database_server_lib",
         "//proto:sealed_memory_rust_proto",
-        "@oak_crates_index//:anyhow",
-        "@oak_crates_index//:googletest",
-        "@oak_crates_index//:prost-types",
-        "@oak_crates_index//:tempfile",
+        "@oak_crates_index//:tokio",
+        "@oak_crates_index//:tonic",
     ],
 )
 

oak_private_memory/proto/database.proto

@@ -84,6 +84,41 @@ message WriteMetadataBlobRequest {
 
 message WriteMetadataBlobResponse {}
 
+message ReadMetadataBlobStreamRequest {
+  string id = 1;
+}
+
+message ReadMetadataBlobStreamResponse {
+  oneof response {
+    // The initial response containing the metadata. The blob field in the
+    // data_blob can be empty or contain the first chunk.
+    //
+    // This variant should only be sent once, as the first request. Sending it
+    // more than once will result in an error.
+    MetadataBlob metadata_blob = 1;
+
+    // Subsequent chunks of the blob. These will be concatenated in the order
+    // received to construct the final payload. Sending this variant before
+    // sending the metadata_blob will result in an error.
+    bytes chunk = 2;
+  }
+}
+
+message WriteMetadataBlobStreamRequest {
+  oneof request {
+    // The initial request containing the metadata. The blob field in the
+    // data_blob can be empty or contain the first chunk.
+    MetadataBlob metadata_blob = 1;
+
+    // Subsequent chunks of the blob. They can be concatenated in the order
+    // received to construct the final payload. The server will indicate the end
+    // of the file by sending a half-close.
+    bytes chunk = 2;
+  }
+}
+
+message WriteMetadataBlobStreamResponse {}
+
 message ResetDatabaseRequest {}
 
 message ResetDatabaseResponse {}
@@ -121,6 +156,26 @@ service SealedMemoryDatabaseService {
   rpc ReadMetadataBlob(ReadMetadataBlobRequest)
       returns (ReadMetadataBlobResponse) {}
 
+  // Write a metadata blob to the database using a stream.
+  // The first request message must contain the metadata_blob.
+  // Subsequent request messages should contain chunks of the blob.
+  //
+  // There's no requirement on the chunk szies; they can very from
+  // request-to-request, and be of any size, as llong as it fits in the message
+  // size limit.
+  rpc WriteMetadataBlobStream(stream WriteMetadataBlobStreamRequest)
+      returns (WriteMetadataBlobStreamResponse) {}
+
+  // Read a metadata blob from the database using a stream.
+  // The first response message will contain the metadata_blob.
+  // Subsequent response messages will contain chunks of the blob.
+  //
+  // The returned chunks can be of any size, and the size can vary from
+  // response-to-response. The client should concatenate them in the order
+  // received to construct the final payload.
+  rpc ReadMetadataBlobStream(ReadMetadataBlobStreamRequest)
+      returns (stream ReadMetadataBlobStreamResponse) {}
+
   // Writes an unencrypted data blob to the database.
   rpc WriteUnencryptedDataBlob(WriteUnencryptedDataBlobRequest)
       returns (WriteUnencryptedDataBlobResponse) {}

oak_private_memory/src/external_db_client.rs

@@ -15,21 +15,27 @@
 
 use anyhow::Context;
 use async_trait::async_trait;
+use futures::StreamExt;
 use log::info;
 use prost::Message;
 use sealed_memory_grpc_proto::oak::private_memory::sealed_memory_database_service_client::SealedMemoryDatabaseServiceClient;
 use sealed_memory_rust_proto::oak::private_memory::{
     DataBlob, DeleteBlobsRequest, EncryptedDataBlob, EncryptedMetadataBlob, MetadataBlob,
     ReadDataBlobRequest, ReadMetadataBlobRequest, ReadMetadataBlobResponse,
-    ReadUnencryptedDataBlobRequest, WriteDataBlobRequest, WriteMetadataBlobRequest,
-    WriteUnencryptedDataBlobRequest,
+    ReadMetadataBlobStreamRequest, ReadUnencryptedDataBlobRequest, WriteDataBlobRequest,
+    WriteMetadataBlobRequest, WriteMetadataBlobStreamRequest, WriteUnencryptedDataBlobRequest,
+    read_metadata_blob_stream_response, write_metadata_blob_stream_request,
 };
 use tonic::{Code, transport::Channel};
 
 pub type ExternalDbClient = SealedMemoryDatabaseServiceClient<Channel>;
 // The unique id for a opaque blob stored in the disk.
 pub type BlobId = String;
 
+// The size of the chunks to use when streaming data to/from the external
+// database. 1MB.
+const CHUNK_SIZE: usize = 1024 * 1024;
+
 // A non-fatal result from an attempt to persist the metadata db.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum MetadataPersistResult {
@@ -68,6 +74,15 @@ pub trait DataBlobHandler {
         &mut self,
         id: &BlobId,
     ) -> anyhow::Result<Option<EncryptedMetadataBlob>>;
+    async fn add_metadata_blob_stream(
+        &mut self,
+        id: &BlobId,
+        metadata_blob: EncryptedMetadataBlob,
+    ) -> anyhow::Result<MetadataPersistResult>;
+    async fn get_metadata_blob_stream(
+        &mut self,
+        id: &BlobId,
+    ) -> anyhow::Result<Option<EncryptedMetadataBlob>>;
     async fn get_blobs(
         &mut self,
         ids: &[BlobId],
@@ -144,6 +159,51 @@ impl DataBlobHandler for ExternalDbClient {
         Ok(MetadataPersistResult::Succeeded)
     }
 
+    async fn add_metadata_blob_stream(
+        &mut self,
+        id: &BlobId,
+        metadata_blob: EncryptedMetadataBlob,
+    ) -> anyhow::Result<MetadataPersistResult> {
+        let blob =
+            metadata_blob.encrypted_data_blob.as_ref().context("no blob contents")?.encode_to_vec();
+        let blob_size = blob.len() as u64;
+
+        let mut chunks = blob.chunks(CHUNK_SIZE);
+        let first_chunk = chunks.next().unwrap_or_default();
+        let mut messages = vec![WriteMetadataBlobStreamRequest {
+            request: Some(write_metadata_blob_stream_request::Request::MetadataBlob(
+                MetadataBlob {
+                    data_blob: Some(DataBlob { id: id.clone(), blob: first_chunk.to_vec() }),
+                    version: metadata_blob.version,
+                },
+            )),
+        }];
+
+        for chunk in chunks {
+            messages.push(WriteMetadataBlobStreamRequest {
+                request: Some(write_metadata_blob_stream_request::Request::Chunk(chunk.to_vec())),
+            });
+        }
+
+        let start_time = tokio::time::Instant::now();
+        let write_result = self.write_metadata_blob_stream(futures::stream::iter(messages)).await;
+
+        if let Err(ref status) = write_result {
+            if status.code() == Code::FailedPrecondition {
+                return Ok(MetadataPersistResult::RetryNeeded);
+            }
+            write_result?;
+        }
+
+        let mut elapsed_time = start_time.elapsed().as_millis() as u64;
+        if elapsed_time == 0 {
+            elapsed_time = 1;
+        }
+        let speed = blob_size / 1024 / elapsed_time;
+        metrics::get_global_metrics().record_db_save_speed(speed);
+        Ok(MetadataPersistResult::Succeeded)
+    }
+
     async fn add_blobs(
         &mut self,
         data_blobs: Vec<EncryptedDataBlob>,
@@ -238,6 +298,68 @@ impl DataBlobHandler for ExternalDbClient {
         }
     }
 
+    async fn get_metadata_blob_stream(
+        &mut self,
+        id: &BlobId,
+    ) -> anyhow::Result<Option<EncryptedMetadataBlob>> {
+        let start_time = tokio::time::Instant::now();
+        match self.read_metadata_blob_stream(ReadMetadataBlobStreamRequest { id: id.clone() }).await
+        {
+            Ok(response) => {
+                let mut response_stream = response.into_inner();
+                let mut full_blob = Vec::new();
+
+                let first_response = match response_stream.next().await {
+                    Some(res) => res?,
+                    None => return Ok(None),
+                };
+
+                let version = match first_response.response {
+                    Some(read_metadata_blob_stream_response::Response::MetadataBlob(metadata)) => {
+                        if let Some(data_blob) = metadata.data_blob {
+                            full_blob.extend_from_slice(&data_blob.blob);
+                        }
+                        metadata.version
+                    }
+                    _ => anyhow::bail!("Expected MetadataBlob as the first message in the stream"),
+                };
+
+                while let Some(response) = response_stream.next().await {
+                    let response = response?;
+                    match response.response {
+                        Some(read_metadata_blob_stream_response::Response::Chunk(chunk)) => {
+                            full_blob.extend_from_slice(&chunk);
+                        }
+                        _ => anyhow::bail!("Expected Chunk message in the stream"),
+                    }
+                }
+
+                let metadata_blob = EncryptedMetadataBlob {
+                    encrypted_data_blob: Some(
+                        EncryptedDataBlob::decode(&*full_blob)
+                            .context("Failed to decode EncryptedMetadataBlob")?,
+                    ),
+                    version,
+                };
+
+                let mut elapsed_time = start_time.elapsed().as_millis() as u64;
+                if elapsed_time == 0 {
+                    elapsed_time = 1;
+                }
+                let speed = full_blob.len() as u64 / 1024 / elapsed_time;
+                metrics::get_global_metrics().record_db_load_speed(speed);
+                Ok(Some(metadata_blob))
+            }
+            Err(status) => {
+                if status.code() == Code::NotFound {
+                    Ok(None)
+                } else {
+                    Err(status.into())
+                }
+            }
+        }
+    }
+
     async fn get_blobs(
         &mut self,
         ids: &[BlobId],

oak_private_memory/test/database/service.rs

@@ -23,12 +23,15 @@ use sealed_memory_grpc_proto::oak::private_memory::sealed_memory_database_servic
 use sealed_memory_rust_proto::oak::private_memory::{
     DataBlob, DeleteBlobsRequest, DeleteBlobsResponse, MetadataBlob, ReadDataBlobRequest,
     ReadDataBlobResponse, ReadMetadataBlobRequest, ReadMetadataBlobResponse,
-    ReadUnencryptedDataBlobRequest, ReadUnencryptedDataBlobResponse, ResetDatabaseRequest,
-    ResetDatabaseResponse, WriteDataBlobRequest, WriteDataBlobResponse, WriteMetadataBlobRequest,
-    WriteMetadataBlobResponse, WriteUnencryptedDataBlobRequest, WriteUnencryptedDataBlobResponse,
+    ReadMetadataBlobStreamRequest, ReadMetadataBlobStreamResponse, ReadUnencryptedDataBlobRequest,
+    ReadUnencryptedDataBlobResponse, ResetDatabaseRequest, ResetDatabaseResponse,
+    WriteDataBlobRequest, WriteDataBlobResponse, WriteMetadataBlobRequest,
+    WriteMetadataBlobResponse, WriteMetadataBlobStreamRequest, WriteMetadataBlobStreamResponse,
+    WriteUnencryptedDataBlobRequest, WriteUnencryptedDataBlobResponse,
+    read_metadata_blob_stream_response, write_metadata_blob_stream_request,
 };
 use tokio::{net::TcpListener, sync::Mutex};
-use tokio_stream::wrappers::TcpListenerStream;
+use tokio_stream::{StreamExt, wrappers::TcpListenerStream};
 
 pub struct SealedMemoryDatabaseServiceTestImpl {
     pub database: Mutex<HashMap<String, DataBlob>>,
@@ -140,6 +143,70 @@ impl SealedMemoryDatabaseService for SealedMemoryDatabaseServiceTestImpl {
         }
     }
 
+    type ReadMetadataBlobStreamStream = std::pin::Pin<
+        Box<
+            dyn tokio_stream::Stream<Item = Result<ReadMetadataBlobStreamResponse, tonic::Status>>
+                + Send,
+        >,
+    >;
+
+    async fn write_metadata_blob_stream(
+        &self,
+        request: tonic::Request<tonic::Streaming<WriteMetadataBlobStreamRequest>>,
+    ) -> Result<tonic::Response<WriteMetadataBlobStreamResponse>, tonic::Status> {
+        let mut stream = request.into_inner();
+        let mut full_blob = Vec::new();
+        let mut version = String::new();
+        let mut id = String::new();
+
+        while let Some(request) = stream.next().await {
+            let request = request?;
+            match request.request {
+                Some(write_metadata_blob_stream_request::Request::MetadataBlob(metadata)) => {
+                    version = metadata.version;
+                    if let Some(data_blob) = metadata.data_blob {
+                        id = data_blob.id;
+                        full_blob.extend_from_slice(&data_blob.blob);
+                    }
+                }
+                Some(write_metadata_blob_stream_request::Request::Chunk(chunk)) => {
+                    full_blob.extend_from_slice(&chunk);
+                }
+                None => {}
+            }
+        }
+
+        let added = self
+            .add_metadata_blob_inner(
+                id.clone(),
+                MetadataBlob { data_blob: Some(DataBlob { id, blob: full_blob }), version },
+            )
+            .await;
+        if !added {
+            Err(tonic::Status::failed_precondition("metadata blob version mismatch"))
+        } else {
+            Ok(tonic::Response::new(WriteMetadataBlobStreamResponse {}))
+        }
+    }
+
+    async fn read_metadata_blob_stream(
+        &self,
+        request: tonic::Request<ReadMetadataBlobStreamRequest>,
+    ) -> Result<tonic::Response<Self::ReadMetadataBlobStreamStream>, tonic::Status> {
+        let request = request.into_inner();
+        let blob = self.get_metdata_blob_inner(&request.id).await;
+
+        if let Some(blob) = blob {
+            let response = ReadMetadataBlobStreamResponse {
+                response: Some(read_metadata_blob_stream_response::Response::MetadataBlob(blob)),
+            };
+            let stream = tokio_stream::iter(vec![Ok(response)]);
+            Ok(tonic::Response::new(Box::pin(stream)))
+        } else {
+            Err(tonic::Status::not_found("Blob not found"))
+        }
+    }
+
     async fn write_unencrypted_data_blob(
         &self,
         request: tonic::Request<WriteUnencryptedDataBlobRequest>,