introduce codec type parameter for Snapshot and Transcation

Signed-off-by: iosmanthus <[email protected]>

Author: iosmanthus

src/transaction/client.rs

@@ -158,7 +158,7 @@ where
     /// transaction.commit().await.unwrap();
     /// # });
     /// ```
-    pub async fn begin_optimistic(&self) -> Result<Transaction<PdRpcClient<C>>> {
+    pub async fn begin_optimistic(&self) -> Result<Transaction<C>> {
         debug!(self.logger, "creating new optimistic transaction");
         let timestamp = self.current_timestamp().await?;
         Ok(self.new_transaction(timestamp, TransactionOptions::new_optimistic()))
@@ -184,7 +184,7 @@ where
     /// transaction.commit().await.unwrap();
     /// # });
     /// ```
-    pub async fn begin_pessimistic(&self) -> Result<Transaction<PdRpcClient<C>>> {
+    pub async fn begin_pessimistic(&self) -> Result<Transaction<C>> {
         debug!(self.logger, "creating new pessimistic transaction");
         let timestamp = self.current_timestamp().await?;
         Ok(self.new_transaction(timestamp, TransactionOptions::new_pessimistic()))
@@ -210,21 +210,14 @@ where
     /// transaction.commit().await.unwrap();
     /// # });
     /// ```
-    pub async fn begin_with_options(
-        &self,
-        options: TransactionOptions,
-    ) -> Result<Transaction<PdRpcClient<C>>> {
+    pub async fn begin_with_options(&self, options: TransactionOptions) -> Result<Transaction<C>> {
         debug!(self.logger, "creating new customized transaction");
         let timestamp = self.current_timestamp().await?;
         Ok(self.new_transaction(timestamp, options))
     }
 
     /// Create a new [`Snapshot`](Snapshot) at the given [`Timestamp`](Timestamp).
-    pub fn snapshot(
-        &self,
-        timestamp: Timestamp,
-        options: TransactionOptions,
-    ) -> Snapshot<PdRpcClient<C>> {
+    pub fn snapshot(&self, timestamp: Timestamp, options: TransactionOptions) -> Snapshot<C> {
         debug!(self.logger, "creating new snapshot");
         let logger = self.logger.new(o!("child" => 1));
         Snapshot::new(self.new_transaction(timestamp, options.read_only()), logger)
@@ -306,11 +299,7 @@ where
         Ok(res)
     }
 
-    fn new_transaction(
-        &self,
-        timestamp: Timestamp,
-        options: TransactionOptions,
-    ) -> Transaction<PdRpcClient<C>> {
+    fn new_transaction(&self, timestamp: Timestamp, options: TransactionOptions) -> Transaction<C> {
         let logger = self.logger.new(o!("child" => 1));
         Transaction::new(timestamp, self.pd.clone(), options, logger)
     }

src/transaction/snapshot.rs

@@ -1,8 +1,13 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
-use crate::{pd::PdClient, BoundRange, Key, KvPair, Result, Transaction, Value};
+use crate::{
+    pd::{PdClient, PdRpcClient},
+    request::codec::RequestCodec,
+    BoundRange, Key, KvPair, Result, Transaction, Value,
+};
 use derive_new::new;
 use slog::Logger;
+use std::marker::PhantomData;
 
 /// A read-only transaction which reads at the given timestamp.
 ///
@@ -12,12 +17,17 @@ use slog::Logger;
 ///
 /// See the [Transaction](struct@crate::Transaction) docs for more information on the methods.
 #[derive(new)]
-pub struct Snapshot<PdC: PdClient> {
-    transaction: Transaction<PdC>,
+pub struct Snapshot<C: RequestCodec, PdC: PdClient<RequestCodec = C> = PdRpcClient<C>> {
+    transaction: Transaction<C, PdC>,
     logger: Logger,
+    _phantom: PhantomData<C>,
 }
 
-impl<PdC: PdClient> Snapshot<PdC> {
+impl<C, PdC> Snapshot<C, PdC>
+where
+    C: RequestCodec,
+    PdC: PdClient<RequestCodec = C>,
+{
     /// Get the value associated with the given key.
     pub async fn get(&mut self, key: impl Into<Key>) -> Result<Option<Value>> {
         debug!(self.logger, "invoking get request on snapshot");

src/transaction/transaction.rs

@@ -1,6 +1,6 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
-use std::{iter, sync::Arc, time::Instant};
+use std::{iter, marker::PhantomData, sync::Arc, time::Instant};
 
 use derive_new::new;
 use fail::fail_point;
@@ -12,9 +12,10 @@ use tikv_client_proto::{kvrpcpb, pdpb::Timestamp};
 
 use crate::{
     backoff::{Backoff, DEFAULT_REGION_BACKOFF},
-    pd::PdClient,
+    pd::{PdClient, PdRpcClient},
     request::{
-        Collect, CollectError, CollectSingle, CollectWithShard, Plan, PlanBuilder, RetryOptions,
+        codec::RequestCodec, Collect, CollectError, CollectSingle, CollectWithShard, Plan,
+        PlanBuilder, RetryOptions,
     },
     timestamp::TimestampExt,
     transaction::{buffer::Buffer, lowering::*},
@@ -62,7 +63,7 @@ use crate::{
 /// txn.commit().await.unwrap();
 /// # });
 /// ```
-pub struct Transaction<PdC: PdClient> {
+pub struct Transaction<C: RequestCodec, PdC: PdClient<RequestCodec = C> = PdRpcClient<C>> {
     status: Arc<RwLock<TransactionStatus>>,
     timestamp: Timestamp,
     buffer: Buffer,
@@ -71,15 +72,20 @@ pub struct Transaction<PdC: PdClient> {
     is_heartbeat_started: bool,
     start_instant: Instant,
     logger: Logger,
+    _phantom: PhantomData<C>,
 }
 
-impl<PdC: PdClient> Transaction<PdC> {
+impl<C, PdC> Transaction<C, PdC>
+where
+    C: RequestCodec,
+    PdC: PdClient<RequestCodec = C>,
+{
     pub(crate) fn new(
         timestamp: Timestamp,
         rpc: Arc<PdC>,
         options: TransactionOptions,
         logger: Logger,
-    ) -> Transaction<PdC> {
+    ) -> Transaction<C, PdC> {
         let status = if options.read_only {
             TransactionStatus::ReadOnly
         } else {
@@ -94,6 +100,7 @@ impl<PdC: PdClient> Transaction<PdC> {
             is_heartbeat_started: false,
             start_instant: std::time::Instant::now(),
             logger,
+            _phantom: PhantomData,
         }
     }
 
@@ -931,7 +938,11 @@ impl<PdC: PdClient> Transaction<PdC> {
     }
 }
 
-impl<PdC: PdClient> Drop for Transaction<PdC> {
+impl<C, PdC> Drop for Transaction<C, PdC>
+where
+    C: RequestCodec,
+    PdC: PdClient<RequestCodec = C>,
+{
     fn drop(&mut self) {
         debug!(self.logger, "dropping transaction");
         if std::thread::panicking() {

tests/integration_tests.rs

@@ -23,9 +23,9 @@ use std::{
 };
 use tikv_client::{
     raw,
-    request::codec::RawCodec,
+    request::codec::{RawCodec, RequestCodec},
     transaction::{self, HeartbeatOption},
-    BoundRange, Error, Key, KvPair, PdClient, RawClient, Result, Transaction, TransactionClient,
+    BoundRange, Error, Key, KvPair, RawClient, Result, Transaction, TransactionClient,
     TransactionOptions, Value,
 };
 
@@ -973,7 +973,7 @@ async fn get_u32<C: RawCodec>(client: &RawClient<C>, key: Vec<u8>) -> Result<u32
 }
 
 // helper function
-async fn get_txn_u32<C: PdClient>(txn: &mut Transaction<C>, key: Vec<u8>) -> Result<u32> {
+async fn get_txn_u32<C: RequestCodec>(txn: &mut Transaction<C>, key: Vec<u8>) -> Result<u32> {
     let x = txn.get(key).await?.unwrap();
     let boxed_slice = x.into_boxed_slice();
     let array: Box<[u8; 4]> = boxed_slice