Merge pull request #262 from ekexium/atomic

Support atomic operations (CAS, put, and delete)

Author: nrc

mock-tikv/src/pd.rs

@@ -336,11 +336,11 @@ impl Pd for MockPd {
         todo!()
     }
 
-    fn get_dc_locations(
+    fn get_dc_location_info(
         &mut self,
-        _: ::grpcio::RpcContext<'_>,
-        _: GetDcLocationsRequest,
-        _: ::grpcio::UnarySink<GetDcLocationsResponse>,
+        _: grpcio::RpcContext,
+        _: GetDcLocationInfoRequest,
+        _: grpcio::UnarySink<GetDcLocationInfoResponse>,
     ) {
         todo!()
     }

mock-tikv/src/server.rs

@@ -5,7 +5,7 @@ use derive_new::new;
 use futures::{FutureExt, TryFutureExt};
 use grpcio::{Environment, Server, ServerBuilder};
 use std::sync::Arc;
-use tikv_client_proto::{kvrpcpb::*, tikvpb::*};
+use tikv_client_proto::{coprocessor_v2::*, kvrpcpb::*, tikvpb::*};
 
 pub const MOCK_TIKV_PORT: u16 = 50019;
 
@@ -514,11 +514,38 @@ impl Tikv for MockTikv {
         todo!()
     }
 
+    fn raw_get_key_ttl(
+        &mut self,
+        _: grpcio::RpcContext,
+        _: RawGetKeyTtlRequest,
+        _: grpcio::UnarySink<RawGetKeyTtlResponse>,
+    ) {
+        todo!()
+    }
+
+    fn raw_compare_and_swap(
+        &mut self,
+        _: grpcio::RpcContext,
+        _: RawCasRequest,
+        _: grpcio::UnarySink<RawCasResponse>,
+    ) {
+        todo!()
+    }
+
     fn coprocessor_v2(
         &mut self,
         _: grpcio::RpcContext,
-        _: tikv_client_proto::coprocessor_v2::RawCoprocessorRequest,
-        _: grpcio::UnarySink<tikv_client_proto::coprocessor_v2::RawCoprocessorResponse>,
+        _: RawCoprocessorRequest,
+        _: grpcio::UnarySink<RawCoprocessorResponse>,
+    ) {
+        todo!()
+    }
+
+    fn get_store_safe_ts(
+        &mut self,
+        _: grpcio::RpcContext,
+        _: StoreSafeTsRequest,
+        _: grpcio::UnarySink<StoreSafeTsResponse>,
     ) {
         todo!()
     }

src/raw/client.rs

@@ -24,6 +24,8 @@ const MAX_RAW_KV_SCAN_LIMIT: u32 = 10240;
 pub struct Client {
     rpc: Arc<PdRpcClient>,
     cf: Option<ColumnFamily>,
+    /// Whether to use the [`atomic mode`](Client::with_atomic_for_cas).
+    atomic: bool,
 }
 
 impl Client {
@@ -63,7 +65,11 @@ impl Client {
     ) -> Result<Client> {
         let pd_endpoints: Vec<String> = pd_endpoints.into_iter().map(Into::into).collect();
         let rpc = Arc::new(PdRpcClient::connect(&pd_endpoints, &config, false).await?);
-        Ok(Client { rpc, cf: None })
+        Ok(Client {
+            rpc,
+            cf: None,
+            atomic: false,
+        })
     }
 
     /// Set the column family of requests.
@@ -89,6 +95,22 @@ impl Client {
         Client {
             rpc: self.rpc.clone(),
             cf: Some(cf),
+            atomic: self.atomic,
+        }
+    }
+
+    /// Set to use the atomic mode.
+    ///
+    /// The only reason of using atomic mode is the
+    /// [`compare_and_swap`](Client::compare_and_swap) operation. To guarantee
+    /// the atomicity of CAS, write operations like [`put`](Client::put) or
+    /// [`delete`](Client::delete) in atomic mode are more expensive. Some
+    /// operations are not supported in the mode.
+    pub fn with_atomic_for_cas(&self) -> Client {
+        Client {
+            rpc: self.rpc.clone(),
+            cf: self.cf.clone(),
+            atomic: true,
         }
     }
 
@@ -171,7 +193,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn put(&self, key: impl Into<Key>, value: impl Into<Value>) -> Result<()> {
-        let request = new_raw_put_request(key.into(), value.into(), self.cf.clone());
+        let request = new_raw_put_request(key.into(), value.into(), self.cf.clone(), self.atomic);
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .single_region()
             .await?
@@ -203,7 +225,11 @@ impl Client {
         &self,
         pairs: impl IntoIterator<Item = impl Into<KvPair>>,
     ) -> Result<()> {
-        let request = new_raw_batch_put_request(pairs.into_iter().map(Into::into), self.cf.clone());
+        let request = new_raw_batch_put_request(
+            pairs.into_iter().map(Into::into),
+            self.cf.clone(),
+            self.atomic,
+        );
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .multi_region()
             .retry_region(DEFAULT_REGION_BACKOFF)
@@ -231,7 +257,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn delete(&self, key: impl Into<Key>) -> Result<()> {
-        let request = new_raw_delete_request(key.into(), self.cf.clone());
+        let request = new_raw_delete_request(key.into(), self.cf.clone(), self.atomic);
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .single_region()
             .await?
@@ -260,6 +286,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn batch_delete(&self, keys: impl IntoIterator<Item = impl Into<Key>>) -> Result<()> {
+        self.assert_non_atomic()?;
         let request =
             new_raw_batch_delete_request(keys.into_iter().map(Into::into), self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
@@ -287,6 +314,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn delete_range(&self, range: impl Into<BoundRange>) -> Result<()> {
+        self.assert_non_atomic()?;
         let request = new_raw_delete_range_request(range.into(), self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .multi_region()
@@ -415,6 +443,40 @@ impl Client {
             .collect())
     }
 
+    /// Create a new *atomic* 'compare and set' request.
+    ///
+    /// Once resolved this request will result in an atomic `compare and set'
+    /// operation for the given key.
+    ///
+    /// If the value retrived is equal to `current_value`, `new_value` is
+    /// written.
+    ///
+    /// # Return Value
+    ///
+    /// A tuple is returned if successful: the previous value and whether the
+    /// value is swapped
+    pub async fn compare_and_swap(
+        &self,
+        key: impl Into<Key>,
+        previous_value: impl Into<Option<Value>>,
+        new_value: impl Into<Value>,
+    ) -> Result<(Option<Value>, bool)> {
+        self.assert_atomic()?;
+        let req = new_cas_request(
+            key.into(),
+            new_value.into(),
+            previous_value.into(),
+            self.cf.clone(),
+        );
+        let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req)
+            .single_region()
+            .await?
+            .retry_region(DEFAULT_REGION_BACKOFF)
+            .post_process_default()
+            .plan();
+        plan.execute().await
+    }
+
     async fn scan_inner(
         &self,
         range: impl Into<BoundRange>,
@@ -467,4 +529,12 @@ impl Client {
             .plan();
         plan.execute().await
     }
+
+    fn assert_non_atomic(&self) -> Result<()> {
+        (!self.atomic).then(|| ()).ok_or(Error::UnsupportedMode)
+    }
+
+    fn assert_atomic(&self) -> Result<()> {
+        self.atomic.then(|| ()).ok_or(Error::UnsupportedMode)
+    }
 }

src/raw/lowering.rs

@@ -22,19 +22,25 @@ pub fn new_raw_put_request(
     key: Key,
     value: Value,
     cf: Option<ColumnFamily>,
+    atomic: bool,
 ) -> kvrpcpb::RawPutRequest {
-    requests::new_raw_put_request(key.into(), value, cf)
+    requests::new_raw_put_request(key.into(), value, cf, atomic)
 }
 
 pub fn new_raw_batch_put_request(
     pairs: impl Iterator<Item = KvPair>,
     cf: Option<ColumnFamily>,
+    atomic: bool,
 ) -> kvrpcpb::RawBatchPutRequest {
-    requests::new_raw_batch_put_request(pairs.map(Into::into).collect(), cf)
+    requests::new_raw_batch_put_request(pairs.map(Into::into).collect(), cf, atomic)
 }
 
-pub fn new_raw_delete_request(key: Key, cf: Option<ColumnFamily>) -> kvrpcpb::RawDeleteRequest {
-    requests::new_raw_delete_request(key.into(), cf)
+pub fn new_raw_delete_request(
+    key: Key,
+    cf: Option<ColumnFamily>,
+    atomic: bool,
+) -> kvrpcpb::RawDeleteRequest {
+    requests::new_raw_delete_request(key.into(), cf, atomic)
 }
 
 pub fn new_raw_batch_delete_request(
@@ -76,3 +82,12 @@ pub fn new_raw_batch_scan_request(
 ) -> kvrpcpb::RawBatchScanRequest {
     requests::new_raw_batch_scan_request(ranges.map(Into::into).collect(), each_limit, key_only, cf)
 }
+
+pub fn new_cas_request(
+    key: Key,
+    value: Value,
+    previous_value: Option<Value>,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawCasRequest {
+    requests::new_cas_request(key.into(), value, previous_value, cf)
+}

src/raw/requests.rs

@@ -76,11 +76,13 @@ pub fn new_raw_put_request(
     key: Vec<u8>,
     value: Vec<u8>,
     cf: Option<ColumnFamily>,
+    atomic: bool,
 ) -> kvrpcpb::RawPutRequest {
     let mut req = kvrpcpb::RawPutRequest::default();
     req.set_key(key);
     req.set_value(value);
     req.maybe_set_cf(cf);
+    req.set_for_cas(atomic);
 
     req
 }
@@ -98,10 +100,12 @@ impl SingleKey for kvrpcpb::RawPutRequest {
 pub fn new_raw_batch_put_request(
     pairs: Vec<kvrpcpb::KvPair>,
     cf: Option<ColumnFamily>,
+    atomic: bool,
 ) -> kvrpcpb::RawBatchPutRequest {
     let mut req = kvrpcpb::RawBatchPutRequest::default();
     req.set_pairs(pairs);
     req.maybe_set_cf(cf);
+    req.set_for_cas(atomic);
 
     req
 }
@@ -132,10 +136,15 @@ impl Shardable for kvrpcpb::RawBatchPutRequest {
     }
 }
 
-pub fn new_raw_delete_request(key: Vec<u8>, cf: Option<ColumnFamily>) -> kvrpcpb::RawDeleteRequest {
+pub fn new_raw_delete_request(
+    key: Vec<u8>,
+    cf: Option<ColumnFamily>,
+    atomic: bool,
+) -> kvrpcpb::RawDeleteRequest {
     let mut req = kvrpcpb::RawDeleteRequest::default();
     req.set_key(key);
     req.maybe_set_cf(cf);
+    req.set_for_cas(atomic);
 
     req
 }
@@ -267,6 +276,46 @@ impl Merge<kvrpcpb::RawBatchScanResponse> for Collect {
     }
 }
 
+pub fn new_cas_request(
+    key: Vec<u8>,
+    value: Vec<u8>,
+    previous_value: Option<Vec<u8>>,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawCasRequest {
+    let mut req = kvrpcpb::RawCasRequest::default();
+    req.set_key(key);
+    req.set_value(value);
+    match previous_value {
+        Some(v) => req.set_previous_value(v),
+        None => req.set_previous_not_exist(true),
+    }
+    req.maybe_set_cf(cf);
+    req
+}
+
+impl KvRequest for kvrpcpb::RawCasRequest {
+    type Response = kvrpcpb::RawCasResponse;
+}
+
+impl SingleKey for kvrpcpb::RawCasRequest {
+    fn key(&self) -> &Vec<u8> {
+        &self.key
+    }
+}
+
+impl Process<kvrpcpb::RawCasResponse> for DefaultProcessor {
+    type Out = (Option<Value>, bool); // (previous_value, swapped)
+
+    fn process(&self, input: Result<kvrpcpb::RawCasResponse>) -> Result<Self::Out> {
+        let input = input?;
+        if input.previous_not_exist {
+            Ok((None, input.succeed))
+        } else {
+            Ok((Some(input.previous_value), input.succeed))
+        }
+    }
+}
+
 macro_rules! impl_raw_rpc_request {
     ($name: ident) => {
         impl RawRpcRequest for kvrpcpb::$name {
@@ -286,6 +335,7 @@ impl_raw_rpc_request!(RawBatchDeleteRequest);
 impl_raw_rpc_request!(RawScanRequest);
 impl_raw_rpc_request!(RawBatchScanRequest);
 impl_raw_rpc_request!(RawDeleteRangeRequest);
+impl_raw_rpc_request!(RawCasRequest);
 
 impl HasLocks for kvrpcpb::RawGetResponse {}
 impl HasLocks for kvrpcpb::RawBatchGetResponse {}
@@ -296,6 +346,7 @@ impl HasLocks for kvrpcpb::RawBatchDeleteResponse {}
 impl HasLocks for kvrpcpb::RawScanResponse {}
 impl HasLocks for kvrpcpb::RawBatchScanResponse {}
 impl HasLocks for kvrpcpb::RawDeleteRangeResponse {}
+impl HasLocks for kvrpcpb::RawCasResponse {}
 
 #[cfg(test)]
 mod test {

src/request/plan.rs

@@ -276,8 +276,11 @@ impl<P: Plan + HasKeys, PdC: PdClient> HasKeys for ResolveLock<P, PdC> {
     }
 }
 
-/// When executed, the plan extracts errors from its inner plan, and
-/// returns an `Err` wrapping the error.
+/// When executed, the plan extracts errors from its inner plan, and returns an
+/// `Err` wrapping the error.
+///
+/// We usually need to apply this plan if (and only if) the output of the inner
+/// plan is of a response type.
 ///
 /// The errors come from two places: `Err` from inner plans, and `Ok(response)`
 /// where `response` contains unresolved errors (`error` and `region_error`).

src/timestamp.rs

@@ -35,6 +35,8 @@ impl TimestampExt for Timestamp {
         Self {
             physical: version >> PHYSICAL_SHIFT_BITS,
             logical: version & LOGICAL_MASK,
+            // Now we only support global transactions
+            suffix_bits: 0,
         }
     }