Skip to content

Instantly share code, notes, and snippets.

@joelgallant
Last active September 6, 2023 08:59
Show Gist options
  • Save joelgallant/f0fe9f593728b106daddbf7fa00e10af to your computer and use it in GitHub Desktop.
Save joelgallant/f0fe9f593728b106daddbf7fa00e10af to your computer and use it in GitHub Desktop.
// GeneratedClient here is codegen-ed
#[tonic::async_trait]
impl TonicClient for GeneratedClient<Channel> {
async fn try_connect(
dest: tonic::transport::Endpoint,
) -> Result<Self, tonic::transport::Error> {
GeneratedClient::connect(dest).await
}
}
use std::sync::Arc;
use tokio::sync::RwLock;
use tonic::transport::Endpoint;
#[tonic::async_trait]
pub trait TonicClient
where
Self: Sized,
{
async fn try_connect(dest: Endpoint) -> Result<Self, tonic::transport::Error>;
}
#[derive(Debug, Clone)]
enum LazyClientImpl<Client> {
NotConnected(Box<Endpoint>),
Connected(Client),
}
// We use an Arc<RwLock<_>> for two reasons
// 1. Arc so that clones are trivial and synchronous
// 2. RwLock so that interior mutations can "propagate" across all copies of this Client
//
// This ensures that we don't do the handshaking work more than 1 time ever.
// While is was technically possible to make a LazyClient before without the RwLock, after
// a clone on that client, connecting from LazyClient variable `a` would be separate from connection from LazyClient variable `b`.
//
// In other words, it's inherently necessary to mutate &self, but we want concurrency, so we use an RwLock for optimistic reads.
#[derive(Debug, Clone)]
pub struct LazyClient<Client>(Arc<RwLock<LazyClientImpl<Client>>>);
impl<Client> LazyClientImpl<Client>
where
Client: TonicClient + Clone,
{
fn create<Dest: Into<Endpoint>>(dest: Dest) -> Self {
Self::NotConnected(Box::new(dest.into()))
}
}
impl<Client> LazyClient<Client>
where
Client: TonicClient + Clone,
{
pub fn create<Dest: Into<Endpoint>>(dest: Dest) -> Self {
Self(Arc::new(RwLock::new(LazyClientImpl::create(dest))))
}
pub async fn conn(&self) -> Result<Client, tonic::transport::Error> {
if let LazyClientImpl::Connected(client) = &*self.0.read().await {
// fast path, don't need a write lock
return Ok(client.clone());
}
log::debug!("LazyClient intializing first connection...");
let mut inner_write = self.0.write().await;
match &*inner_write {
LazyClientImpl::Connected(_) => {
unreachable!("we determined above that self is NotConnected")
}
LazyClientImpl::NotConnected(dest) => {
let client = Client::try_connect(*dest.clone()).await?;
*inner_write = LazyClientImpl::Connected(client.clone());
Ok(client)
}
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment