Merge pull request #268 from ekexium/multi-reigon-test

Support multi-reigon test

Author: nrc

.travis.yml

@@ -43,18 +43,18 @@ script:
   - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "stable" ]]; then cargo clippy -- -D clippy::all; fi
   - cargo build --all
   - cargo build --examples
-  - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then cargo doc --workspace --exclude tikv-client-proto --document-private-items --no-deps; fi
+  - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then make doc; fi
   - if [[ $TRAVIS_RUST_VERSION == "nightly" ]]; then cargo test --all -- --nocapture; fi
   # For now we only run full integration tests on Linux. Here's why:
   # * Docker on OS X is not supported by Travis.
   # * Docker on Windows seems to not have the correct binary at `"/c/Program Files/Docker/Docker/DockerCli.exe" to switch it to Linux containers.
-  - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then docker run -d --net=host --name pd --rm pingcap/pd --name "pd" --data-dir "pd" --client-urls "http://127.0.0.1:2379" --advertise-client-urls "http://127.0.0.1:2379"; fi
-  - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then docker run -d --net=host --name kv --rm --ulimit nofile=90000:90000 pingcap/tikv --pd-endpoints "127.0.0.1:2379" --addr "127.0.0.1:2378" --data-dir "kv"; fi
+  - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then make docker-pd; fi
+  - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then make docker-kv; fi
   - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then docker ps; fi
   - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then docker logs pd; fi
   - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then docker logs kv; fi
   - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then sleep 60; fi
-  - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then PD_ADDRS="127.0.0.1:2379" cargo test --all --features integration-tests -- --nocapture; fi
+  - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then MULTI_REGION=1 make integration-test; fi
   - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then cargo run --example raw -- --pd="127.0.0.1:2379"; fi
   - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then cargo run --example transaction -- --pd="127.0.0.1:2379"; fi
   - if [[ $TRAVIS_OS_NAME == "linux" && $TRAVIS_RUST_VERSION == "nightly" ]]; then cargo run --example pessimistic -- --pd="127.0.0.1:2379"; fi

Cargo.toml

@@ -23,7 +23,7 @@ derive-new = "0.5"
 fail = "0.4"
 futures = { version = "0.3", features = ["async-await", "thread-pool"] }
 futures-timer = "3.0"
-grpcio = { version = "0.8", features = [ "secure", "prost-codec", "use-bindgen" ], default-features = false }
+grpcio = { version = "0.8", features = [ "secure", "prost-codec", "use-bindgen", "openssl-vendored" ], default-features = false }
 lazy_static = "1"
 log = "0.4"
 prometheus = { version = "0.12", features = [ "push", "process" ], default-features = false } 
@@ -49,6 +49,8 @@ proptest-derive = "0.3"
 serial_test = "0.5.0"
 simple_logger = "1"
 tokio = { version = "1.0", features = [ "sync", "rt-multi-thread", "macros" ] }
