Add try one pc to config

Signed-off-by: Nick Cameron <nrc@ncameron.org>

src/config.rs

@@ -27,6 +27,7 @@ pub struct Config {
     pub cert_path: Option<PathBuf>,
     pub key_path: Option<PathBuf>,
     pub timeout: Duration,
+    pub try_one_pc: bool,
 }
 
 const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2);
@@ -38,6 +39,7 @@ impl Default for Config {
             cert_path: None,
             key_path: None,
             timeout: DEFAULT_REQUEST_TIMEOUT,
+            try_one_pc: false,
         }
     }
 }
@@ -79,4 +81,9 @@ impl Config {
         self.timeout = timeout;
         self
     }
+
+    pub fn try_one_pc(mut self) -> Self {
+        self.try_one_pc = true;
+        self
+    }
 }

src/transaction/client.rs

@@ -37,6 +37,7 @@ pub struct Client {
     /// The thread pool for background tasks including committing secondary keys and failed
     /// transaction cleanups.
     bg_worker: ThreadPool,
+    config: Config,
 }
 
 impl Client {
@@ -77,7 +78,11 @@ impl Client {
         let pd_endpoints: Vec<String> = pd_endpoints.into_iter().map(Into::into).collect();
         let bg_worker = ThreadPool::new()?;
         let pd = Arc::new(PdRpcClient::connect(&pd_endpoints, &config, true).await?);
-        Ok(Client { pd, bg_worker })
+        Ok(Client {
+            pd,
+            bg_worker,
+            config,
+        })
     }
 
     /// Creates a new [`Transaction`](Transaction) in optimistic mode.
@@ -102,7 +107,11 @@ impl Client {
     /// ```
     pub async fn begin(&self) -> Result<Transaction> {
         let timestamp = self.current_timestamp().await?;
-        Ok(self.new_transaction(timestamp, TransactionStyle::new_optimistic(), false))
+        Ok(self.new_transaction(
+            timestamp,
+            TransactionStyle::new_optimistic(self.config.try_one_pc),
+            false,
+        ))
     }
 
     /// Creates a new [`Transaction`](Transaction) in pessimistic mode.
@@ -124,12 +133,20 @@ impl Client {
     /// ```
     pub async fn begin_pessimistic(&self) -> Result<Transaction> {
         let timestamp = self.current_timestamp().await?;
-        Ok(self.new_transaction(timestamp, TransactionStyle::new_pessimistic(), false))
+        Ok(self.new_transaction(
+            timestamp,
+            TransactionStyle::new_pessimistic(self.config.try_one_pc),
+            false,
+        ))
     }
 
     /// Creates a new [`Snapshot`](Snapshot) at the given [`Timestamp`](Timestamp).
     pub fn snapshot(&self, timestamp: Timestamp) -> Snapshot {
-        Snapshot::new(self.new_transaction(timestamp, TransactionStyle::new_optimistic(), true))
+        Snapshot::new(self.new_transaction(
+            timestamp,
+            TransactionStyle::new_optimistic(self.config.try_one_pc),
+            true,
+        ))
     }
 
     /// Retrieves the current [`Timestamp`](Timestamp).

src/transaction/transaction.rs

@@ -43,27 +43,22 @@ pub struct TransactionStyle {
 }
 
 impl TransactionStyle {
-    pub fn new_optimistic() -> TransactionStyle {
+    pub fn new_optimistic(try_one_pc: bool) -> TransactionStyle {
         TransactionStyle {
             kind: TransactionKind::Optimistic,
-            try_one_pc: false,
+            try_one_pc,
             async_commit: false,
         }
     }
 
-    pub fn new_pessimistic() -> TransactionStyle {
+    pub fn new_pessimistic(try_one_pc: bool) -> TransactionStyle {
         TransactionStyle {
             kind: TransactionKind::Pessimistic(0),
-            try_one_pc: false,
+            try_one_pc,
             async_commit: false,
         }
     }
 
-    pub fn try_one_pc(mut self) -> TransactionStyle {
-        self.try_one_pc = true;
-        self
-    }
-
     pub fn async_commit(mut self) -> TransactionStyle {
         self.async_commit = true;
         self