test: test multi-region batch_get_for_update

Signed-off-by: ekexium <ekexium@gmail.com>
This commit is contained in:
ekexium 2021-03-15 16:14:24 +08:00
parent 555981eae7
commit 260bcd70cd
1 changed files with 29 additions and 3 deletions

View File

@ -225,12 +225,15 @@ async fn raw_bank_transfer() -> Result<()> {
/// Tests transactional API when there are multiple regions.
/// Write large volumes of data to enforce region splitting.
/// In order to test `scan`, data is uniformly inserted.
// FIXME: this test is stupid. We should use pd-ctl or config files to make
// multiple regions, instead of bulk writing.
#[tokio::test]
#[serial]
async fn txn_write_million() -> Result<()> {
const NUM_BITS_TXN: u32 = 7;
const NUM_BITS_KEY_PER_TXN: u32 = 3;
const NUM_BITS_TXN: u32 = 12;
const NUM_BITS_KEY_PER_TXN: u32 = 5;
let interval = 2u32.pow(32 - NUM_BITS_TXN - NUM_BITS_KEY_PER_TXN);
let value = "large_value".repeat(10);
clear_tikv().await;
let client = TransactionClient::new(pd_addrs()).await?;
@ -246,7 +249,7 @@ async fn txn_write_million() -> Result<()> {
.take(2usize.pow(NUM_BITS_KEY_PER_TXN))
.collect::<Vec<_>>(); // each txn puts 2 ^ 12 keys. 12 = 25 - 13
let mut txn = client.begin_optimistic().await?;
for (k, v) in keys.iter().zip(iter::repeat(1u32.to_be_bytes().to_vec())) {
for (k, v) in keys.iter().zip(iter::repeat(value.clone())) {
txn.put(k.clone(), v).await?;
}
txn.commit().await?;
@ -298,6 +301,29 @@ async fn txn_write_million() -> Result<()> {
assert_eq!(sum, 2usize.pow(NUM_BITS_KEY_PER_TXN + NUM_BITS_TXN));
// test batch_get and batch_get_for_update
const SKIP_BITS: u32 = 6; // do not retrive all because there's a limit of message size
let mut cur = 0u32;
let keys = iter::repeat_with(|| {
let v = cur;
cur = cur.overflowing_add(interval * 2u32.pow(SKIP_BITS)).0;
v
})
.map(|u| u.to_be_bytes().to_vec())
.take(2usize.pow(NUM_BITS_KEY_PER_TXN + NUM_BITS_TXN - SKIP_BITS))
.collect::<Vec<_>>();
let mut txn = client.begin_pessimistic().await?;
let res = txn.batch_get(keys.clone()).await?.collect::<Vec<_>>();
assert_eq!(res.len(), keys.len());
let res = txn
.batch_get_for_update(keys.clone())
.await?
.collect::<Vec<_>>();
assert_eq!(res.len(), keys.len());
txn.commit().await?;
Ok(())
}