chore: add direct connect handshake (#483)

Momento direct connect uses an http+json -> protosocket negotiation
handshake. The http connection retrieves a list of socket addresses
that can be used with the connection-oriented protosocket rpc
protocol.

The ProtosocketCacheClient queries the address list on initial
connection, and updates the local list in the background. When a new
connection is required, for example due to a disconnect, the local
list is used to choose a new socket address for that "slot."

Direct connections are lighter than grpc connections, but currently
have some compromises. For example, they do not currently support
graceful GOAWAY mechanism. Also, the cached local address list may
be out of date with reality for a few seconds, causing some possible
re-connect attempts to fail until list refresh.

This connection scheme will be improved over time, and is expected
to outperform grpc for both cost & latency now, and match grpc
practically for resilience in the future.

Author: kvcache

Cargo.lock

@@ -36,12 +36,14 @@ base64 = "0.22"
 derive_more = { version = "2.0.1", features = ["full"] }
 futures = "0"
 h2 = { version = "0.4" }
+http = { version = "1" }
 hyper = { version = "1.6" }
 log = "0.4"
 protosocket = "0.11.0"
 protosocket-prost = "0.11.0"
 protosocket-rpc = "0.11.0"
 rand = "0.9"
+reqwest = { version = "0.12", default-features = false, features = ["http2", "json", "rustls-tls-native-roots"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 thiserror = "2.0"

Cargo.toml

@@ -11,7 +11,7 @@ struct V1Token {
     pub endpoint: String,
 }
 
-#[derive(PartialEq, Eq, Clone)]
+#[derive(PartialEq, Eq, Clone, Debug)]
 pub(crate) enum EndpointSecurity {
     Insecure,
     Unverified,
@@ -25,6 +25,7 @@ pub struct CredentialProvider {
     pub(crate) auth_token: String,
     pub(crate) control_endpoint: String,
     pub(crate) cache_endpoint: String,
+    pub(crate) cache_http_endpoint: String,
     pub(crate) token_endpoint: String,
     pub(crate) endpoint_security: EndpointSecurity,
 }
@@ -46,6 +47,7 @@ impl Debug for CredentialProvider {
             .field("cache_endpoint", &self.cache_endpoint)
             .field("control_endpoint", &self.control_endpoint)
             .field("token_endpoint", &self.token_endpoint)
+            .field("endpoint_security", &self.endpoint_security)
             .finish()
     }
 }
@@ -83,6 +85,11 @@ impl CredentialProvider {
         decode_auth_token(token_to_process)
     }
 
+    /// Returns the hostname that can be used with momento HTTP apis
+    pub fn cache_http_endpoint(&self) -> &str {
+        &self.cache_http_endpoint
+    }
+
     /// Returns a Credential Provider from the provided API key
     ///
     /// # Arguments
@@ -175,6 +182,7 @@ fn process_v1_token(auth_token_bytes: Vec<u8>) -> MomentoResult<CredentialProvid
     Ok(CredentialProvider {
         auth_token: json.api_key,
         cache_endpoint: https_endpoint(get_cache_endpoint(&json.endpoint)),
+        cache_http_endpoint: https_endpoint(get_cache_http_endpoint(&json.endpoint)),
         control_endpoint: https_endpoint(get_control_endpoint(&json.endpoint)),
         token_endpoint: https_endpoint(get_token_endpoint(&json.endpoint)),
         endpoint_security: EndpointSecurity::Tls,
@@ -185,6 +193,10 @@ fn get_cache_endpoint(endpoint: &str) -> String {
     format!("cache.{endpoint}")
 }
 
+fn get_cache_http_endpoint(endpoint: &str) -> String {
+    format!("api.cache.{endpoint}")
+}
+
 fn get_control_endpoint(endpoint: &str) -> String {
     format!("control.{endpoint}")
 }

src/credential_provider.rs

@@ -0,0 +1,127 @@
+use std::{
+    collections::HashMap,
+    net::SocketAddr,
+    sync::{Arc, Mutex},
+};
+
+use crate::CredentialProvider;
+
+// Todo: should make the connections try to balance better than random
+#[derive(serde::Deserialize, serde::Serialize, Debug, Default)]
+pub(crate) struct Addresses {
+    #[serde(flatten)]
+    azs: HashMap<AzId, Vec<Address>>,
+}
+
+impl Addresses {
+    /// Get a list of socket addresses, optionally filtered by availability zone ID.
+    /// If an az_id is provided, only addresses in that availability zone will be returned,
+    pub fn for_az(&self, az_id: Option<&str>) -> Vec<SocketAddr> {
+        if let Some(az_id) = az_id {
+            if let Some(addresses) = self.azs.get(&AzId(az_id.to_string())) {
+                if !addresses.is_empty() {
+                    return addresses.iter().map(|a| a.socket_address).collect();
+                }
+            }
+        }
+        self.azs
+            .values()
+            .flat_map(|addresses| addresses.iter().map(|a| a.socket_address))
+            .collect()
+    }
+}
+
+#[derive(serde::Deserialize, serde::Serialize, Debug, Eq, PartialEq, Hash)]
+pub(crate) struct AzId(String);
+
+#[derive(serde::Deserialize, serde::Serialize, Debug)]
+pub(crate) struct Address {
+    socket_address: SocketAddr,
+}
+
+#[derive(Debug)]
+pub(crate) struct AddressProvider {
+    addresses: Mutex<Arc<Addresses>>,
+    client: reqwest::Client,
+    credential_provider: CredentialProvider,
+}
+
+#[derive(Debug, thiserror::Error)]
+pub(crate) enum RefreshError {
+    Reqwest(#[from] reqwest::Error),
+    Json(#[from] serde_json::Error),
+    Uri(#[from] http::uri::InvalidUri),
+    BadStatus((reqwest::StatusCode, String)),
+}
+impl std::fmt::Display for RefreshError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            RefreshError::Reqwest(e) => write!(f, "Reqwest error: {e}"),
+            RefreshError::Json(e) => write!(f, "JSON error: {e}"),
+            RefreshError::Uri(e) => write!(f, "URI error: {e}"),
+            RefreshError::BadStatus((status, text)) => write!(f, "Bad status: {status}, {text}"),
+        }
+    }
+}
+
+impl AddressProvider {
+    /// Looks for an address list from the provided endpoint.
+    #[allow(clippy::expect_used)]
+    pub fn new(credential_provider: CredentialProvider) -> Self {
+        let client = reqwest::Client::builder()
+            .tls_built_in_native_certs(true)
+            .tls_built_in_root_certs(true)
+            .build()
+            .expect("must be able to build client");
+        Self {
+            addresses: Default::default(),
+            client,
+            credential_provider,
+        }
+    }
+
+    #[allow(clippy::expect_used)]
+    pub fn get_addresses(&self) -> impl std::ops::Deref<Target = Addresses> {
+        self.addresses
+            .lock()
+            .expect("local mutex must not be poisoned")
+            .clone()
+    }
+
+    #[allow(clippy::expect_used)]
+    pub async fn try_refresh_addresses(&self) -> Result<(), RefreshError> {
+        let request = self
+            .client
+            .get(format!(
+                "{}/endpoints",
+                self.credential_provider
+                    .cache_http_endpoint()
+                    .trim_end_matches('/')
+            ))
+            .header("authorization", &self.credential_provider.auth_token)
+            .build()?;
+        let response = self.client.execute(request).await?;
+
+        if !response.status().is_success() {
+            let status = response.status();
+            let text = response.text().await?;
+            return Err(RefreshError::BadStatus((status, text)));
+        }
+
+        let response = response.text().await?;
+        let addresses = match serde_json::from_str(&response) {
+            Ok(addresses) => addresses,
+            Err(e) => {
+                log::warn!("error parsing address list JSON: {response}");
+                return Err(RefreshError::Json(e));
+            }
+        };
+        log::debug!("refreshed address list: {addresses:?}");
+        let addresses = Arc::new(addresses);
+        *self
+            .addresses
+            .lock()
+            .expect("local mutex must not be poisoned") = addresses;
+        Ok(())
+    }
+}

src/protosocket/cache/address_provider.rs

@@ -3,7 +3,7 @@ use protosocket_rpc::client::{ConnectionPool, RpcClient};
 
 use crate::cache::{GetRequest, SetRequest};
 use crate::protosocket::cache::cache_client_builder::NeedsDefaultTtl;
-use crate::protosocket::cache::utils::ProtosocketConnectionManager;
+use crate::protosocket::cache::connection_manager::ProtosocketConnectionManager;
 use crate::protosocket::cache::{Configuration, MomentoProtosocketRequest};
 use crate::{utils, IntoBytes, MomentoError, MomentoResult, ProtosocketCacheClientBuilder};
 use std::convert::TryInto;
@@ -103,7 +103,7 @@ impl ProtosocketCacheClient {
     /// # }
     /// ```
     pub fn builder() -> ProtosocketCacheClientBuilder<NeedsDefaultTtl> {
-        ProtosocketCacheClientBuilder(NeedsDefaultTtl(()))
+        super::cache_client_builder::initial()
     }
 
     /// Gets an item from a Momento Cache
@@ -211,7 +211,7 @@ impl ProtosocketCacheClient {
         &self,
     ) -> MomentoResult<RpcClient<CacheCommand, CacheResponse>> {
         let pooled_client = self.client_pool.get_connection().await.map_err(|e| {
-            MomentoError::unknown_error("protosocket_connection", Some(e.to_string()))
+            MomentoError::unknown_error("protosocket_connection", Some(format!("{e:?}")))
         })?;
         Ok(pooled_client.clone())
     }

src/protosocket/cache/cache_client.rs

@@ -1,4 +1,4 @@
-use crate::protosocket::cache::utils::ProtosocketConnectionManager;
+use crate::protosocket::cache::connection_manager::ProtosocketConnectionManager;
 use crate::protosocket::cache::Configuration;
 use crate::{CredentialProvider, MomentoResult, ProtosocketCacheClient};
 use momento_protos::protosocket::cache::CacheCommand;
@@ -11,11 +11,19 @@ pub type Serializer = ProstSerializer<CacheResponse, CacheCommand>;
 
 /// The initial state of the ProtosocketCacheClientBuilder.
 #[derive(PartialEq, Eq, Clone, Debug)]
-pub struct ProtosocketCacheClientBuilder<State>(pub State);
+pub struct ProtosocketCacheClientBuilder<State> {
+    state: State,
+}
+
+pub(crate) fn initial() -> ProtosocketCacheClientBuilder<NeedsDefaultTtl> {
+    ProtosocketCacheClientBuilder {
+        state: NeedsDefaultTtl,
+    }
+}
 
 /// The state of the ProtosocketCacheClientBuilder when it is waiting for a default TTL.
 #[derive(PartialEq, Eq, Clone, Debug)]
-pub struct NeedsDefaultTtl(pub ());
+pub struct NeedsDefaultTtl;
 
 /// The state of the ProtosocketCacheClientBuilder when it is waiting for a configuration.
 #[derive(PartialEq, Eq, Clone, Debug)]
@@ -53,7 +61,9 @@ impl ProtosocketCacheClientBuilder<NeedsDefaultTtl> {
         self,
         default_ttl: Duration,
     ) -> ProtosocketCacheClientBuilder<NeedsConfiguration> {
-        ProtosocketCacheClientBuilder(NeedsConfiguration { default_ttl })
+        ProtosocketCacheClientBuilder {
+            state: NeedsConfiguration { default_ttl },
+        }
     }
 }
 
@@ -63,10 +73,12 @@ impl ProtosocketCacheClientBuilder<NeedsConfiguration> {
         self,
         configuration: impl Into<Configuration>,
     ) -> ProtosocketCacheClientBuilder<NeedsCredentialProvider> {
-        ProtosocketCacheClientBuilder(NeedsCredentialProvider {
-            default_ttl: self.0.default_ttl,
-            configuration: configuration.into(),
-        })
+        ProtosocketCacheClientBuilder {
+            state: NeedsCredentialProvider {
+                default_ttl: self.state.default_ttl,
+                configuration: configuration.into(),
+            },
+        }
     }
 }
 
@@ -76,11 +88,13 @@ impl ProtosocketCacheClientBuilder<NeedsCredentialProvider> {
         self,
         credential_provider: CredentialProvider,
     ) -> ProtosocketCacheClientBuilder<NeedsRuntime> {
-        ProtosocketCacheClientBuilder(NeedsRuntime {
-            default_ttl: self.0.default_ttl,
-            configuration: self.0.configuration,
-            credential_provider,
-        })
+        ProtosocketCacheClientBuilder {
+            state: NeedsRuntime {
+                default_ttl: self.state.default_ttl,
+                configuration: self.state.configuration,
+                credential_provider,
+            },
+        }
     }
 }
 
@@ -90,28 +104,38 @@ impl ProtosocketCacheClientBuilder<NeedsRuntime> {
         self,
         runtime: tokio::runtime::Handle,
     ) -> ProtosocketCacheClientBuilder<ReadyToBuild> {
-        ProtosocketCacheClientBuilder(ReadyToBuild {
-            default_ttl: self.0.default_ttl,
-            runtime,
-            credential_provider: self.0.credential_provider,
-            configuration: self.0.configuration,
-        })
+        ProtosocketCacheClientBuilder {
+            state: ReadyToBuild {
+                default_ttl: self.state.default_ttl,
+                runtime,
+                credential_provider: self.state.credential_provider,
+                configuration: self.state.configuration,
+            },
+        }
     }
 }
 
 impl ProtosocketCacheClientBuilder<ReadyToBuild> {
     /// Constructs a new CacheClientBuilder in the ReadyToBuild state.
     pub async fn build(self) -> MomentoResult<ProtosocketCacheClient> {
-        let client_connector =
-            ProtosocketConnectionManager::new(self.0.credential_provider, self.0.runtime)?;
+        let ReadyToBuild {
+            default_ttl,
+            credential_provider,
+            runtime,
+            configuration,
+        } = self.state;
+        let client_connector = ProtosocketConnectionManager::new(
+            credential_provider,
+            runtime,
+            configuration.az_id.clone(),
+        )?;
 
-        let client_pool =
-            ConnectionPool::new(client_connector, self.0.configuration.connection_count());
+        let client_pool = ConnectionPool::new(client_connector, configuration.connection_count());
 
         Ok(ProtosocketCacheClient::new(
             client_pool,
-            self.0.default_ttl,
-            self.0.configuration,
+            default_ttl,
+            configuration,
         ))
     }
 }