mirror of
https://github.com/neondatabase/neon.git
synced 2026-05-28 10:30:40 +00:00
3ca47bc37b
Split the modules responsible for passing data and connecting to compute from auth and waking for PGLB. This PR just moves files. The waking is going to get removed from pglb after this.
104 lines
3.2 KiB
Rust
104 lines
3.2 KiB
Rust
use futures::FutureExt;
|
|
use smol_str::SmolStr;
|
|
use tokio::io::{AsyncRead, AsyncWrite};
|
|
use tracing::debug;
|
|
use utils::measured_stream::MeasuredStream;
|
|
|
|
use super::copy_bidirectional::ErrorSource;
|
|
use crate::cancellation;
|
|
use crate::compute::PostgresConnection;
|
|
use crate::config::ComputeConfig;
|
|
use crate::control_plane::messages::MetricsAuxInfo;
|
|
use crate::metrics::{Direction, Metrics, NumClientConnectionsGuard, NumConnectionRequestsGuard};
|
|
use crate::stream::Stream;
|
|
use crate::usage_metrics::{Ids, MetricCounterRecorder, USAGE_METRICS};
|
|
|
|
/// Forward bytes in both directions (client <-> compute).
|
|
#[tracing::instrument(skip_all)]
|
|
pub(crate) async fn proxy_pass(
|
|
client: impl AsyncRead + AsyncWrite + Unpin,
|
|
compute: impl AsyncRead + AsyncWrite + Unpin,
|
|
aux: MetricsAuxInfo,
|
|
private_link_id: Option<SmolStr>,
|
|
) -> Result<(), ErrorSource> {
|
|
// we will report ingress at a later date
|
|
let usage_tx = USAGE_METRICS.register(Ids {
|
|
endpoint_id: aux.endpoint_id,
|
|
branch_id: aux.branch_id,
|
|
private_link_id,
|
|
});
|
|
|
|
let metrics = &Metrics::get().proxy.io_bytes;
|
|
let m_sent = metrics.with_labels(Direction::Tx);
|
|
let mut client = MeasuredStream::new(
|
|
client,
|
|
|_| {},
|
|
|cnt| {
|
|
// Number of bytes we sent to the client (outbound).
|
|
metrics.get_metric(m_sent).inc_by(cnt as u64);
|
|
usage_tx.record_egress(cnt as u64);
|
|
},
|
|
);
|
|
|
|
let m_recv = metrics.with_labels(Direction::Rx);
|
|
let mut compute = MeasuredStream::new(
|
|
compute,
|
|
|_| {},
|
|
|cnt| {
|
|
// Number of bytes the client sent to the compute node (inbound).
|
|
metrics.get_metric(m_recv).inc_by(cnt as u64);
|
|
usage_tx.record_ingress(cnt as u64);
|
|
},
|
|
);
|
|
|
|
// Starting from here we only proxy the client's traffic.
|
|
debug!("performing the proxy pass...");
|
|
let _ = crate::pglb::copy_bidirectional::copy_bidirectional_client_compute(
|
|
&mut client,
|
|
&mut compute,
|
|
)
|
|
.await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub(crate) struct ProxyPassthrough<S> {
|
|
pub(crate) client: Stream<S>,
|
|
pub(crate) compute: PostgresConnection,
|
|
pub(crate) aux: MetricsAuxInfo,
|
|
pub(crate) session_id: uuid::Uuid,
|
|
pub(crate) private_link_id: Option<SmolStr>,
|
|
pub(crate) cancel: cancellation::Session,
|
|
|
|
pub(crate) _req: NumConnectionRequestsGuard<'static>,
|
|
pub(crate) _conn: NumClientConnectionsGuard<'static>,
|
|
}
|
|
|
|
impl<S: AsyncRead + AsyncWrite + Unpin> ProxyPassthrough<S> {
|
|
pub(crate) async fn proxy_pass(
|
|
self,
|
|
compute_config: &ComputeConfig,
|
|
) -> Result<(), ErrorSource> {
|
|
let res = proxy_pass(
|
|
self.client,
|
|
self.compute.stream,
|
|
self.aux,
|
|
self.private_link_id,
|
|
)
|
|
.await;
|
|
if let Err(err) = self
|
|
.compute
|
|
.cancel_closure
|
|
.try_cancel_query(compute_config)
|
|
.boxed()
|
|
.await
|
|
{
|
|
tracing::warn!(session_id = ?self.session_id, ?err, "could not cancel the query in the database");
|
|
}
|
|
|
|
drop(self.cancel.remove_cancel_key()); // we don't need a result. If the queue is full, we just log the error
|
|
|
|
res
|
|
}
|
|
}
|