Skip to content

Commit 4ae1bc8

Browse files
drdrshdrdrsh
authored
Add SHOW CLIENTS / SHOW SERVERS + Stats refactor and tests (#159)
* wip * Main Thread Panic when swarmed with clients * fix * fix * 1024 * fix * remove test * Add SHOW CLIENTS * revert * fmt * Refactor + tests * fmt * add test * Add SHOW SERVERS + Make PR unreviewable * prometheus * add state to clients and servers * fmt * Add application_name to server stats * Add tests for waiting clients * Docs * remove comment * comments * typo * cleanup * CI
1 parent 0751674 commit 4ae1bc8

File tree

9 files changed

+1203
-372
lines changed

9 files changed

+1203
-372
lines changed

‎src/admin.rs

Lines changed: 155 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,15 @@
22
use bytes::{Buf, BufMut, BytesMut};
33
use log::{info, trace};
44
use std::collections::HashMap;
5+
use tokio::time::Instant;
56

67
use crate::config::{get_config, reload_config, VERSION};
78
use crate::errors::Error;
89
use crate::messages::*;
910
use crate::pool::get_all_pools;
10-
use crate::stats::get_stats;
11+
use crate::stats::{
12+
get_address_stats, get_client_stats, get_pool_stats, get_server_stats, ClientState, ServerState,
13+
};
1114
use crate::ClientServerMap;
1215

1316
pub fn generate_server_info_for_admin() -> BytesMut {
@@ -72,6 +75,14 @@ where
7275
trace!("SHOW POOLS");
7376
show_pools(stream).await
7477
}
78+
"CLIENTS" => {
79+
trace!("SHOW CLIENTS");
80+
show_clients(stream).await
81+
}
82+
"SERVERS" => {
83+
trace!("SHOW SERVERS");
84+
show_servers(stream).await
85+
}
7586
"STATS" => {
7687
trace!("SHOW STATS");
7788
show_stats(stream).await
@@ -91,7 +102,8 @@ async fn show_lists<T>(stream: &mut T) -> Result<(), Error>
91102
where
92103
T: tokio::io::AsyncWrite + std::marker::Unpin,
93104
{
94-
let stats = get_stats();
105+
let client_stats = get_client_stats();
106+
let server_stats = get_server_stats();
95107

96108
let columns = vec![("list", DataType::Text), ("items", DataType::Int4)];
97109

@@ -111,18 +123,18 @@ where
111123
res.put(data_row(&vec!["pools".to_string(), databases.to_string()]));
112124
res.put(data_row(&vec![
113125
"free_clients".to_string(),
114-
stats
126+
client_stats
115127
.keys()
116-
.map(|address_id| stats[&address_id]["cl_idle"])
117-
.sum::<i64>()
128+
.filter(|client_id| client_stats.get(client_id).unwrap().state == ClientState::Idle)
129+
.count()
118130
.to_string(),
119131
]));
120132
res.put(data_row(&vec![
121133
"used_clients".to_string(),
122-
stats
134+
client_stats
123135
.keys()
124-
.map(|address_id| stats[&address_id]["cl_active"])
125-
.sum::<i64>()
136+
.filter(|client_id| client_stats.get(client_id).unwrap().state == ClientState::Active)
137+
.count()
126138
.to_string(),
127139
]));
128140
res.put(data_row(&vec![
@@ -131,18 +143,18 @@ where
131143
]));
132144
res.put(data_row(&vec![
133145
"free_servers".to_string(),
134-
stats
146+
server_stats
135147
.keys()
136-
.map(|address_id| stats[&address_id]["sv_idle"])
137-
.sum::<i64>()
148+
.filter(|server_id| server_stats.get(server_id).unwrap().state == ServerState::Idle)
149+
.count()
138150
.to_string(),
139151
]));
140152
res.put(data_row(&vec![
141153
"used_servers".to_string(),
142-
stats
154+
server_stats
143155
.keys()
144-
.map(|address_id| stats[&address_id]["sv_active"])
145-
.sum::<i64>()
156+
.filter(|server_id| server_stats.get(server_id).unwrap().state == ServerState::Active)
157+
.count()
146158
.to_string(),
147159
]));
148160
res.put(data_row(&vec!["dns_names".to_string(), "0".to_string()]));
@@ -182,11 +194,12 @@ async fn show_pools<T>(stream: &mut T) -> Result<(), Error>
182194
where
183195
T: tokio::io::AsyncWrite + std::marker::Unpin,
184196
{
185-
let stats = get_stats();
197+
let all_pool_stats = get_pool_stats();
186198

187199
let columns = vec![
188200
("database", DataType::Text),
189201
("user", DataType::Text),
202+
("pool_mode", DataType::Text),
190203
("cl_idle", DataType::Numeric),
191204
("cl_active", DataType::Numeric),
192205
("cl_waiting", DataType::Numeric),
@@ -198,32 +211,27 @@ where
198211
("sv_login", DataType::Numeric),
199212
("maxwait", DataType::Numeric),
200213
("maxwait_us", DataType::Numeric),
201-
("pool_mode", DataType::Text),
202214
];
203215

204216
let mut res = BytesMut::new();
205217
res.put(row_description(&columns));
206-
for (_, pool) in get_all_pools() {
207-
let pool_config = &pool.settings;
208-
for shard in 0..pool.shards() {
209-
for server in 0..pool.servers(shard) {
210-
let address = pool.address(shard, server);
211-
let stats = match stats.get(&address.id) {
212-
Some(stats) => stats.clone(),
213-
None => HashMap::new(),
214-
};
218+
for ((pool_name, username), pool) in get_all_pools() {
219+
let def = HashMap::default();
220+
let pool_stats = all_pool_stats
221+
.get(&(pool_name.clone(), username.clone()))
222+
.unwrap_or(&def);
215223

216-
let mut row = vec![address.name(), pool_config.user.username.clone()];
217-
218-
for column in &columns[2..columns.len() - 1] {
219-
let value = stats.get(column.0).unwrap_or(&0).to_string();
220-
row.push(value);
221-
}
222-
223-
row.push(pool_config.pool_mode.to_string());
224-
res.put(data_row(&row));
225-
}
224+
let pool_config = &pool.settings;
225+
let mut row = vec![
226+
pool_name.clone(),
227+
username.clone(),
228+
pool_config.pool_mode.to_string(),
229+
];
230+
for column in &columns[3..columns.len()] {
231+
let value = pool_stats.get(column.0).unwrap_or(&0).to_string();
232+
row.push(value);
226233
}
234+
res.put(data_row(&row));
227235
}
228236

229237
res.put(command_complete("SHOW"));
@@ -387,6 +395,7 @@ where
387395
T: tokio::io::AsyncWrite + std::marker::Unpin,
388396
{
389397
let columns = vec![
398+
("instance", DataType::Text),
390399
("database", DataType::Text),
391400
("user", DataType::Text),
392401
("total_xact_count", DataType::Numeric),
@@ -396,32 +405,32 @@ where
396405
("total_xact_time", DataType::Numeric),
397406
("total_query_time", DataType::Numeric),
398407
("total_wait_time", DataType::Numeric),
408+
("total_errors", DataType::Numeric),
399409
("avg_xact_count", DataType::Numeric),
400410
("avg_query_count", DataType::Numeric),
401411
("avg_recv", DataType::Numeric),
402412
("avg_sent", DataType::Numeric),
413+
("avg_errors", DataType::Numeric),
403414
("avg_xact_time", DataType::Numeric),
404415
("avg_query_time", DataType::Numeric),
405416
("avg_wait_time", DataType::Numeric),
406417
];
407418

408-
let stats = get_stats();
419+
let all_stats = get_address_stats();
409420
let mut res = BytesMut::new();
410421
res.put(row_description(&columns));
411422

412-
for ((_db_name, username), pool) in get_all_pools() {
423+
for ((db, username), pool) in get_all_pools() {
413424
for shard in 0..pool.shards() {
414425
for server in 0..pool.servers(shard) {
415426
let address = pool.address(shard, server);
416-
let stats = match stats.get(&address.id) {
427+
let stats = match all_stats.get(&address.id) {
417428
Some(stats) => stats.clone(),
418429
None => HashMap::new(),
419430
};
420431

421-
let mut row = vec![address.name()];
422-
row.push(username.clone());
423-
424-
for column in &columns[2..] {
432+
let mut row = vec![address.name(), db.clone(), username.clone()];
433+
for column in &columns[3..] {
425434
row.push(stats.get(column.0).unwrap_or(&0).to_string());
426435
}
427436

@@ -439,3 +448,107 @@ where
439448

440449
write_all_half(stream, res).await
441450
}
451+
452+
/// Show currently connected clients
453+
async fn show_clients<T>(stream: &mut T) -> Result<(), Error>
454+
where
455+
T: tokio::io::AsyncWrite + std::marker::Unpin,
456+
{
457+
let columns = vec![
458+
("client_id", DataType::Text),
459+
("database", DataType::Text),
460+
("user", DataType::Text),
461+
("application_name", DataType::Text),
462+
("state", DataType::Text),
463+
("transaction_count", DataType::Numeric),
464+
("query_count", DataType::Numeric),
465+
("error_count", DataType::Numeric),
466+
("age_seconds", DataType::Numeric),
467+
];
468+
469+
let new_map = get_client_stats();
470+
let mut res = BytesMut::new();
471+
res.put(row_description(&columns));
472+
473+
for (_, client) in new_map {
474+
let row = vec![
475+
format!("{:#010X}", client.client_id),
476+
client.pool_name,
477+
client.username,
478+
client.application_name.clone(),
479+
client.state.to_string(),
480+
client.transaction_count.to_string(),
481+
client.query_count.to_string(),
482+
client.error_count.to_string(),
483+
Instant::now()
484+
.duration_since(client.connect_time)
485+
.as_secs()
486+
.to_string(),
487+
];
488+
489+
res.put(data_row(&row));
490+
}
491+
492+
res.put(command_complete("SHOW"));
493+
494+
// ReadyForQuery
495+
res.put_u8(b'Z');
496+
res.put_i32(5);
497+
res.put_u8(b'I');
498+
499+
write_all_half(stream, res).await
500+
}
501+
502+
/// Show currently connected servers
503+
async fn show_servers<T>(stream: &mut T) -> Result<(), Error>
504+
where
505+
T: tokio::io::AsyncWrite + std::marker::Unpin,
506+
{
507+
let columns = vec![
508+
("server_id", DataType::Text),
509+
("database_name", DataType::Text),
510+
("user", DataType::Text),
511+
("address_id", DataType::Text),
512+
("application_name", DataType::Text),
513+
("state", DataType::Text),
514+
("transaction_count", DataType::Numeric),
515+
("query_count", DataType::Numeric),
516+
("bytes_sent", DataType::Numeric),
517+
("bytes_received", DataType::Numeric),
518+
("age_seconds", DataType::Numeric),
519+
];
520+
521+
let new_map = get_server_stats();
522+
let mut res = BytesMut::new();
523+
res.put(row_description(&columns));
524+
525+
for (_, server) in new_map {
526+
let row = vec![
527+
format!("{:#010X}", server.server_id),
528+
server.pool_name,
529+
server.username,
530+
server.address_name,
531+
server.application_name,
532+
server.state.to_string(),
533+
server.transaction_count.to_string(),
534+
server.query_count.to_string(),
535+
server.bytes_sent.to_string(),
536+
server.bytes_received.to_string(),
537+
Instant::now()
538+
.duration_since(server.connect_time)
539+
.as_secs()
540+
.to_string(),
541+
];
542+
543+
res.put(data_row(&row));
544+
}
545+
546+
res.put(command_complete("SHOW"));
547+
548+
// ReadyForQuery
549+
res.put_u8(b'Z');
550+
res.put_i32(5);
551+
res.put_u8(b'I');
552+
553+
write_all_half(stream, res).await
554+
}

0 commit comments

Comments
 (0)