From b6d59ec19a3d41ce581716cf0dda5d47c2785843 Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Mon, 4 Jul 2022 14:00:02 +0200
Subject: [PATCH] Fix poll item when item didn't change

---
 src/k2v-client/bin/k2v-cli.rs | 9 ++++++++-
 src/k2v-client/lib.rs         | 8 ++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/src/k2v-client/bin/k2v-cli.rs b/src/k2v-client/bin/k2v-cli.rs
index 884e7438..925ebeb8 100644
--- a/src/k2v-client/bin/k2v-cli.rs
+++ b/src/k2v-client/bin/k2v-cli.rs
@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use k2v_client::*;
 
 use garage_util::formater::format_table;
@@ -64,6 +66,9 @@ enum Command {
 		/// Causality information
 		#[clap(short, long)]
 		causality: String,
+		/// Timeout, in seconds
+		#[clap(short, long)]
+		timeout: Option<u64>,
 		/// Output formating
 		#[clap(flatten)]
 		output_kind: ReadOutputKind,
@@ -341,10 +346,12 @@ async fn main() -> Result<(), Error> {
 			partition_key,
 			sort_key,
 			causality,
+			timeout,
 			output_kind,
 		} => {
+			let timeout = timeout.map(Duration::from_secs);
 			let res_opt = client
-				.poll_item(&partition_key, &sort_key, causality.into(), None)
+				.poll_item(&partition_key, &sort_key, causality.into(), timeout)
 				.await?;
 			if let Some(res) = res_opt {
 				output_kind.display_output(res);
diff --git a/src/k2v-client/lib.rs b/src/k2v-client/lib.rs
index 95974d7a..c2606af4 100644
--- a/src/k2v-client/lib.rs
+++ b/src/k2v-client/lib.rs
@@ -122,14 +122,14 @@ impl K2vClient {
 
 		let res = self.dispatch(req, Some(timeout + DEFAULT_TIMEOUT)).await?;
 
-		let causality = res
-			.causality_token
-			.ok_or_else(|| Error::InvalidResponse("missing causality token".into()))?;
-
 		if res.status == StatusCode::NOT_MODIFIED {
 			return Ok(None);
 		}
 
+		let causality = res
+			.causality_token
+			.ok_or_else(|| Error::InvalidResponse("missing causality token".into()))?;
+
 		if res.status == StatusCode::NO_CONTENT {
 			return Ok(Some(CausalValue {
 				causality,
-- 
2.45.2