Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 036a1b4

Browse files
mempiratemempirate
authored
Merge pull request #88 from 7suyash7/main
feat(sub): add socket-wide stats including dropped messages
2 parents d0d6a80 + a272b5e commit 036a1b4

19 files changed

Lines changed: 192 additions & 95 deletions

File tree

‎msg-socket/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ use tokio::io::{AsyncRead, AsyncWrite};
77

88
use msg_transport::Address;
99

10+
pub mod stats;
11+
1012
#[path = "pub/mod.rs"]
1113
mod pubs;
1214
pub use pubs::{PubError, PubOptions, PubSocket};

‎msg-socket/src/pub/driver.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ where
7979
}
8080
Err(e) => {
8181
error!(err = ?e, "Error authenticating client");
82-
this.state.stats.decrement_active_clients();
82+
this.state.stats.specific.decrement_active_clients();
8383
}
8484
}
8585

@@ -93,15 +93,15 @@ where
9393
Ok(io) => {
9494
if let Err(e) = this.on_incoming(io) {
9595
error!(err = ?e, "Error accepting incoming connection");
96-
this.state.stats.decrement_active_clients();
96+
this.state.stats.specific.decrement_active_clients();
9797
}
9898
}
9999
Err(e) => {
100100
error!(err = ?e, "Error accepting incoming connection");
101101

102102
// Active clients have already been incremented in the initial call to
103103
// `poll_accept`, so we need to decrement them here.
104-
this.state.stats.decrement_active_clients();
104+
this.state.stats.specific.decrement_active_clients();
105105
}
106106
}
107107

