rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-24 17:00:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: reduce unit test suite wall time (#7657)

Browse Source 
* feat: reduce wait timt of  from 57s to 0.6s

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* , , ,

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* test_query_concurrently

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2026-02-03 21:32:11 +08:00
committed by GitHub
parent 969a64d483
commit 0f2f20d4b7
7 changed files with 391 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
src/frontend/src/frontend.rs
View File

@@ -207,6 +207,10 @@ mod tests {
let mut requests = request.into_inner();
let suspend = self.suspend.clone();
async move {
// Make the heartbeat interval short in unit tests to reduce the waiting time.
// Only the handshake response needs to populate it (as metasrv does).
let heartbeat_interval_ms = Duration::from_millis(200).as_millis() as u64;
let mut is_handshake = true;
while let Some(request) = requests.next().await {
if let Err(e) = request {
let _ = tx.send(Err(e)).await;
@@ -220,9 +224,16 @@ mod tests {
)),
..Default::default()
});
let heartbeat_config =
is_handshake.then_some(api::v1::meta::HeartbeatConfig {
heartbeat_interval_ms,
retry_interval_ms: heartbeat_interval_ms,
});
is_handshake = false;
let response = HeartbeatResponse {
header: Some(ResponseHeader::success()),
mailbox_message,
heartbeat_config,
..Default::default()
};
@@ -376,6 +387,21 @@ mod tests {
}
}
async fn wait_for_suspend_state(frontend: &Frontend, expected: bool) {
let check = || frontend.instance.is_suspended() == expected;
if check() {
return;
}
tokio::time::timeout(Duration::from_secs(5), async move {
while !check() {
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_suspend_frontend() -> Result<()> {
common_telemetry::init_default_ut_logging();
@@ -412,12 +438,6 @@ mod tests {
let meta_client = create_meta_client(&meta_client_options, server.clone()).await;
let frontend = create_frontend(&options, meta_client).await?;
use common_meta::distributed_time_constants::{
BASE_HEARTBEAT_INTERVAL, frontend_heartbeat_interval,
};
let frontend_heartbeat_interval =
frontend_heartbeat_interval(BASE_HEARTBEAT_INTERVAL) + Duration::from_secs(1);
tokio::time::sleep(frontend_heartbeat_interval).await;
// initial state: not suspend:
assert!(!frontend.instance.is_suspended());
verify_suspend_state_by_http(&frontend, Ok(r#"[{"records":{"schema":{"column_schemas":[{"name":"Int64(1)","data_type":"Int64"}]},"rows":[[1]],"total_rows":1}}]"#)).await;
@@ -434,7 +454,7 @@ mod tests {
// make heartbeat server returned "suspend" instruction,
server.suspend.store(true, Ordering::Relaxed);
tokio::time::sleep(frontend_heartbeat_interval).await;
wait_for_suspend_state(&frontend, true).await;
// ... then the frontend is suspended:
assert!(frontend.instance.is_suspended());
verify_suspend_state_by_http(
@@ -450,7 +470,7 @@ mod tests {
// make heartbeat server NOT returned "suspend" instruction,
server.suspend.store(false, Ordering::Relaxed);
tokio::time::sleep(frontend_heartbeat_interval).await;
wait_for_suspend_state(&frontend, false).await;
// ... then frontend's suspend state is cleared:
assert!(!frontend.instance.is_suspended());
verify_suspend_state_by_http(&frontend, Ok(r#"[{"records":{"schema":{"column_schemas":[{"name":"Int64(1)","data_type":"Int64"}]},"rows":[[1]],"total_rows":1}}]"#)).await;

19
src/servers/tests/mysql/mysql_server_test.rs
View File

@@ -26,7 +26,6 @@ use datatypes::schema::{ColumnSchema, Schema};
use datatypes::value::Value;
use mysql_async::prelude::*;
use mysql_async::{Conn, Row, SslOpts};
use rand::Rng;
use servers::error::Result;
use servers::install_ring_crypto_provider;
use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
@@ -436,19 +435,23 @@ async fn test_query_concurrently() -> Result<()> {
let server_port = server_addr.port();
let threads = 4;
let expect_executed_queries_per_worker = 1000;
let expect_executed_queries_per_worker = 200;
let queries = Arc::new(
(0..100u32)
.map(|expected| format!("SELECT uint32s FROM numbers WHERE uint32s = {expected}"))
.collect::<Vec<_>>(),
);
let mut join_handles = vec![];
for _ in 0..threads {
for worker_id in 0..threads {
let queries = queries.clone();
join_handles.push(tokio::spawn(async move {
let mut connection = create_connection_default_db_name(server_port, false)
.await
.unwrap();
for _ in 0..expect_executed_queries_per_worker {
let expected: u32 = rand::rng().random_range(0..100);
for i in 0..expect_executed_queries_per_worker {
let expected: u32 = ((i + worker_id) % 100) as u32;
let result: u32 = connection
.query_first(format!(
"SELECT uint32s FROM numbers WHERE uint32s = {expected}"
))
.query_first(&queries[expected as usize])
.await
.unwrap()
.unwrap();

154
src/sql/src/partition.rs
View File

@@ -120,8 +120,6 @@ fn partition_rules_for_uuid(partition_num: u32, ident: &str) -> Result<Vec<Expr>
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use sqlparser::ast::{Expr, ValueWithSpan};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;
@@ -180,20 +178,16 @@ mod tests {
}
fn check_distribution(test_partition: u32, test_uuid_num: usize) -> bool {
// Generate test_uuid_num random uuids.
let uuids = (0..test_uuid_num)
.map(|_| Uuid::new_v4().to_string().replace("-", "").to_lowercase())
.collect::<Vec<String>>();
// Generate the partition rules.
let rules = partition_rules_for_uuid(test_partition, "test_trace_id").unwrap();
let upper_bounds = upper_bounds_from_rules(&rules);
// Collect the number of partitions for each uuid.
let mut stats = HashMap::new();
for uuid in uuids {
let partition = allocate_partition_for_uuid(uuid.clone(), &rules);
// Count the number of uuids in each partition.
*stats.entry(partition).or_insert(0) += 1;
let mut counts = vec![0usize; test_partition as usize];
for _ in 0..test_uuid_num {
let uuid = Uuid::new_v4().simple().to_string();
let partition = allocate_partition_for_uuid(uuid.as_str(), &upper_bounds);
counts[partition] += 1;
}
// Check if the partition distribution is uniform.
@@ -203,7 +197,7 @@ mod tests {
let tolerance = 100.0 / test_partition as f64 * 0.30;
// For each partition, its ratio should be as close as possible to the expected ratio.
for (_, count) in stats {
for count in counts {
let ratio = (count as f64 / test_uuid_num as f64) * 100.0;
if (ratio - expected_ratio).abs() >= tolerance {
return false;
@@ -213,58 +207,94 @@ mod tests {
true
}
fn allocate_partition_for_uuid(uuid: String, rules: &[Expr]) -> usize {
for (i, rule) in rules.iter().enumerate() {
if let Expr::BinaryOp { left, op: _, right } = rule {
if i == 0 {
// Hit the leftmost rule.
if let Expr::Value(ValueWithSpan {
value: Value::SingleQuotedString(leftmost),
..
}) = *right.clone()
&& uuid < leftmost
{
return i;
}
} else if i == rules.len() - 1 {
// Hit the rightmost rule.
if let Expr::Value(ValueWithSpan {
value: Value::SingleQuotedString(rightmost),
..
}) = *right.clone()
&& uuid >= rightmost
{
return i;
}
} else {
// Hit the middle rules.
if let Expr::BinaryOp {
left: _,
op: _,
right: inner_right,
} = *left.clone()
&& let Expr::Value(ValueWithSpan {
value: Value::SingleQuotedString(lower),
..
}) = *inner_right.clone()
&& let Expr::BinaryOp {
left: _,
op: _,
right: inner_right,
} = *right.clone()
&& let Expr::Value(ValueWithSpan {
value: Value::SingleQuotedString(upper),
..
}) = *inner_right.clone()
&& uuid >= lower
&& uuid < upper
{
return i;
}
fn upper_bounds_from_rules<'a>(rules: &'a [Expr]) -> Vec<&'a str> {
let mut upper_bounds = Vec::with_capacity(rules.len().saturating_sub(1));
let mut prev_upper: Option<&'a str> = None;
for (idx, rule) in rules.iter().enumerate() {
let (lower, upper) = extract_rule_bounds(rule);
match idx {
0 => {
assert!(lower.is_none());
assert!(upper.is_some());
}
idx if idx == rules.len() - 1 => {
assert_eq!(lower, prev_upper);
assert!(upper.is_none());
}
_ => {
assert_eq!(lower, prev_upper);
assert!(upper.is_some());
}
}
if idx < rules.len() - 1 {
upper_bounds.push(upper.unwrap());
}
prev_upper = upper;
}
upper_bounds
}
fn extract_rule_bounds(rule: &Expr) -> (Option<&str>, Option<&str>) {
fn extract_single_quoted_string(expr: &Expr) -> Option<&str> {
match expr {
Expr::Value(ValueWithSpan {
value: Value::SingleQuotedString(value),
..
}) => Some(value.as_str()),
_ => None,
}
}
panic!("No partition found for uuid: {}, rules: {:?}", uuid, rules);
match rule {
Expr::BinaryOp {
op: BinaryOperator::Lt,
right,
..
} => (None, extract_single_quoted_string(right)),
Expr::BinaryOp {
op: BinaryOperator::GtEq,
right,
..
} => (extract_single_quoted_string(right), None),
Expr::BinaryOp {
op: BinaryOperator::And,
left,
right,
} => {
let lower = match left.as_ref() {
Expr::BinaryOp {
op: BinaryOperator::GtEq,
right,
..
} => extract_single_quoted_string(right),
_ => None,
};
let upper = match right.as_ref() {
Expr::BinaryOp {
op: BinaryOperator::Lt,
right,
..
} => extract_single_quoted_string(right),
_ => None,
};
(lower, upper)
}
_ => (None, None),
}
}
fn allocate_partition_for_uuid(uuid: &str, upper_bounds: &[&str]) -> usize {
upper_bounds.partition_point(|upper| uuid >= *upper)
}
#[test]
fn test_extract_rule_bounds() {
let rules = partition_rules_for_uuid(16, "trace_id").unwrap();
let upper_bounds = upper_bounds_from_rules(&rules);
assert_eq!(upper_bounds.len(), 15);
}
}