From 0d05446976f34ad6a303b5dc5d0ad0ac7cac1999 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Mon, 18 Jul 2022 18:56:08 +0100
Subject: [PATCH] Ensure TiKV transactions do not panic when a transaction is
 dropped

---
 lib/src/kvs/tikv/mod.rs | 44 ++++++++++++++++++++++++++---------------
 1 file changed, 28 insertions(+), 16 deletions(-)

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