@@ -112,15 +112,15 @@ where
112112
// incoming connection tasks.
113113
if let Poll::Ready(accept) = Pin::new(&mut this.transport).poll_accept(cx) {
114114
if let Some(max) = this.options.max_clients {
115-
if this.state.stats.active_clients() >= max {
115+
if this.state.stats.specific.active_clients() >= max {
116116
warn!("Max connections reached ({}), rejecting incoming connection", max);
117117
continue;
118118
}
119119
}
120120

121121
// Increment the active clients counter. If the authentication fails,
122122
// this counter will be decremented.
123-
this.state.stats.increment_active_clients();
123+
this.state.stats.specific.increment_active_clients();
124124

125125
this.conn_tasks.push(accept);
126126

‎msg-socket/src/pub/mod.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@ mod socket;
1111
pub use socket::*;
1212

1313
mod stats;
14-
use stats::SocketStats;
14+
use crate::stats::SocketStats;
15+
use stats::PubStats;
1516

1617
mod trie;
1718

@@ -167,7 +168,7 @@ impl PubMessage {
167168
/// The publisher socket state, shared between the backend task and the socket.
168169
#[derive(Debug, Default)]
169170
pub(crate) struct SocketState {
170-
pub(crate) stats: SocketStats,
171+
pub(crate) stats: SocketStats<PubStats>,
171172
}
172173

173174
#[cfg(test)]

‎msg-socket/src/pub/session.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ impl<Io: AsyncRead + AsyncWrite + Unpin> SubscriberSession<Io> {
9696

9797
impl<Io> Drop for SubscriberSession<Io> {
9898
fn drop(&mut self) {
99-
self.state.stats.decrement_active_clients();
99+
self.state.stats.specific.decrement_active_clients();
100100
}
101101
}
102102

@@ -157,7 +157,7 @@ impl<Io: AsyncRead + AsyncWrite + Unpin> Future for SubscriberSession<Io> {
157157

158158
match this.conn.start_send_unpin(msg) {
159159
Ok(_) => {
160-
this.state.stats.increment_tx(msg_len);
160+
this.state.stats.specific.increment_tx(msg_len);
161161

162162
this.should_flush = true;
163163
// We might be able to send more queued messages

‎msg-socket/src/pub/socket.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,14 @@ use tokio::{
99
};
1010
use tracing::{debug, trace, warn};
1111

12-
use super::{driver::PubDriver, stats::SocketStats, PubError, PubMessage, PubOptions, SocketState};
12+
use super::{driver::PubDriver, stats::PubStats, PubError, PubMessage, PubOptions, SocketState};
1313
use crate::Authenticator;
1414

1515
use msg_transport::{Address, Transport};
1616
use msg_wire::compression::Compressor;
1717

1818
/// A publisher socket. This is thread-safe and can be cloned.
19-
#[derive(Clone, Default)]
19+
#[derive(Clone)]
2020
pub struct PubSocket<T: Transport<A>, A: Address> {
2121
/// The reply socket options, shared with the driver.
2222
options: Arc<PubOptions>,
@@ -190,8 +190,8 @@ where
190190
Ok(())
191191
}
192192

193-
pub fn stats(&self) -> &SocketStats {
194-
&self.state.stats
193+
pub fn stats(&self) -> &PubStats {
194+
&self.state.stats.specific
195195
}
196196

197197
/// Returns the local address this socket is bound to. `None` if the socket is not bound.

‎msg-socket/src/pub/stats.rs‎

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,36 @@
11
use std::sync::atomic::{AtomicUsize, Ordering};
22

3-
/// Statistics for a reply socket. These are shared between the driver task
4-
/// and the socket.
53
#[derive(Debug, Default)]
6-
pub struct SocketStats {
4+
pub struct PubStats {
75
/// Total bytes sent
86
bytes_tx: AtomicUsize,
97
/// Total number of active request clients
108
active_clients: AtomicUsize,
119
}
1210

13-
impl SocketStats {
11+
impl PubStats {
1412
#[inline]
15-
pub(crate) fn increment_tx(&self, bytes: usize) {
16-
self.bytes_tx.fetch_add(bytes, Ordering::Relaxed);
13+
pub fn bytes_tx(&self) -> usize {
14+
self.bytes_tx.load(Ordering::Relaxed)
1715
}
1816

1917
#[inline]
20-
pub(crate) fn increment_active_clients(&self) {
21-
self.active_clients.fetch_add(1, Ordering::Relaxed);
18+
pub fn active_clients(&self) -> usize {
19+
self.active_clients.load(Ordering::Relaxed)
2220
}
2321

2422
#[inline]
25-
pub(crate) fn decrement_active_clients(&self) {
26-
self.active_clients.fetch_sub(1, Ordering::Relaxed);
23+
pub(crate) fn increment_tx(&self, bytes: usize) {
24+
self.bytes_tx.fetch_add(bytes, Ordering::Relaxed);
2725
}
2826

2927
#[inline]
30-
pub fn bytes_tx(&self) -> usize {
31-
self.bytes_tx.load(Ordering::Relaxed)
28+
pub(crate) fn increment_active_clients(&self) {
29+
self.active_clients.fetch_add(1, Ordering::Relaxed);
3230
}
3331

3432
#[inline]
35-
pub fn active_clients(&self) -> usize {
36-
self.active_clients.load(Ordering::Relaxed)
33+
pub(crate) fn decrement_active_clients(&self) {
34+
self.active_clients.fetch_sub(1, Ordering::Relaxed);
3735
}
3836
}

‎msg-socket/src/rep/driver.rs‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -89,15 +89,15 @@ where
8989
}
9090
}
9191

92-
this.state.stats.increment_rx(size);
92+
this.state.stats.specific.increment_rx(size);
9393
let _ = this.to_socket.try_send(request);
9494
}
9595
Some(Err(e)) => {
9696
error!(err = ?e, "Error receiving message from peer {:?}", peer);
9797
}
9898
None => {
9999
warn!("Peer {:?} disconnected", peer);
100-
this.state.stats.decrement_active_clients();
100+
this.state.stats.specific.decrement_active_clients();
101101
}
102102
}
103103

@@ -125,7 +125,7 @@ where
125125
}
126126
Err(e) => {
127127
error!(err = ?e, "Error authenticating client");
128-
this.state.stats.decrement_active_clients();
128+
this.state.stats.specific.decrement_active_clients();
129129
}
130130
}
131131

@@ -137,15 +137,15 @@ where
137137
Ok(io) => {
138138
if let Err(e) = this.on_incoming(io) {
139139
error!(err = ?e, "Error accepting incoming connection");
140-
this.state.stats.decrement_active_clients();
140+
this.state.stats.specific.decrement_active_clients();
141141
}
142142
}
143143
Err(e) => {
144144
error!(err = ?e, "Error accepting incoming connection");
145145

146146
// Active clients have already been incremented in the initial call to
147147
// `poll_accept`, so we need to decrement them here.
148-
this.state.stats.decrement_active_clients();
148+
this.state.stats.specific.decrement_active_clients();
149149
}
150150
}
151151

@@ -156,7 +156,7 @@ where
156156
// incoming connection tasks.
157157
if let Poll::Ready(accept) = Pin::new(&mut this.transport).poll_accept(cx) {
158158
if let Some(max) = this.options.max_clients {
159-
if this.state.stats.active_clients() >= max {
159+
if this.state.stats.specific.active_clients() >= max {
160160
warn!(
161161
"Max connections reached ({}), rejecting new incoming connection",
162162
max
@@ -168,7 +168,7 @@ where
168168

169169
// Increment the active clients counter. If the authentication fails, this counter
170170
// will be decremented.
171-
this.state.stats.increment_active_clients();
171+
this.state.stats.specific.increment_active_clients();
172172

173173
this.conn_tasks.push(accept);
174174

@@ -271,14 +271,14 @@ impl<T: AsyncRead + AsyncWrite + Unpin, A: Address + Unpin> Stream for PeerState
271271
let msg_len = msg.size();
272272
match this.conn.start_send_unpin(msg) {
273273
Ok(_) => {
274-
this.state.stats.increment_tx(msg_len);
274+
this.state.stats.specific.increment_tx(msg_len);
275275
this.should_flush = true;
276276

277277
// We might be able to send more queued messages
278278
continue;
279279
}
280280
Err(e) => {
281-
this.state.stats.increment_failed_requests();
281+
this.state.stats.specific.increment_failed_requests();
282282
error!(err = ?e, "Failed to send message to socket");
283283
// End this stream as we can't send any more messages
284284
return Poll::Ready(None);

‎msg-socket/src/rep/mod.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,9 @@ use tokio::sync::oneshot;
66
mod driver;
77
mod socket;
88
mod stats;
9+
use crate::stats::SocketStats;
910
pub use socket::*;
10-
use stats::SocketStats;
11+
use stats::RepStats;
1112

1213
/// Errors that can occur when using a reply socket.
1314
#[derive(Debug, Error)]
@@ -55,7 +56,7 @@ impl RepOptions {
5556
/// The request socket state, shared between the backend task and the socket.
5657
#[derive(Debug, Default)]
5758
pub(crate) struct SocketState {
58-
pub(crate) stats: SocketStats,
59+
pub(crate) stats: SocketStats<RepStats>,
5960
}
6061

6162
/// A request received by the socket.

‎msg-socket/src/rep/socket.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,16 @@ use tokio_stream::StreamMap;
1616
use tracing::{debug, warn};
1717

1818
use crate::{
19-
rep::{driver::RepDriver, RepError, SocketState, SocketStats},
19+
rep::{driver::RepDriver, RepError, SocketState},
2020
Authenticator, RepOptions, Request, DEFAULT_BUFFER_SIZE,
2121
};
2222

2323
use msg_transport::{Address, Transport};
2424
use msg_wire::compression::Compressor;
2525

26+
use super::stats::RepStats;
27+
2628
/// A reply socket. This socket implements [`Stream`] and yields incoming [`Request`]s.
27-
#[derive(Default)]
2829
pub struct RepSocket<T: Transport<A>, A: Address> {
2930
/// The reply socket options, shared with the driver.
3031
options: Arc<RepOptions>,
@@ -143,8 +144,8 @@ where
143144
}
144145

145146
/// Returns the statistics for this socket.
146-
pub fn stats(&self) -> &SocketStats {
147-
&self.state.stats
147+
pub fn stats(&self) -> &RepStats {
148+
&self.state.stats.specific
148149
}
149150

150151
/// Returns the local address this socket is bound to. `None` if the socket is not bound.

‎msg-socket/src/rep/stats.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
33
/// Statistics for a reply socket.
44
/// These are shared between the driver task and the socket.
55
#[derive(Debug, Default)]
6-
pub struct SocketStats {
6+
pub struct RepStats {
77
/// Total bytes sent
88
bytes_tx: AtomicUsize,
99
/// Total bytes received
@@ -14,7 +14,7 @@ pub struct SocketStats {
1414
failed_requests: AtomicUsize,
1515
}
1616

17-
impl SocketStats {
17+
impl RepStats {
1818
#[inline]
1919
pub(crate) fn increment_tx(&self, bytes: usize) {
2020
self.bytes_tx.fetch_add(bytes, Ordering::Relaxed);

0 commit comments

Comments
 (0)