ADD: Add configurable heartbeat support to clients

threecgreen · threecgreen · commit 0b589ed51ccc · 2024-05-31T15:00:24.000Z
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,10 +3,16 @@
 ## 0.19.0 - TBD
 
 ### Enhancements
+- Added configurable `heartbeat_interval` parameter for live clients that determines the
+  timeout before heartbeat `SystemMsg` records will be sent. It can be configured via
+  the `SetHeartbeatInterval` method of the `LiveBuilder`
+- Added `SetAddress` method to `LiveBuilder` for configuring a custom gateway address
+  without using the constructor directly
 - Added new `UncrossingPrice` `StatType` variant
 - Added new publisher values for `XNAS.BASIC`
 
 ### Breaking changes
+- Added `heartbeat_interval` parameter to the `Live` constructors
 - Removed `start_date` and `end_date` fields from `DatasetRange` struct
   in favor of `start` and `end`
 - Removed live `Subscribe` method overloads with `use_snapshot`
diff --git a/include/databento/live.hpp b/include/databento/live.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <chrono>
 #include <string>
 
 #include "databento/enums.hpp"  // VersionUpgradePolicy
@@ -29,6 +30,10 @@ class LiveBuilder {
   LiveBuilder& SetUpgradePolicy(VersionUpgradePolicy upgrade_policy);
   // Sets the receiver of the logs to be used by the client.
   LiveBuilder& SetLogReceiver(ILogReceiver* log_receiver);
+  // Overrides the heartbeat interval.
+  LiveBuilder& SetHeartbeatInterval(std::chrono::seconds heartbeat_interval);
+  // Overrides the gateway and port. This is an advanced method.
+  LiveBuilder& SetAddress(std::string gateway, std::uint16_t port);
   // Attempts to construct an instance of a blocking live client or throws an
   // exception.
   LiveBlocking BuildBlocking();
@@ -40,9 +45,12 @@ class LiveBuilder {
   void Validate();
 
   ILogReceiver* log_receiver_{};
+  std::string gateway_{};
+  std::uint16_t port_{};
   std::string key_;
   std::string dataset_;
   bool send_ts_out_{false};
   VersionUpgradePolicy upgrade_policy_{VersionUpgradePolicy::Upgrade};
+  std::chrono::seconds heartbeat_interval_{};
 };
 }  // namespace databento
diff --git a/include/databento/live_blocking.hpp b/include/databento/live_blocking.hpp
@@ -4,6 +4,7 @@
 #include <chrono>  // milliseconds
 #include <cstdint>
 #include <string>
+#include <utility>  // pair
 #include <vector>
 
 #include "databento/datetime.hpp"           // UnixNanos
@@ -22,10 +23,12 @@ class ILogReceiver;
 class LiveBlocking {
  public:
   LiveBlocking(ILogReceiver* log_receiver, std::string key, std::string dataset,
-               bool send_ts_out, VersionUpgradePolicy upgrade_policy);
+               bool send_ts_out, VersionUpgradePolicy upgrade_policy,
+               std::chrono::seconds heartbeat_interval);
   LiveBlocking(ILogReceiver* log_receiver, std::string key, std::string dataset,
                std::string gateway, std::uint16_t port, bool send_ts_out,
-               VersionUpgradePolicy upgrade_policy);
+               VersionUpgradePolicy upgrade_policy,
+               std::chrono::seconds heartbeat_interval);
   /*
    * Getters
    */
