cli/
bench.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{BTreeMap, HashMap};
16use std::future::Future;
17use std::sync::Arc;
18use std::time::Duration;
19
20use async_trait::async_trait;
21use clap::Parser;
22use common_error::ext::BoxedError;
23use common_meta::key::{TableMetadataManager, TableMetadataManagerRef};
24use common_meta::kv_backend::etcd::EtcdStore;
25use common_meta::kv_backend::memory::MemoryKvBackend;
26#[cfg(feature = "mysql_kvbackend")]
27use common_meta::kv_backend::rds::MySqlStore;
28#[cfg(feature = "pg_kvbackend")]
29use common_meta::kv_backend::rds::PgStore;
30use common_meta::peer::Peer;
31use common_meta::rpc::router::{Region, RegionRoute};
32use common_telemetry::info;
33use common_wal::options::WalOptions;
34use datatypes::data_type::ConcreteDataType;
35use datatypes::schema::{ColumnSchema, RawSchema};
36use rand::Rng;
37use store_api::storage::RegionNumber;
38use table::metadata::{RawTableInfo, RawTableMeta, TableId, TableIdent, TableType};
39use table::table_name::TableName;
40
41use self::metadata::TableMetadataBencher;
42use crate::Tool;
43
44mod metadata;
45
46async fn bench_self_recorded<F, Fut>(desc: &str, f: F, count: u32)
47where
48    F: Fn(u32) -> Fut,
49    Fut: Future<Output = Duration>,
50{
51    let mut total = Duration::default();
52
53    for i in 1..=count {
54        total += f(i).await;
55    }
56
57    let cost = total.as_millis() as f64 / count as f64;
58    info!("{desc}, average operation cost: {cost:.2} ms");
59}
60
61/// Command to benchmark table metadata operations.
62#[derive(Debug, Default, Parser)]
63pub struct BenchTableMetadataCommand {
64    #[clap(long)]
65    etcd_addr: Option<String>,
66    #[cfg(feature = "pg_kvbackend")]
67    #[clap(long)]
68    postgres_addr: Option<String>,
69    #[cfg(feature = "pg_kvbackend")]
70    #[clap(long)]
71    postgres_schema: Option<String>,
72    #[cfg(feature = "mysql_kvbackend")]
73    #[clap(long)]
74    mysql_addr: Option<String>,
75    #[clap(long)]
76    count: u32,
77}
78
79impl BenchTableMetadataCommand {
80    pub async fn build(&self) -> std::result::Result<Box<dyn Tool>, BoxedError> {
81        let kv_backend = if let Some(etcd_addr) = &self.etcd_addr {
82            info!("Using etcd as kv backend");
83            EtcdStore::with_endpoints([etcd_addr], 128).await.unwrap()
84        } else {
85            Arc::new(MemoryKvBackend::new())
86        };
87
88        #[cfg(feature = "pg_kvbackend")]
89        let kv_backend = if let Some(postgres_addr) = &self.postgres_addr {
90            info!("Using postgres as kv backend");
91            PgStore::with_url(postgres_addr, "greptime_metakv", 128)
92                .await
93                .unwrap()
94        } else {
95            kv_backend
96        };
97
98        #[cfg(feature = "mysql_kvbackend")]
99        let kv_backend = if let Some(mysql_addr) = &self.mysql_addr {
100            info!("Using mysql as kv backend");
101            MySqlStore::with_url(mysql_addr, "greptime_metakv", 128)
102                .await
103                .unwrap()
104        } else {
105            kv_backend
106        };
107
108        let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend));
109
110        let tool = BenchTableMetadata {
111            table_metadata_manager,
112            count: self.count,
113        };
114        Ok(Box::new(tool))
115    }
116}
117
118struct BenchTableMetadata {
119    table_metadata_manager: TableMetadataManagerRef,
120    count: u32,
121}
122
123#[async_trait]
124impl Tool for BenchTableMetadata {
125    async fn do_work(&self) -> std::result::Result<(), BoxedError> {
126        let bencher = TableMetadataBencher::new(self.table_metadata_manager.clone(), self.count);
127        bencher.bench_create().await;
128        bencher.bench_get().await;
129        bencher.bench_rename().await;
130        bencher.bench_delete().await;
131        Ok(())
132    }
133}
134
135fn create_table_info(table_id: TableId, table_name: TableName) -> RawTableInfo {
136    let columns = 100;
137    let mut column_schemas = Vec::with_capacity(columns);
138    column_schemas.push(
139        ColumnSchema::new(
140            "ts",
141            ConcreteDataType::timestamp_millisecond_datatype(),
142            true,
143        )
144        .with_time_index(true),
145    );
146
147    for i in 1..columns {
148        let column_name = format!("my_column_{i}");
149        column_schemas.push(ColumnSchema::new(
150            column_name,
151            ConcreteDataType::string_datatype(),
152            true,
153        ));
154    }
155
156    let meta = RawTableMeta {
157        schema: RawSchema::new(column_schemas),
158        engine: "mito".to_string(),
159        created_on: chrono::DateTime::default(),
160        primary_key_indices: vec![],
161        next_column_id: columns as u32 + 1,
162        value_indices: vec![],
163        options: Default::default(),
164        region_numbers: (1..=100).collect(),
165        partition_key_indices: vec![],
166        column_ids: vec![],
167    };
168
169    RawTableInfo {
170        ident: TableIdent {
171            table_id,
172            version: 1,
173        },
174        name: table_name.table_name,
175        desc: Some("blah".to_string()),
176        catalog_name: table_name.catalog_name,
177        schema_name: table_name.schema_name,
178        meta,
179        table_type: TableType::Base,
180    }
181}
182
183fn create_region_routes(regions: Vec<RegionNumber>) -> Vec<RegionRoute> {
184    let mut region_routes = Vec::with_capacity(100);
185    let mut rng = rand::rng();
186
187    for region_id in regions.into_iter().map(u64::from) {
188        region_routes.push(RegionRoute {
189            region: Region {
190                id: region_id.into(),
191                name: String::new(),
192                partition: None,
193                attrs: BTreeMap::new(),
194                partition_expr: Default::default(),
195            },
196            leader_peer: Some(Peer {
197                id: rng.random_range(0..10),
198                addr: String::new(),
199            }),
200            follower_peers: vec![],
201            leader_state: None,
202            leader_down_since: None,
203        });
204    }
205
206    region_routes
207}
208
209fn create_region_wal_options(regions: Vec<RegionNumber>) -> HashMap<RegionNumber, WalOptions> {
210    // TODO(niebayes): construct region wal options for benchmark.
211    let _ = regions;
212    HashMap::default()
213}
cli/bench.rs

cli/
bench.rs
