Ensure TiKV transactions do not panic when a transaction is dropped

This commit is contained in:
Tobie Morgan Hitchcock 2022-07-18 18:56:08 +01:00
parent 10bc3b2ddd
commit 0d05446976

View file

@ -4,6 +4,8 @@ use crate::err::Error;
use crate::kvs::Key; use crate::kvs::Key;
use crate::kvs::Val; use crate::kvs::Val;
use std::ops::Range; use std::ops::Range;
use tikv::CheckLevel;
use tikv::TransactionOptions;
pub struct Datastore { pub struct Datastore {
db: tikv::TransactionClient, db: tikv::TransactionClient,
@ -31,22 +33,32 @@ impl Datastore {
// Start a new transaction // Start a new transaction
pub async fn transaction(&self, write: bool, lock: bool) -> Result<Transaction, Error> { pub async fn transaction(&self, write: bool, lock: bool) -> Result<Transaction, Error> {
match lock { match lock {
true => match self.db.begin_optimistic().await { true => {
Ok(tx) => Ok(Transaction { // Set the behaviour when dropping an unfinished transaction
ok: false, let opt = TransactionOptions::new_optimistic().drop_check(CheckLevel::Warn);
rw: write, // Create a new optimistic transaction
tx, match self.db.begin_with_options(opt).await {
}), Ok(tx) => Ok(Transaction {
Err(e) => Err(Error::Tx(e.to_string())), ok: false,
}, rw: write,
false => match self.db.begin_pessimistic().await { tx,
Ok(tx) => Ok(Transaction { }),
ok: false, Err(e) => Err(Error::Tx(e.to_string())),
rw: write, }
tx, }
}), false => {
Err(e) => Err(Error::Tx(e.to_string())), // Set the behaviour when dropping an unfinished transaction
}, let opt = TransactionOptions::new_pessimistic().drop_check(CheckLevel::Warn);
// Create a new pessimistic transaction
match self.db.begin_with_options(opt).await {
Ok(tx) => Ok(Transaction {
ok: false,
rw: write,
tx,
}),
Err(e) => Err(Error::Tx(e.to_string())),
}
}
} }
} }
} }