@@ -36,6 +39,11 @@ class LiveBlocking {
   std::uint16_t Port() const { return port_; }
   bool SendTsOut() const { return send_ts_out_; }
   VersionUpgradePolicy UpgradePolicy() const { return upgrade_policy_; }
+  // The the first member of the pair will be true, when the heartbeat interval
+  // was overridden.
+  std::pair<bool, std::chrono::seconds> HeartbeatInterval() const {
+    return {heartbeat_interval_.count() > 0, heartbeat_interval_};
+  }
 
   /*
    * Methods
@@ -95,6 +103,7 @@ class LiveBlocking {
   bool send_ts_out_;
   std::uint8_t version_{};
   VersionUpgradePolicy upgrade_policy_;
+  std::chrono::seconds heartbeat_interval_;
   detail::TcpClient client_;
   // Must be 8-byte aligned for records
   alignas(RecordHeader) std::array<char, kMaxStrLen> read_buffer_{};
diff --git a/include/databento/live_threaded.hpp b/include/databento/live_threaded.hpp
@@ -4,6 +4,7 @@
 #include <functional>  // function
 #include <memory>      // unique_ptr
 #include <string>
+#include <utility>  // pair
 #include <vector>
 
 #include "databento/datetime.hpp"              // UnixNanos
@@ -31,10 +32,12 @@ class LiveThreaded {
       std::function<ExceptionAction(const std::exception&)>;
 
   LiveThreaded(ILogReceiver* log_receiver, std::string key, std::string dataset,
-               bool send_ts_out, VersionUpgradePolicy upgrade_policy);
+               bool send_ts_out, VersionUpgradePolicy upgrade_policy,
+               std::chrono::seconds heartbeat_interval);
   LiveThreaded(ILogReceiver* log_receiver, std::string key, std::string dataset,
                std::string gateway, std::uint16_t port, bool send_ts_out,
-               VersionUpgradePolicy upgrade_policy);
+               VersionUpgradePolicy upgrade_policy,
+               std::chrono::seconds heartbeat_interval);
   LiveThreaded(const LiveThreaded&) = delete;
   LiveThreaded& operator=(const LiveThreaded&) = delete;
   LiveThreaded(LiveThreaded&& other) noexcept;
@@ -51,6 +54,9 @@ class LiveThreaded {
   std::uint16_t Port() const;
   bool SendTsOut() const;
   VersionUpgradePolicy UpgradePolicy() const;
+  // The the first member of the pair will be true, when the heartbeat interval
+  // was overridden.
+  std::pair<bool, std::chrono::seconds> HeartbeatInterval() const;
 
   /*
    * Methods
diff --git a/src/live.cpp b/src/live.cpp
@@ -1,5 +1,6 @@
 #include "databento/live.hpp"
 
+#include <chrono>
 #include <utility>  // move
 
 #include "databento/constants.hpp"      // kApiKeyLength
@@ -51,16 +52,40 @@ LiveBuilder& LiveBuilder::SetLogReceiver(
   return *this;
 }
 
+LiveBuilder& LiveBuilder::SetHeartbeatInterval(
+    std::chrono::seconds heartbeat_interval) {
+  heartbeat_interval_ = heartbeat_interval;
+  return *this;
+}
+
+LiveBuilder& LiveBuilder::SetAddress(std::string gateway, std::uint16_t port) {
+  gateway_ = std::move(gateway);
+  port_ = port;
+  return *this;
+}
+
 databento::LiveBlocking LiveBuilder::BuildBlocking() {
   Validate();
-  return databento::LiveBlocking{log_receiver_, key_, dataset_, send_ts_out_,
-                                 upgrade_policy_};
+  if (gateway_.empty()) {
+    return databento::LiveBlocking{log_receiver_,   key_,
+                                   dataset_,        send_ts_out_,
+                                   upgrade_policy_, heartbeat_interval_};
+  }
+  return databento::LiveBlocking{
+      log_receiver_, key_,         dataset_,        gateway_,
+      port_,         send_ts_out_, upgrade_policy_, heartbeat_interval_};
 }
 
 databento::LiveThreaded LiveBuilder::BuildThreaded() {
   Validate();
-  return databento::LiveThreaded{log_receiver_, key_, dataset_, send_ts_out_,
-                                 upgrade_policy_};
+  if (gateway_.empty()) {
+    return databento::LiveThreaded{log_receiver_,   key_,
+                                   dataset_,        send_ts_out_,
+                                   upgrade_policy_, heartbeat_interval_};
+  }
+  return databento::LiveThreaded{
+      log_receiver_, key_,         dataset_,        gateway_,
+      port_,         send_ts_out_, upgrade_policy_, heartbeat_interval_};
 }
 
 void LiveBuilder::Validate() {
diff --git a/src/live_blocking.cpp b/src/live_blocking.cpp
@@ -4,7 +4,8 @@
 
 #include <algorithm>  // copy
 #include <cctype>     // tolower
-#include <cstddef>    // ptrdiff_t
+#include <chrono>
+#include <cstddef>  // ptrdiff_t
 #include <cstdlib>
 #include <ios>  //hex, setfill, setw
 #include <sstream>
@@ -26,7 +27,8 @@ constexpr std::size_t kBucketIdLength = 5;
 
 LiveBlocking::LiveBlocking(ILogReceiver* log_receiver, std::string key,
                            std::string dataset, bool send_ts_out,
-                           VersionUpgradePolicy upgrade_policy)
+                           VersionUpgradePolicy upgrade_policy,
+                           std::chrono::seconds heartbeat_interval)
 
     : log_receiver_{log_receiver},
       key_{std::move(key)},
@@ -35,20 +37,23 @@ LiveBlocking::LiveBlocking(ILogReceiver* log_receiver, std::string key,
       port_{13000},
       send_ts_out_{send_ts_out},
       upgrade_policy_{upgrade_policy},
+      heartbeat_interval_{heartbeat_interval},
       client_{gateway_, port_},
       session_id_{this->Authenticate()} {}
 
 LiveBlocking::LiveBlocking(ILogReceiver* log_receiver, std::string key,
                            std::string dataset, std::string gateway,
                            std::uint16_t port, bool send_ts_out,
-                           VersionUpgradePolicy upgrade_policy)
+                           VersionUpgradePolicy upgrade_policy,
+                           std::chrono::seconds heartbeat_interval)
     : log_receiver_{log_receiver},
       key_{std::move(key)},
       dataset_{std::move(dataset)},
       gateway_{std::move(gateway)},
       port_{port},
       send_ts_out_{send_ts_out},
       upgrade_policy_{upgrade_policy},
+      heartbeat_interval_{heartbeat_interval},
       client_{gateway_, port_},
       session_id_{this->Authenticate()} {}
 
@@ -261,11 +266,14 @@ std::string LiveBlocking::GenerateCramReply(const std::string& challenge_key) {
 }
 
 std::string LiveBlocking::EncodeAuthReq(const std::string& auth) {
-  std::ostringstream reply_stream;
-  reply_stream << "auth=" << auth << "|dataset=" << dataset_ << "|encoding=dbn|"
-               << "ts_out=" << send_ts_out_
-               << "|client=C++ " DATABENTO_VERSION "\n";
-  return reply_stream.str();
+  std::ostringstream req_stream;
+  req_stream << "auth=" << auth << "|dataset=" << dataset_ << "|encoding=dbn|"
+             << "ts_out=" << send_ts_out_ << "|client=C++ " DATABENTO_VERSION;
+  if (heartbeat_interval_.count() > 0) {
+    req_stream << "|heartbeat_interval_s=" << heartbeat_interval_.count();
+  }
+  req_stream << '\n';
+  return req_stream.str();
 }
 
 std::uint64_t LiveBlocking::DecodeAuthResp() {
diff --git a/src/live_threaded.cpp b/src/live_threaded.cpp
@@ -57,16 +57,19 @@ LiveThreaded::~LiveThreaded() {
 
 LiveThreaded::LiveThreaded(ILogReceiver* log_receiver, std::string key,
                            std::string dataset, bool send_ts_out,
-                           VersionUpgradePolicy upgrade_policy)
+                           VersionUpgradePolicy upgrade_policy,
+                           std::chrono::seconds heartbeat_interval)
     : impl_{new Impl{log_receiver, std::move(key), std::move(dataset),
-                     send_ts_out, upgrade_policy}} {}
+                     send_ts_out, upgrade_policy, heartbeat_interval}} {}
 
 LiveThreaded::LiveThreaded(ILogReceiver* log_receiver, std::string key,
                            std::string dataset, std::string gateway,
                            std::uint16_t port, bool send_ts_out,
-                           VersionUpgradePolicy upgrade_policy)
+                           VersionUpgradePolicy upgrade_policy,
+                           std::chrono::seconds heartbeat_interval)
     : impl_{new Impl{log_receiver, std::move(key), std::move(dataset),
-                     std::move(gateway), port, send_ts_out, upgrade_policy}} {}
+                     std::move(gateway), port, send_ts_out, upgrade_policy,
+                     heartbeat_interval}} {}
 
 const std::string& LiveThreaded::Key() const { return impl_->blocking.Key(); }
 
@@ -86,6 +89,10 @@ databento::VersionUpgradePolicy LiveThreaded::UpgradePolicy() const {
   return impl_->blocking.UpgradePolicy();
 }
 
+std::pair<bool, std::chrono::seconds> LiveThreaded::HeartbeatInterval() const {
+  return impl_->blocking.HeartbeatInterval();
+}
+
 void LiveThreaded::Subscribe(const std::vector<std::string>& symbols,
                              Schema schema, SType stype_in) {
   impl_->blocking.Subscribe(symbols, schema, stype_in);
diff --git a/test/include/mock/mock_lsg_server.hpp b/test/include/mock/mock_lsg_server.hpp
@@ -10,6 +10,7 @@ using ssize_t = SSIZE_T;
 #include <sys/socket.h>  // send
 #endif
 
+#include <chrono>
 #include <condition_variable>
 #include <functional>  // function
 #include <mutex>
@@ -28,6 +29,9 @@ class MockLsgServer {
  public:
   MockLsgServer(std::string dataset, bool ts_out,
                 std::function<void(MockLsgServer&)> serve_fn);
+  MockLsgServer(std::string dataset, bool ts_out,
+                std::chrono::seconds heartbeat_interval,
+                std::function<void(MockLsgServer&)> serve_fn);
 
   std::uint16_t Port() const { return port_; }
 
@@ -83,6 +87,7 @@ class MockLsgServer {
 
   std::string dataset_;
   bool ts_out_;
+  std::chrono::seconds heartbeat_interval_;
   std::uint16_t port_{};
   detail::ScopedFd socket_{};
   detail::ScopedFd conn_fd_{};
diff --git a/test/src/live_blocking_tests.cpp b/test/src/live_blocking_tests.cpp
diff --git a/test/src/live_threaded_tests.cpp b/test/src/live_threaded_tests.cpp
diff --git a/test/src/mock_lsg_server.cpp b/test/src/mock_lsg_server.cpp
