Merge pull request #79 from nrc/ts

Some more refactoring

Author: nrc

Cargo.toml

@@ -23,7 +23,7 @@ derive-new = "0.5"
 failure = "0.1"
 futures-preview = { version = "0.3.0-alpha.15", features = ["compat"] }
 grpcio = { version = "0.5.0-alpha", features = [ "secure", "prost-codec" ], default-features = false }
-lazy_static = "0.2.1"
+lazy_static = "1"
 log = "0.3.9"
 serde = "1.0"
 serde_derive = "1.0"
@@ -44,4 +44,7 @@ tempdir = "0.3"
 runtime = "0.3.0-alpha.3"
 runtime-tokio = "0.3.0-alpha.3"
 proptest = "0.9"
-proptest-derive = "0.1.0"
+proptest-derive = "0.1.0"
+
+[replace]
+"prost:0.5.0" = { git = "https://github.com/danburkert/prost", rev = "1944c27c3029d01ff216e7b126d846b6cf8c7d77" }

rust-toolchain

@@ -0,0 +1 @@
+nightly-2019-07-09

src/compat.rs

@@ -4,8 +4,8 @@
 //! from futures 0.1 to 1.0 easier.
 
 use futures::prelude::*;
+use futures::ready;
 use futures::task::{Context, Poll};
-use futures::{ready, try_ready};
 use std::pin::Pin;
 
 /// The status of a `loop_fn` loop.
