Skip to content

Commit

Permalink
feat: RemoteTransactionProver lazy connect (#937)
Browse files Browse the repository at this point in the history
* feat: connect lazily on prove

* chore: update CHANGELOG

* feat: create separate makefile install command for testing

* fix: wasm build

* feat: mantain connection between prove calls

* fix: remove new line

* fix: address suggestions

* fix: update makefile doc comment

---------

Co-authored-by: Bobbin Threadbare <[email protected]>
  • Loading branch information
tomyrd and bobbinth authored Oct 29, 2024
1 parent e784448 commit ede0b78
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 20 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

## 0.6.0 (TBD)

- Changed `RemoteTransactionProver` to lazily connect on prove (#937).
- [BREAKING] Made `TransactionMastForest` and `BasicAuthenticator` be `Send`, `Sync`; made scripts lazily initialize (#939).
- Added `RemoteTransactionProver` struct to `miden-tx-prover` (#921).
- [BREAKING] Changed `TransactionProver` trait to be `maybe_async_trait` based on the `async` feature (#913).
Expand Down
4 changes: 4 additions & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -106,3 +106,7 @@ bench-tx: ## Run transaction benchmarks
.PHONY: install-prover
install-prover: ## Installs prover
cargo install --path bin/tx-prover --locked

.PHONY: install-prover-testing
install-prover-testing: ## Installs prover intended for testing purposes
cargo install --path bin/tx-prover --locked --features testing
13 changes: 13 additions & 0 deletions bin/tx-prover/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,3 +22,16 @@ pub enum RemoteTransactionProverError {
/// Indicates that the connection to the server failed.
ConnectionFailed(String),
}

impl std::fmt::Display for RemoteTransactionProverError {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
match self {
RemoteTransactionProverError::InvalidEndpoint(endpoint) => {
write!(f, "Invalid endpoint: {}", endpoint)
},
RemoteTransactionProverError::ConnectionFailed(endpoint) => {
write!(f, "Failed to connect to transaction prover at: {}", endpoint)
},
}
}
}
70 changes: 50 additions & 20 deletions bin/tx-prover/src/prover.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
use alloc::{boxed::Box, string::ToString};
use alloc::{
boxed::Box,
string::{String, ToString},
};
use core::cell::RefCell;

use miden_objects::transaction::{ProvenTransaction, TransactionWitness};
Expand All @@ -11,39 +14,57 @@ use crate::{generated::api_client::ApiClient, RemoteTransactionProverError};

/// A [RemoteTransactionProver] is a transaction prover that sends witness data to a remote
/// gRPC server and receives a proven transaction.
///
/// When compiled for the `wasm32-unknown-unknown` target, it uses the `tonic_web_wasm_client`
/// transport. Otherwise, it uses the built-in `tonic::transport` for native platforms.
///
/// The transport layer connection is established lazily when the first transaction is proven.
#[derive(Clone)]
pub struct RemoteTransactionProver {
#[cfg(target_arch = "wasm32")]
client: RefCell<ApiClient<tonic_web_wasm_client::Client>>,
client: RefCell<Option<ApiClient<tonic_web_wasm_client::Client>>>,

#[cfg(not(target_arch = "wasm32"))]
client: RefCell<ApiClient<tonic::transport::Channel>>,
client: RefCell<Option<ApiClient<tonic::transport::Channel>>>,

endpoint: String,
}

impl RemoteTransactionProver {
/// Creates a new [RemoteTransactionProver] with the specified gRPC server endpoint.
/// This instantiates a tonic client that attempts connecting with the server.
///
/// When compiled for the `wasm32-unknown-unknown` target, it uses the `tonic_web_wasm_client`
/// transport. Otherwise, it uses the built-in `tonic::transport` for native platforms.
///
/// # Errors
///
/// This function will return an error if the endpoint is invalid or if the gRPC
/// connection to the server cannot be established.
pub async fn new(endpoint: &str) -> Result<Self, RemoteTransactionProverError> {
/// Creates a new [RemoteTransactionProver] with the specified gRPC server endpoint. The
/// endpoint should be in the format `{protocol}://{hostname}:{port}`.
pub fn new(endpoint: &str) -> Self {
RemoteTransactionProver {
endpoint: endpoint.to_string(),
client: RefCell::new(None),
}
}

/// Establishes a connection to the remote transaction prover server. The connection is
/// mantained for the lifetime of the prover. If the connection is already established, this
/// method does nothing.
async fn connect(&self) -> Result<(), RemoteTransactionProverError> {
let mut client = self.client.borrow_mut();
if client.is_some() {
return Ok(());
}

#[cfg(target_arch = "wasm32")]
let client = {
let web_client = tonic_web_wasm_client::Client::new(endpoint.to_string());
let new_client = {
let web_client = tonic_web_wasm_client::Client::new(self.endpoint.clone());
ApiClient::new(web_client)
};

#[cfg(not(target_arch = "wasm32"))]
let client = ApiClient::connect(endpoint.to_string())
.await
.map_err(|_| RemoteTransactionProverError::ConnectionFailed(endpoint.to_string()))?;
let new_client = {
ApiClient::connect(self.endpoint.clone()).await.map_err(|_| {
RemoteTransactionProverError::ConnectionFailed(self.endpoint.to_string())
})?
};

Ok(RemoteTransactionProver { client: RefCell::new(client) })
*client = Some(new_client);

Ok(())
}
}

Expand All @@ -54,13 +75,22 @@ impl TransactionProver for RemoteTransactionProver {
tx_witness: TransactionWitness,
) -> Result<ProvenTransaction, TransactionProverError> {
use miden_objects::utils::Serializable;
self.connect().await.map_err(|err| {
TransactionProverError::InternalError(format!(
"Failed to connect to the remote prover: {}",
err
))
})?;

let mut client = self.client.borrow_mut();

let request = tonic::Request::new(crate::generated::ProveTransactionRequest {
transaction_witness: tx_witness.to_bytes(),
});

let response = client
.as_mut()
.expect("client should be connected")
.prove_transaction(request)
.await
.map_err(|err| TransactionProverError::InternalError(err.to_string()))?;
Expand Down

0 comments on commit ede0b78

Please sign in to comment.