Add Txn to do heartbeat, track reads and writes and commit for transa…

…ction

Resolves #34.

Cargo.toml

@@ -56,9 +56,11 @@ futures = "0.3.28"
 thiserror = "1.0.48"
 tokio-stream = "0.1.15"
 asyncs = "0.3.0"
+async-io = "2.3.4"
 
 [dev-dependencies]
 assertor = "0.0.2"
+asyncs = { version = "0.3.0", features = ["test", "tokio"] }
 env_logger = "0.10.0"
 serial_test = "2.0.0"
 speculoos = "0.11.0"

src/lib.rs

@@ -21,4 +21,6 @@ pub mod keys;
 pub mod log;
 pub mod protos;
 pub mod tablet;
+pub mod timer;
+pub mod txn;
 pub mod utils;

src/protos/mod.rs

@@ -441,6 +441,20 @@ impl BatchResponse {
             },
         }
     }
+
+    #[allow(clippy::result_large_err)]
+    pub fn into_increment(mut self) -> Result<IncrementResponse, Self> {
+        if self.responses.len() != 1 {
+            return Err(self);
+        }
+        match self.responses.remove(0) {
+            ShardResponse { response: DataResponse::Increment(increment), .. } => Ok(increment),
+            response => {
+                self.responses.push(response);
+                Err(self)
+            },
+        }
+    }
 }
 
 impl DataResponse {

src/protos/temporal.rs

@@ -38,9 +38,18 @@ impl Temporal {
             Temporal::Transaction(txn) => txn,
         }
     }
+
+    pub fn txn(&self) -> &Transaction {
+        match self {
+            Temporal::Timestamp(ts) => panic!("expect transaction, got timestamp {ts}"),
+            Temporal::Transaction(txn) => txn,
+        }
+    }
 }
 
 impl Transaction {
+    pub const HEARTBEAT_INTERVAL: Duration = Duration::from_millis(500);
+
     pub fn comparer() -> impl Fn(&Transaction, &Transaction) -> Ordering {
         let comparer = TxnMeta::comparer();
         move |a, b| comparer(&a.meta, &b.meta)
@@ -70,7 +79,7 @@ impl Transaction {
     }
 
     pub fn next_heartbeat_ts(&self) -> Timestamp {
-        self.heartbeat_ts.max(self.start_ts()).into_physical() + Duration::from_millis(500)
+        self.heartbeat_ts.max(self.start_ts()).into_physical() + Self::HEARTBEAT_INTERVAL
     }
 
     pub fn restart(&mut self) {

src/tablet/client.rs

@@ -46,6 +46,7 @@ use crate::protos::{
     TabletServiceClient,
     Temporal,
     Timestamp,
+    TimestampedValue,
     Transaction,
     TxnMeta,
     Uuid,
@@ -501,6 +502,81 @@ impl TabletClient {
         }
     }
 
+    pub async fn transactional_get(
+        &self,
+        txn: Transaction,
+        key: &[u8],
+        sequence: u32,
+    ) -> Result<(Transaction, Option<TimestampedValue>)> {
+        let (shard, mut service) = self.service(key).await?;
+        let mut response = service
+            .batch(BatchRequest {
+                tablet_id: shard.tablet_id().into(),
+                temporal: Temporal::Transaction(txn),
+                requests: vec![ShardRequest {
+                    shard_id: shard.shard_id().into(),
+                    request: DataRequest::Get(GetRequest { key: key.to_owned(), sequence }),
+                }],
+                ..Default::default()
+            })
+            .await?
+            .into_inner();
+        let txn = std::mem::take(&mut response.temporal).into_transaction();
+        let get = response.into_get().map_err(|r| anyhow!("expect get response, get {r:?}"))?;
+        Ok((txn, get.value))
+    }
+
+    pub async fn transactional_put(
+        &self,
+        txn: Transaction,
+        key: &[u8],
+        value: Option<Value>,
+        sequence: u32,
+        expect_ts: Option<Timestamp>,
+    ) -> Result<Transaction> {
+        let (shard, mut service) = self.service(key).await?;
+        let mut response = service
+            .batch(BatchRequest {
+                tablet_id: shard.tablet_id().into(),
+                temporal: Temporal::Transaction(txn),
+                requests: vec![ShardRequest {
+                    shard_id: shard.shard_id().into(),
+                    request: DataRequest::Put(PutRequest { key: key.to_owned(), value, sequence, expect_ts }),
+                }],
+                ..Default::default()
+            })
+            .await?
+            .into_inner();
+        let txn = std::mem::take(&mut response.temporal).into_transaction();
+        response.into_put().map_err(|r| anyhow!("expect put response, get {r:?}"))?;
+        Ok(txn)
+    }
+
+    pub async fn transactional_increment(
+        &self,
+        txn: Transaction,
+        key: &[u8],
+        increment: i64,
+        sequence: u32,
+    ) -> Result<(Transaction, i64)> {
+        let (shard, mut service) = self.service(key).await?;
+        let mut response = service
+            .batch(BatchRequest {
+                tablet_id: shard.tablet_id().into(),
+                temporal: Temporal::Transaction(txn),
+                requests: vec![ShardRequest {
+                    shard_id: shard.shard_id().into(),
+                    request: DataRequest::Increment(IncrementRequest { key: key.to_owned(), increment, sequence }),
+                }],
+                ..Default::default()
+            })
+            .await?
+            .into_inner();
+        let txn = std::mem::take(&mut response.temporal).into_transaction();
+        let increment = response.into_increment().map_err(|r| anyhow!("expect increment response, get {r:?}"))?;
+        Ok((txn, increment.value))
+    }
+
     pub async fn open_participate_txn(
         &self,
         request: ParticipateTxnRequest,

src/timer.rs

@@ -0,0 +1,63 @@
+// Copyright 2023 The SeamDB Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::time::{Duration, Instant};
+
+use async_io::Timer as IoTimer;
+
+enum Inner {
+    Ready(Instant),
+    Never,
+    Timer(IoTimer),
+}
+
+pub struct Timer {
+    inner: Inner,
+}
+
+impl Timer {
+    pub fn never() -> Timer {
+        Self { inner: Inner::Never }
+    }
+
+    pub fn ready() -> Timer {
+        Self { inner: Inner::Ready(Instant::now()) }
+    }
+
+    pub fn after(duration: Duration) -> Timer {
+        Self { inner: Inner::Timer(IoTimer::after(duration)) }
+    }
+
+    pub fn interval(period: Duration) -> Timer {
+        Self { inner: Inner::Timer(IoTimer::interval(period)) }
+    }
+}
+
+impl Future for Timer {
+    type Output = Instant;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match &mut self.inner {
+            Inner::Ready(instant) => Poll::Ready(*instant),
+            Inner::Never => Poll::Pending,
+            Inner::Timer(timer) => {
+                let timer = unsafe { Pin::new_unchecked(timer) };
+                timer.poll(cx)
+            },
+        }
+    }
+}