@@ -52,7 +52,7 @@ where
         loop {
             unsafe {
                 let this = Pin::get_unchecked_mut(self);
-                match try_ready!(Pin::new_unchecked(&mut this.future).poll(cx)) {
+                match ready!(Pin::new_unchecked(&mut this.future).poll(cx))? {
                     Loop::Break(x) => return Poll::Ready(Ok(x)),
                     Loop::Continue(s) => this.future = (this.func)(s),
                 }
@@ -112,7 +112,7 @@ where
     fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<T, E>> {
         unsafe {
             let this = Pin::get_unchecked_mut(self);
-            let result = try_ready!(Pin::new_unchecked(&mut this.future).poll(cx));
+            let result = ready!(Pin::new_unchecked(&mut this.future).poll(cx))?;
             Poll::Ready((this.func)(result))
         }
     }
@@ -145,13 +145,13 @@ pub(crate) trait SinkCompat<I, E> {
     fn send_all_compat<S>(self, stream: S) -> SendAllCompat<Self, S>
     where
         S: Stream<Item = I> + Unpin,
-        Self: Sink<I, SinkError = E> + Sized + Unpin,
+        Self: Sink<I, Error = E> + Sized + Unpin,
     {
         SendAllCompat::new(self, stream)
     }
 }
 
-impl<T, E, S: Sink<T, SinkError = E>> SinkCompat<T, E> for S {}
+impl<T, E, S: Sink<T, Error = E>> SinkCompat<T, E> for S {}
 
 #[derive(Debug)]
 #[must_use = "futures do nothing unless polled"]
@@ -219,7 +219,7 @@ where
         &mut self,
         item: St::Item,
         cx: &mut Context,
-    ) -> Poll<Result<(()), Si::SinkError>> {
+    ) -> Poll<Result<(()), Si::Error>> {
         debug_assert!(self.buffered.is_none());
         match self.sink_mut().poll_ready(cx) {
             Poll::Ready(Ok(())) => Poll::Ready(self.sink_mut().start_send(item)),
@@ -237,22 +237,22 @@ where
     Si: Sink<St::Item> + Unpin,
     St: Stream + Unpin,
 {
-    type Output = Result<((Si, St)), Si::SinkError>;
+    type Output = Result<((Si, St)), Si::Error>;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<((Si, St)), Si::SinkError>> {
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<((Si, St)), Si::Error>> {
         if let Some(item) = self.buffered.take() {
-            try_ready!(self.try_start_send(item, cx))
+            ready!(self.try_start_send(item, cx))?
         }
 
         loop {
             match self.stream_mut().poll_next(cx) {
-                Poll::Ready(Some(item)) => try_ready!(self.try_start_send(item, cx)),
+                Poll::Ready(Some(item)) => ready!(self.try_start_send(item, cx))?,
                 Poll::Ready(None) => {
-                    try_ready!(self.sink_mut().poll_close(cx));
+                    ready!(self.sink_mut().poll_close(cx))?;
                     return Poll::Ready(Ok(self.take_result()));
                 }
                 Poll::Pending => {
-                    try_ready!(self.sink_mut().poll_flush(cx));
+                    ready!(self.sink_mut().poll_flush(cx))?;
                     return Poll::Pending;
                 }
             }

src/rpc/mod.rs

@@ -8,4 +8,5 @@ mod pd;
 mod security;
 mod tikv;
 
-pub(crate) use crate::rpc::client::RpcClient;
+pub(crate) use self::client::RpcClient;
+pub use self::pd::Timestamp;

src/rpc/pd/client.rs

@@ -9,14 +9,14 @@ use std::{
 use futures::compat::Compat01As03;
 use futures::prelude::*;
 use grpcio::{CallOption, Environment};
-use kvproto::{metapb, pdpb, pdpb::PdClient as RpcClient};
+use kvproto::{pdpb, pdpb::PdClient as RpcClient};
 
 use crate::{
     rpc::{
         context::RequestContext,
         pd::{
-            context::request_context, leader::LeaderClient, request::Request, PdTimestamp, Region,
-            RegionId, Store, StoreId,
+            context::request_context, leader::LeaderClient, request::Request, Region, Store,
+            StoreId, Timestamp,
         },
         security::SecurityManager,
     },
@@ -74,10 +74,7 @@ impl PdClient {
         self.leader.read().unwrap().members.get_leader().clone()
     }
 
-    fn get_region_and_leader(
-        &self,
-        key: &[u8],
-    ) -> impl Future<Output = Result<(metapb::Region, Option<metapb::Peer>)>> {
+    pub fn get_region(&self, key: &[u8]) -> impl Future<Output = Result<Region>> {
         let mut req = pd_request!(self.cluster_id, pdpb::GetRegionRequest);
         req.set_region_key(key.to_owned());
         let key = req.get_region_key().to_owned();
@@ -99,14 +96,11 @@ impl PdClient {
             } else {
                 None
             };
-            future::ready(Ok((region, leader)))
+            future::ready(Ok(Region::new(region, leader)))
         })
     }
 
-    fn get_region_and_leader_by_id(
-        &self,
-        region_id: u64,
-    ) -> impl Future<Output = Result<(metapb::Region, Option<metapb::Peer>)>> {
+    pub fn get_region_by_id(&self, region_id: u64) -> impl Future<Output = Result<Region>> {
         let mut req = pd_request!(self.cluster_id, pdpb::GetRegionByIdRequest);
         req.set_region_id(region_id);
 
@@ -128,7 +122,7 @@ impl PdClient {
             } else {
                 None
             };
-            future::ready(Ok((region, leader)))
+            future::ready(Ok(Region::new(region, leader)))
         })
     }
 
@@ -147,7 +141,7 @@ impl PdClient {
             let option = CallOption::default().timeout(timeout);
             let cli = &cli.read().unwrap().client;
             executor(cli, option).unwrap().map(|r| match r {
-                Err(e) => Err(ErrorKind::Grpc(e))?,
+                Err(e) => Err(ErrorKind::Grpc(e).into()),
                 Ok(r) => {
                     {
                         let header = r.header();
@@ -195,17 +189,7 @@ impl PdClient {
         .map_ok(|mut resp| resp.take_store())
     }
 
-    pub fn get_region(&self, key: &[u8]) -> impl Future<Output = Result<Region>> {
-        self.get_region_and_leader(key)
-            .map_ok(|x| Region::new(x.0, x.1))
-    }
-
-    pub fn get_region_by_id(&self, id: RegionId) -> impl Future<Output = Result<Region>> {
-        self.get_region_and_leader_by_id(id)
-            .map_ok(|x| Region::new(x.0, x.1))
-    }
-
-    pub fn get_ts(&self) -> impl Future<Output = Result<PdTimestamp>> {
+    pub fn get_ts(&self) -> impl Future<Output = Result<Timestamp>> {
         self.leader.write().unwrap().get_ts()
     }
 }

src/rpc/pd/leader.rs

@@ -12,6 +12,7 @@ use futures::channel::{
     oneshot,
 };
 use futures::compat::{Compat01As03, Compat01As03Sink};
+use futures::future::TryFutureExt;
 use futures::prelude::*;
 use grpcio::{CallOption, Environment, WriteFlags};
 use kvproto::pdpb;
@@ -22,7 +23,7 @@ use crate::{
     rpc::{
         pd::{
             context::{observe_tso_batch, request_context},
-            PdTimestamp,
+            Timestamp,
         },
         security::SecurityManager,
     },
@@ -39,12 +40,12 @@ macro_rules! pd_request {
     }};
 }
 
-type TsoChannel = oneshot::Sender<PdTimestamp>;
+type TsoChannel = oneshot::Sender<Timestamp>;
 
 enum PdTask {
     Init,
     Request,
-    Response(Vec<oneshot::Sender<PdTimestamp>>, pdpb::TsoResponse),
+    Response(Vec<oneshot::Sender<Timestamp>>, pdpb::TsoResponse),
 }
 
 struct PdReactor {
@@ -189,7 +190,7 @@ impl PdReactor {
         let timestamp = response.get_timestamp();
         for (offset, request) in requests.drain(..).enumerate() {
             request
-                .send(PdTimestamp {
+                .send(Timestamp {
                     physical: timestamp.physical,
                     logical: timestamp.logical + offset as i64,
                 })
@@ -206,9 +207,9 @@ impl PdReactor {
         }
     }
 
-    fn get_ts(&mut self) -> impl Future<Output = Result<PdTimestamp>> {
+    fn get_ts(&mut self) -> impl Future<Output = Result<Timestamp>> {
         let context = request_context("get_ts", ());
-        let (tx, rx) = oneshot::channel::<PdTimestamp>();
+        let (tx, rx) = oneshot::channel::<Timestamp>();
         self.tso_batch.push(tx);
         if self.tso_pending.is_none() {
             // Schedule tso request to run.
@@ -260,7 +261,7 @@ impl LeaderClient {
         Ok(client)
     }
 
-    pub fn get_ts(&mut self) -> impl Future<Output = Result<PdTimestamp>> {
+    pub fn get_ts(&mut self) -> impl Future<Output = Result<Timestamp>> {
         self.reactor.get_ts()
     }
 

src/rpc/pd/mod.rs

@@ -86,14 +86,10 @@ impl Region {
             .map(Into::into)
             .ok_or_else(|| Error::stale_epoch(None))
     }
-
-    pub fn meta(&self) -> metapb::Region {
-        self.region.clone()
-    }
 }
 
 #[derive(Eq, PartialEq, Debug)]
-pub struct PdTimestamp {
+pub struct Timestamp {
     pub physical: i64,
     pub logical: i64,
 }

src/rpc/tikv/client.rs

@@ -666,7 +666,7 @@ impl KvClient {
         )
         .unwrap()
         .map(|r| match r {
-            Err(e) => Err(ErrorKind::Grpc(e))?,
+            Err(e) => Err(ErrorKind::Grpc(e).into()),
             Ok(mut r) => {
                 if let Some(e) = r.region_error() {
                     Err(e)

src/transaction/mod.rs

@@ -12,6 +12,7 @@
 pub use self::client::{Client, Connect};
 pub use self::requests::{BatchGet, Commit, Delete, Get, LockKeys, Rollback, Scanner, Set};
 pub use self::transaction::{IsolationLevel, Snapshot, Transaction, TxnInfo};
+pub use super::rpc::Timestamp;
 
 use crate::{Key, Value};
 
@@ -26,27 +27,3 @@ pub enum Mutation {
     Lock(Key),
     Rollback(Key),
 }
-
-/// A logical timestamp produced by PD.
-#[derive(Copy, Clone)]
-pub struct Timestamp(u64);
-
-impl From<u64> for Timestamp {
-    fn from(v: u64) -> Self {
-        Timestamp(v)
-    }
-}
-
-impl Timestamp {
-    pub fn timestamp(self) -> u64 {
-        self.0
-    }
-
-    pub fn physical(self) -> i64 {
-        (self.0 >> 16) as i64
-    }
-
-    pub fn logical(self) -> i64 {
-        (self.0 & 0xFFFF as u64) as i64
-    }
-}