make ComputeConnectBackend dyn

This commit is contained in:
Conrad Ludgate
2024-10-24 11:55:31 +01:00
parent 2d34fec39b
commit c8108a4b84
4 changed files with 26 additions and 39 deletions

View File

@@ -389,7 +389,7 @@ pub struct ControlPlaneComputeBackend<'a> {
}
#[async_trait::async_trait]
impl ComputeConnectBackend for ControlPlaneComputeBackend<'_> {
impl ComputeConnectBackend for ControlPlaneComputeBackend<'static> {
async fn wake_compute(
&self,
ctx: &RequestMonitoring,

View File

@@ -56,7 +56,7 @@ pub(crate) trait ConnectMechanism {
}
#[async_trait]
pub(crate) trait ComputeConnectBackend {
pub(crate) trait ComputeConnectBackend: Send + Sync + 'static {
async fn wake_compute(
&self,
ctx: &RequestMonitoring,
@@ -98,10 +98,10 @@ impl ConnectMechanism for TcpMechanism<'_> {
/// Try to connect to the compute node, retrying if necessary.
#[tracing::instrument(skip_all)]
pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: ComputeConnectBackend>(
pub(crate) async fn connect_to_compute<M: ConnectMechanism>(
ctx: &RequestMonitoring,
mechanism: &M,
user_info: &B,
user_info: &dyn ComputeConnectBackend,
allow_self_signed_compute: bool,
wake_compute_retry_config: RetryConfig,
connect_to_compute_retry_config: RetryConfig,

View File

@@ -11,10 +11,10 @@ use crate::metrics::{
};
use crate::proxy::retry::{retry_after, should_retry};
pub(crate) async fn wake_compute<B: ComputeConnectBackend>(
pub(crate) async fn wake_compute(
num_retries: &mut u32,
ctx: &RequestMonitoring,
api: &B,
api: &dyn ComputeConnectBackend,
config: RetryConfig,
) -> Result<CachedNodeInfo, WakeComputeError> {
let retry_type = RetryType::WakeCompute;

View File

@@ -30,7 +30,7 @@ use crate::control_plane::provider::ApiLockError;
use crate::control_plane::{Api, CachedNodeInfo};
use crate::error::{ErrorKind, ReportableError, UserFacingError};
use crate::intern::EndpointIdInt;
use crate::proxy::connect_compute::ConnectMechanism;
use crate::proxy::connect_compute::{ComputeConnectBackend, ConnectMechanism};
use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute};
use crate::rate_limiter::EndpointRateLimiter;
use crate::types::{EndpointId, Host};
@@ -190,40 +190,27 @@ impl PoolingBackend {
tracing::Span::current().record("conn_id", display(conn_id));
info!(%conn_id, "pool: opening a new connection '{conn_info}'");
match &self.auth_backend {
let api = match &self.auth_backend {
ServerlessBackend::ControlPlane(cplane) => {
crate::proxy::connect_compute::connect_to_compute(
ctx,
&TokioMechanism {
conn_id,
conn_info,
pool: self.pool.clone(),
locks: &self.config.connect_compute_locks,
},
&cplane.attach_to_credentials(keys),
false, // do not allow self signed compute for http flow
self.config.wake_compute_retry_config,
self.config.connect_to_compute_retry_config,
)
.await
&cplane.attach_to_credentials(keys) as &dyn ComputeConnectBackend
}
ServerlessBackend::Local(local_proxy) => {
crate::proxy::connect_compute::connect_to_compute(
ctx,
&TokioMechanism {
conn_id,
conn_info,
pool: self.pool.clone(),
locks: &self.config.connect_compute_locks,
},
&**local_proxy,
false, // do not allow self signed compute for http flow
self.config.wake_compute_retry_config,
self.config.connect_to_compute_retry_config,
)
.await
}
}
ServerlessBackend::Local(local_proxy) => &**local_proxy as &dyn ComputeConnectBackend,
};
crate::proxy::connect_compute::connect_to_compute(
ctx,
&TokioMechanism {
conn_id,
conn_info,
pool: self.pool.clone(),
locks: &self.config.connect_compute_locks,
},
api,
false, // do not allow self signed compute for http flow
self.config.wake_compute_retry_config,
self.config.connect_to_compute_retry_config,
)
.await
}
// Wake up the destination if needed