feat: add ComputeFile bao RPC method

This commit is contained in:
Derrick Hammer 2023-05-15 15:45:05 -04:00
parent 35878a2427
commit 687f26cc77
Signed by: pcfreak30
GPG Key ID: C997C339BE476FF2
4 changed files with 57 additions and 4 deletions

View File

@ -23,6 +23,7 @@ type Bao interface {
Write(id uint32, data []byte) error Write(id uint32, data []byte) error
Finalize(id uint32) ([]byte, error) Finalize(id uint32) ([]byte, error)
Destroy(id uint32) error Destroy(id uint32) error
ComputeFile(path string) ([]byte, error)
} }
func init() { func init() {
@ -88,8 +89,7 @@ func init() {
} }
func ComputeBaoTree(reader io.Reader) ([]byte, error) { func ComputeTreeStreaming(reader io.Reader) ([]byte, error) {
instance, err := baoInstance.Init() instance, err := baoInstance.Init()
if err != nil { if err != nil {
return nil, err return nil, err
@ -119,6 +119,15 @@ func ComputeBaoTree(reader io.Reader) ([]byte, error) {
} }
} }
func ComputeTreeFile(file *os.File) ([]byte, error) {
tree, err := baoInstance.ComputeFile(file.Name())
if err != nil {
return nil, err
}
return tree, nil
}
func write(instance uint32, bytes *[]byte) error { func write(instance uint32, bytes *[]byte) error {
err := baoInstance.Write(instance, *bytes) err := baoInstance.Write(instance, *bytes)
if err != nil { if err != nil {

View File

@ -45,3 +45,11 @@ func (g *GRPCClient) Destroy(id uint32) error {
return nil return nil
} }
func (g *GRPCClient) ComputeFile(path string) ([]byte, error) {
tree, err := g.client.ComputeFile(context.Background(), &wrappers.StringValue{Value: path})
if err != nil {
return nil, err
}
return tree.Value, nil
}

View File

@ -17,4 +17,5 @@ service bao {
rpc Write(WriteRequest) returns (google.protobuf.Empty); rpc Write(WriteRequest) returns (google.protobuf.Empty);
rpc Finalize (google.protobuf.UInt32Value) returns (google.protobuf.BytesValue); rpc Finalize (google.protobuf.UInt32Value) returns (google.protobuf.BytesValue);
rpc Destroy (google.protobuf.UInt32Value) returns (google.protobuf.Empty); rpc Destroy (google.protobuf.UInt32Value) returns (google.protobuf.Empty);
rpc ComputeFile (google.protobuf.StringValue) returns (google.protobuf.BytesValue);
} }

View File

@ -1,8 +1,11 @@
#![feature(async_fn_in_trait)] #![feature(async_fn_in_trait)]
#![allow(incomplete_features)] #![allow(incomplete_features)]
use io::Read;
use std::collections::hash_map::Entry; use std::collections::hash_map::Entry;
use std::collections::HashMap; use std::collections::HashMap;
use std::fs::{File};
use std::io;
use std::io::{Cursor, Write}; use std::io::{Cursor, Write};
use std::sync::{Arc}; use std::sync::{Arc};
@ -16,7 +19,7 @@ use tonic_health::server::HealthReporter;
use crate::proto::bao::bao_server::{Bao, BaoServer}; use crate::proto::bao::bao_server::{Bao, BaoServer};
use crate::proto::bao::WriteRequest; use crate::proto::bao::WriteRequest;
use crate::proto::google::protobuf::{BytesValue, Empty, UInt32Value}; use crate::proto::google::protobuf::{BytesValue, Empty, StringValue, UInt32Value};
use crate::unique_port::UniquePort; use crate::unique_port::UniquePort;
mod proto; mod proto;
@ -64,7 +67,7 @@ impl Bao for BaoService {
let next_id = self.counter.inc() as u32; let next_id = self.counter.inc() as u32;
let tree = Vec::new(); let tree = Vec::new();
let cursor = Cursor::new(tree); let cursor = Cursor::new(tree);
let encoder = Encoder::new(cursor); let encoder = Encoder::new_outboard(cursor);
let mut req = self.requests.lock(); let mut req = self.requests.lock();
req.insert(next_id, encoder); req.insert(next_id, encoder);
@ -113,4 +116,36 @@ impl Bao for BaoService {
Ok(Response::new(Empty::default())) Ok(Response::new(Empty::default()))
} }
async fn compute_file(&self, request: Request<StringValue>) -> Result<Response<BytesValue>, Status> {
let r = request.into_inner();
let tree = Vec::new();
let cursor = Cursor::new(tree);
let mut encoder = Encoder::new_outboard(cursor);
let mut input = File::open(r.value)?;
copy_reader_to_writer(&mut input, &mut encoder)?;
let ret = encoder.finalize().unwrap();
let bytes = ret.as_bytes().to_vec();
Ok(Response::new(BytesValue { value: bytes }))
}
}
fn copy_reader_to_writer(
reader: &mut impl Read,
writer: &mut impl Write,
) -> io::Result<u64> {
// At least 16 KiB is necessary to use AVX-512 with BLAKE3.
let mut buf = [0; 65536];
let mut written = 0;
loop {
let len = match reader.read(&mut buf) {
Ok(0) => return Ok(written),
Ok(len) => len,
Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => return Err(e),
};
writer.write_all(&buf[..len])?;
written += len as u64;
}
} }