+reqwest = {version = "0.11", default-features = false, features = ["native-tls-vendored"]}
+serde_json = "1"
 
 [workspace]
 members = [

Makefile

@@ -0,0 +1,32 @@
+.PHONY: default check unit-test integration-tests test doc docker-pd docker-kv docker all
+
+default: check
+
+check:
+	cargo check --all
+	cargo fmt -- --check
+	cargo clippy -- -D clippy::all
+
+unit-test:
+	cargo test --all
+
+integration-test:
+# MULTI_REGION shall be set manually if needed
+	PD_ADDRS="127.0.0.1:2379" cargo test txn_ --all --features integration-tests -- --nocapture
+	PD_ADDRS="127.0.0.1:2379" cargo test raw_ --all --features integration-tests -- --nocapture
+	PD_ADDRS="127.0.0.1:2379" cargo test misc_ --all --features integration-tests -- --nocapture
+
+test: unit-test integration-test
+
+doc: 
+	cargo doc --workspace --exclude tikv-client-proto --document-private-items --no-deps
+
+docker-pd:
+	docker run -d -v $(shell pwd)/config:/config --net=host --name pd --rm pingcap/pd:latest --name "pd" --data-dir "pd" --client-urls "http://127.0.0.1:2379" --advertise-client-urls "http://127.0.0.1:2379" --config /config/pd.toml
+
+docker-kv:
+	docker run -d -v $(shell pwd)/config:/config --net=host --name kv --rm --ulimit nofile=90000:90000 pingcap/tikv:latest --pd-endpoints "127.0.0.1:2379" --addr "127.0.0.1:2378" --data-dir "kv" --config /config/tikv.toml
+
+docker: docker-pd docker-kv
+
+all: check doc test

config/pd.toml

@@ -0,0 +1,3 @@
+[schedule]
+max-merge-region-size = 1
+max-merge-region-keys = 3

config/tikv.toml

@@ -0,0 +1,10 @@
+[coprocessor]
+region-max-keys = 5
+region-split-keys = 3
+batch-split-limit = 100
+
+[raftstore]
+region-split-check-diff = "0KB"
+pd-heartbeat-tick-interval = "2s"
+pd-store-heartbeat-tick-interval = "5s"
+split-region-check-tick-interval = "1s"

src/pd/retry.rs

@@ -25,7 +25,7 @@ use tokio::sync::RwLock;
 // FIXME: these numbers and how they are used are all just cargo-culted in, there
 // may be more optimal values.
 const RECONNECT_INTERVAL_SEC: u64 = 1;
-const MAX_REQUEST_COUNT: usize = 3;
+const MAX_REQUEST_COUNT: usize = 5;
 const LEADER_CHANGE_RETRY: usize = 10;
 
 /// Client for communication with a PD cluster. Has the facility to reconnect to the cluster.

src/util/iter.rs

@@ -74,7 +74,7 @@ mod test {
             .into_iter()
             .flat_map_ok(|i| Some(i).into_iter())
             .collect();
-        assert_eq!(result.unwrap(), vec![]);
+        assert_eq!(result.unwrap(), Vec::<i32>::new());
 
         let result: Result<Vec<i32>, ()> = vec![Result::<i32, ()>::Ok(0), Ok(1), Ok(2)]
             .into_iter()

tests/common/ctl.rs

@@ -0,0 +1,28 @@
+// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
+//! The module provides some utility functions to control and get information
+//! from PD, using its HTTP API.
+
+use super::pd_addrs;
+use tikv_client_common::{Error, Result};
+
+pub async fn get_region_count() -> Result<u64> {
+    let res = reqwest::get(format!("http://{}/pd/api/v1/regions", pd_addrs()[0]))
+        .await
+        .map_err(|e| Error::StringError(e.to_string()))?;
+
+    let body = res
+        .text()
+        .await
+        .map_err(|e| Error::StringError(e.to_string()))?;
+    let value: serde_json::Value = serde_json::from_str(body.as_ref()).unwrap();
+    value["count"]
+        .as_u64()
+        .ok_or_else(|| Error::StringError("pd region count does not return an integer".to_owned()))
+}
+
+#[tokio::test]
+#[ignore]
+async fn test_get_region_count() -> Result<()> {
+    println!("{}", get_region_count().await?);
+    Ok(())
+}

tests/common/mod.rs

@@ -1,5 +1,13 @@
-use std::env;
-use tikv_client::{ColumnFamily, RawClient};
+mod ctl;
+
+use futures_timer::Delay;
+use log::{info, warn};
+use std::{env, time::Duration};
+use tikv_client::{ColumnFamily, Key, RawClient, Result, TransactionClient};
+
+const ENV_PD_ADDRS: &str = "PD_ADDRS";
+const ENV_ENABLE_MULIT_REGION: &str = "MULTI_REGION";
+const REGION_SPLIT_TIME_LIMIT: Duration = Duration::from_secs(15);
 
 // Delete all entries in TiKV to leave a clean space for following tests.
 pub async fn clear_tikv() {
@@ -14,7 +22,69 @@ pub async fn clear_tikv() {
     }
 }
 
-const ENV_PD_ADDRS: &str = "PD_ADDRS";
+// To test with multiple regions, prewrite some data. Tests that hope to test
+// with multiple regions should use keys in the corresponding ranges.
+pub async fn init() -> Result<()> {
+    // ignore SetLoggerError
+    let _ = simple_logger::SimpleLogger::new()
+        .with_level(log::LevelFilter::Warn)
+        .init();
+
+    if env::var(ENV_ENABLE_MULIT_REGION).is_ok() {
+        // 1000 keys: 0..1000
+        let keys_1 = std::iter::successors(Some(0u32), |x| Some(x + 1))
+            .take(1000)
+            .map(|x| x.to_be_bytes().to_vec());
+        // 1024 keys: 0..u32::MAX
+        let count = 1024;
+        let step = u32::MAX / count;
+        let keys_2 = std::iter::successors(Some(0u32), |x| Some(x + step))
+            .take(count as usize - 1)
+            .map(|x| x.to_be_bytes().to_vec());
+
+        ensure_region_split(keys_1.chain(keys_2), 100).await?;
+    }
+
+    clear_tikv().await;
+    Ok(())
+}
+
+async fn ensure_region_split(
+    keys: impl IntoIterator<Item = impl Into<Key>>,
+    region_count: u32,
+) -> Result<()> {
+    if ctl::get_region_count().await? as u32 >= region_count {
+        return Ok(());
+    }
+
+    // 1. write plenty transactional keys
+    // 2. wait until regions split
+
+    let client = TransactionClient::new(pd_addrs()).await?;
+    let mut txn = client.begin_optimistic().await?;
+    for key in keys.into_iter() {
+        txn.put(key.into(), vec![0, 0, 0, 0]).await?;
+    }
+    txn.commit().await?;
+    let mut txn = client.begin_optimistic().await?;
+    let _ = txn.scan(vec![].., 2048).await?;
+    txn.commit().await?;
+
+    info!("splitting regions...");
+    let start_time = std::time::Instant::now();
+    loop {
+        if ctl::get_region_count().await? as u32 >= region_count {
+            break;
+        }
+        if start_time.elapsed() > REGION_SPLIT_TIME_LIMIT {
+            warn!("Stop splitting regions: time limit exceeded");
+            break;
+        }
+        Delay::new(Duration::from_millis(200)).await;
+    }
+
+    Ok(())
+}
 
 pub fn pd_addrs() -> Vec<String> {
     env::var(ENV_PD_ADDRS)

tests/failpoint_tests.rs

@@ -1,15 +1,15 @@
 #![cfg(feature = "integration-tests")]
 
 mod common;
-use common::{clear_tikv, pd_addrs};
+use common::{init, pd_addrs};
 use fail::FailScenario;
 use serial_test::serial;
 use tikv_client::{Result, TransactionClient, TransactionOptions};
 
 #[tokio::test]
 #[serial]
-async fn optimistic_heartbeat() -> Result<()> {
-    clear_tikv().await;
+async fn txn_optimistic_heartbeat() -> Result<()> {
+    init().await?;
     let scenario = FailScenario::setup();
     fail::cfg("after-prewrite", "sleep(10000)").unwrap();