rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-07-07 06:00:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP: Revert the MeasuredStream changes

Browse Source 
To demonstrate why the MeasuredStream change makes things simpler.

This doesn't compile. Not that I tried very hard, but I couldn't
figure out how to fill in the generic type parameters to make this
compile.
This commit is contained in:
Heikki Linnakangas
2023-01-26 00:34:49 +02:00
parent 95fd68d76e
commit 1e1816676c
2 changed files with 18 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
proxy/src/proxy.rs
View File

@@ -347,10 +347,16 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> EstablishedConnection<'_, S> {
.await?;
let m_sent = NUM_BYTES_PROXIED_COUNTER.with_label_values(&node.aux.traffic_labels("tx"));
let client_stream = MeasuredStream::new(stream.into_inner(), m_sent);
let client_stream = MeasuredStream::new(stream.into_inner(), |cnt| {
// Number of bytes we sent to the client (outbound).
m_sent.inc_by(cnt as u64);
});
let m_recv = NUM_BYTES_PROXIED_COUNTER.with_label_values(&node.aux.traffic_labels("rx"));
let db_stream = MeasuredStream::new(db.stream, m_recv);
let db_stream = MeasuredStream::new(db.stream, |cnt| {
// Number of bytes the client sent to the compute node (inbound).
m_recv.inc_by(cnt as u64);
});
Ok(EstablishedConnection {
client_stream,

20
proxy/src/stream.rs
View File

@@ -231,27 +231,27 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for Stream<S> {
}
pin_project! {
/// This stream tracks all writes, and whenever the stream is flushed,
/// increments the user-provided counter by the number of bytes flushed.
pub struct MeasuredStream<S> {
/// This stream tracks all writes and calls user provided
/// callback when the underlying stream is flushed.
pub struct MeasuredStream<S, W> {
#[pin]
stream: S,
write_count: usize,
write_counter: prometheus::IntCounter,
inc_write_count: W,
}
}
impl<S> MeasuredStream<S> {
pub fn new(stream: S, write_counter: prometheus::IntCounter) -> Self {
impl<S, W> MeasuredStream<S, W> {
pub fn new(stream: S, inc_write_count: W) -> Self {
Self {
stream,
write_count: 0,
write_counter,
inc_write_count,
}
}
}
impl<S: AsyncRead + Unpin> AsyncRead for MeasuredStream<S> {
impl<S: AsyncRead + Unpin, W> AsyncRead for MeasuredStream<S, W> {
fn poll_read(
self: Pin<&mut Self>,
context: &mut task::Context<'_>,
@@ -261,7 +261,7 @@ impl<S: AsyncRead + Unpin> AsyncRead for MeasuredStream<S> {
}
}
impl<S: AsyncWrite + Unpin> AsyncWrite for MeasuredStream<S> {
impl<S: AsyncWrite + Unpin, W: FnMut(usize)> AsyncWrite for MeasuredStream<S, W> {
fn poll_write(
self: Pin<&mut Self>,
context: &mut task::Context<'_>,
@@ -282,7 +282,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for MeasuredStream<S> {
let this = self.project();
this.stream.poll_flush(context).map_ok(|()| {
// Call the user provided callback and reset the write count.
this.write_counter.inc_by(*this.write_count as u64);
(this.inc_write_count)(*this.write_count);
*this.write_count = 0;
})
}