From 6a7bfe78a6455ec506c5c862921d4b87492c5ce1 Mon Sep 17 00:00:00 2001
From: Shachar Langbeheim <shachlan@amazon.com>
Date: Sun, 6 Aug 2023 13:42:04 +0000
Subject: [PATCH] Put CMD heartbeat behind a feature flag.

We're not sure we want this for beta, since there are problems with
blocking commands.
---
 babushka-core/src/client/client_cmd.rs | 7 ++++++-
 babushka-core/tests/test_client_cmd.rs | 1 +
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/babushka-core/src/client/client_cmd.rs b/babushka-core/src/client/client_cmd.rs
index 8af9e1e8a4..00b6fbeee8 100644
--- a/babushka-core/src/client/client_cmd.rs
+++ b/babushka-core/src/client/client_cmd.rs
@@ -3,10 +3,13 @@ use super::reconnecting_connection::ReconnectingConnection;
 use crate::connection_request::{AddressInfo, ConnectionRequest, TlsMode};
 use crate::retry_strategies::RetryStrategy;
 use futures::{stream, StreamExt};
-use logger_core::{log_debug, log_trace, log_warn};
+#[cfg(cmd_heartbeat)]
+use logger_core::log_debug;
+use logger_core::{log_trace, log_warn};
 use protobuf::EnumOrUnknown;
 use redis::{RedisError, RedisResult, Value};
 use std::sync::Arc;
+#[cfg(cmd_heartbeat)]
 use tokio::task;
 
 enum ReadFromReplicaStrategy {
@@ -119,6 +122,7 @@ impl ClientCMD {
         let read_from_replica_strategy =
             get_read_from_replica_strategy(&connection_request.read_from_replica_strategy);
 
+        #[cfg(cmd_heartbeat)]
         for node in nodes.iter() {
             Self::start_heartbeat(node.clone());
         }
@@ -207,6 +211,7 @@ impl ClientCMD {
         }
     }
 
+    #[cfg(cmd_heartbeat)]
     fn start_heartbeat(reconnecting_connection: ReconnectingConnection) {
         task::spawn(async move {
             loop {
diff --git a/babushka-core/tests/test_client_cmd.rs b/babushka-core/tests/test_client_cmd.rs
index fc0011fd35..b6b2798090 100644
--- a/babushka-core/tests/test_client_cmd.rs
+++ b/babushka-core/tests/test_client_cmd.rs
@@ -60,6 +60,7 @@ mod client_cmd_tests {
 
     #[rstest]
     #[timeout(LONG_CMD_TEST_TIMEOUT)]
+    #[cfg(cmd_heartbeat)]
     fn test_detect_disconnect_and_reconnect_using_heartbeat(#[values(false, true)] use_tls: bool) {
         let (sender, receiver) = tokio::sync::oneshot::channel();
         block_on_all(async move {