Unify client reconnection thread in CollectorOutput

Author: Molter73

collector/lib/CollectorOutput.cpp

@@ -3,14 +3,16 @@
 #include "internalapi/sensor/collector.pb.h"
 #include "internalapi/sensor/collector_iservice.pb.h"
 
-#include "HostInfo.h"
+#include "GRPCUtil.h"
 
 namespace collector {
 
 CollectorOutput::CollectorOutput(const CollectorConfig& config) {
   if (config.grpc_channel != nullptr) {
-    auto sensor_client = std::make_unique<SensorClient>(config.grpc_channel);
-    auto signal_client = std::make_unique<SignalServiceClient>(config.grpc_channel);
+    channel_ = config.grpc_channel;
+
+    auto sensor_client = std::make_unique<SensorClient>(channel_);
+    auto signal_client = std::make_unique<SignalServiceClient>(channel_);
     sensor_clients_.emplace_back(std::move(sensor_client));
     signal_clients_.emplace_back(std::move(signal_client));
   }
@@ -26,39 +28,62 @@ CollectorOutput::CollectorOutput(const CollectorConfig& config) {
     CLOG(FATAL) << "No available output configured";
   }
 
-  StartClients(sensor_clients_);
-  StartClients(signal_clients_);
+  thread_.Start([this] { EstablishGrpcStream(); });
+}
+
+void CollectorOutput::HandleOutputError() {
+  CLOG(ERROR) << "GRPC stream interrupted";
+  stream_active_.store(false, std::memory_order_release);
+  stream_interrupted_.notify_one();
 }
 
-namespace {
-SignalHandler::Result SensorOutput(std::vector<std::unique_ptr<ISensorClient>>& clients, const sensor::MsgFromCollector& msg) {
-  for (auto& client : clients) {
+SignalHandler::Result CollectorOutput::SensorOutput(const sensor::MsgFromCollector& msg) {
+  for (auto& client : sensor_clients_) {
     auto res = client->SendMsg(msg);
-    if (res != SignalHandler::PROCESSED) {
-      return res;
+    switch (res) {
+      case SignalHandler::PROCESSED:
+        break;
+
+      case SignalHandler::ERROR:
+        HandleOutputError();
+        return res;
+
+      case SignalHandler::IGNORED:
+      case SignalHandler::NEEDS_REFRESH:
+      case SignalHandler::FINISHED:
+        return res;
     }
   }
   return SignalHandler::PROCESSED;
 }
 
-SignalHandler::Result SignalOutput(std::vector<std::unique_ptr<ISignalServiceClient>>& clients, const sensor::SignalStreamMessage& msg) {
-  for (auto& client : clients) {
+SignalHandler::Result CollectorOutput::SignalOutput(const sensor::SignalStreamMessage& msg) {
+  for (auto& client : signal_clients_) {
     auto res = client->PushSignals(msg);
-    if (res != SignalHandler::PROCESSED) {
-      return res;
+    switch (res) {
+      case SignalHandler::PROCESSED:
+        break;
+
+      case SignalHandler::ERROR:
+        HandleOutputError();
+        return res;
+
+      case SignalHandler::IGNORED:
+      case SignalHandler::NEEDS_REFRESH:
+      case SignalHandler::FINISHED:
+        return res;
     }
   }
   return SignalHandler::PROCESSED;
 }
-}  // namespace
 
 SignalHandler::Result CollectorOutput::SendMsg(const MessageType& msg) {
   auto visitor = [this](auto&& m) {
     using T = std::decay_t<decltype(m)>;
     if constexpr (std::is_same_v<T, sensor::MsgFromCollector>) {
-      return SensorOutput(sensor_clients_, m);
+      return SensorOutput(m);
     } else if constexpr (std::is_same_v<T, sensor::SignalStreamMessage>) {
-      return SignalOutput(signal_clients_, m);
+      return SignalOutput(m);
     }
 
     // Unknown type
@@ -82,4 +107,43 @@ void CollectorOutput::Register() {
     }
   }
 }
+
+void CollectorOutput::EstablishGrpcStream() {
+  while (EstablishGrpcStreamSingle()) {
+  }
+  CLOG(INFO) << "Signal service client terminating.";
+}
+
+bool CollectorOutput::EstablishGrpcStreamSingle() {
+  std::mutex mtx;
+  std::unique_lock<std::mutex> lock(mtx);
+  stream_interrupted_.wait(lock, [this]() { return !stream_active_.load(std::memory_order_acquire) || thread_.should_stop(); });
+  if (thread_.should_stop()) {
+    return false;
+  }
+
+  CLOG(INFO) << "Trying to establish GRPC stream...";
+
+  if (!WaitForChannelReady(channel_, [this]() { return thread_.should_stop(); })) {
+    return false;
+  }
+  if (thread_.should_stop()) {
+    return false;
+  }
+
+  // Refresh all clients
+  auto success = true;
+  for (const auto& client : signal_clients_) {
+    success &= client->Refresh();
+  }
+
+  for (const auto& client : sensor_clients_) {
+    success &= client->Refresh();
+  }
+
+  if (success) {
+    stream_active_.store(true, std::memory_order_release);
+  }
+  return true;
+}
 }  // namespace collector

collector/lib/CollectorOutput.h

@@ -17,42 +17,40 @@ using MessageType = std::variant<sensor::MsgFromCollector, sensor::SignalStreamM
 
 class CollectorOutput {
  public:
-  CollectorOutput(const CollectorOutput&) = default;
+  CollectorOutput(const CollectorOutput&) = delete;
   CollectorOutput(CollectorOutput&&) = delete;
   CollectorOutput& operator=(const CollectorOutput&) = delete;
   CollectorOutput& operator=(CollectorOutput&&) = delete;
 
   CollectorOutput(const CollectorConfig& config);
 
   ~CollectorOutput() {
-    StopClients(sensor_clients_);
-    StopClients(signal_clients_);
+    stream_interrupted_.notify_one();
+    thread_.Stop();
   }
 
   SignalHandler::Result SendMsg(const MessageType& msg);
   void Register();
 
-  bool UseSensorClient() { return use_sensor_client_; }
+  bool UseSensorClient() const { return use_sensor_client_; }
 
  private:
-  template <typename T>
-  void StartClients(std::vector<T>& clients) {
-    for (auto& client : clients) {
-      client->Start();
-    }
-  }
+  void EstablishGrpcStream();
+  bool EstablishGrpcStreamSingle();
 
-  template <typename T>
-  void StopClients(std::vector<T>& clients) {
-    for (auto& client : clients) {
-      client->Stop();
-    }
-  }
+  void HandleOutputError();
+  SignalHandler::Result SensorOutput(const sensor::MsgFromCollector& msg);
+  SignalHandler::Result SignalOutput(const sensor::SignalStreamMessage& msg);
 
   std::vector<std::unique_ptr<ISensorClient>> sensor_clients_;
   std::vector<std::unique_ptr<ISignalServiceClient>> signal_clients_;
 
   bool use_sensor_client_{true};
+
+  StoppableThread thread_;
+  std::atomic<bool> stream_active_{false};
+  std::condition_variable stream_interrupted_;
+  std::shared_ptr<grpc::Channel> channel_;
 };
 
 }  // namespace collector

collector/lib/CollectorService.cpp

@@ -96,6 +96,8 @@ void CollectorService::RunForever() {
   // Start monitoring services.
   config_loader_.Start();
 
+  output_.Register();
+
   CLOG(INFO) << "Network scrape interval set to " << config_.ScrapeInterval() << " seconds";
 
   if (net_status_notifier_) {
@@ -106,8 +108,6 @@ void CollectorService::RunForever() {
     CLOG(FATAL) << "Unable to start collector stats exporter";
   }
 
-  output_.Register();
-
   system_inspector_.Start();
 
   ControlValue cv;

collector/lib/SensorClient.cpp

@@ -1,58 +1,23 @@
 #include "SensorClient.h"
 
-#include "GRPCUtil.h"
 #include "Logging.h"
 
 namespace collector {
-bool SensorClient::EstablishGRPCStreamSingle() {
-  std::mutex mtx;
-  std::unique_lock<std::mutex> lock(mtx);
-  stream_interrupted_.wait(lock, [this]() { return !stream_active_.load(std::memory_order_acquire) || thread_.should_stop(); });
-  if (thread_.should_stop()) {
-    return false;
-  }
-
-  CLOG(INFO) << "Trying to establish GRPC stream for signals ...";
-
-  if (!WaitForChannelReady(channel_, [this]() { return thread_.should_stop(); })) {
-    return false;
-  }
-  if (thread_.should_stop()) {
-    return false;
-  }
-
-  // stream writer
+bool SensorClient::Refresh() {
   context_ = std::make_unique<grpc::ClientContext>();
   writer_ = DuplexClient::CreateWithReadsIgnored(&sensor::CollectorService::Stub::AsyncCommunicate, channel_, context_.get());
   if (!writer_->WaitUntilStarted(std::chrono::seconds(30))) {
     CLOG(ERROR) << "Signal stream not ready after 30 seconds. Retrying ...";
     CLOG(ERROR) << "Error message: " << writer_->FinishNow().error_message();
     writer_.reset();
-    return true;
+    return false;
   }
-  CLOG(INFO) << "Successfully established GRPC stream for signals.";
 
   first_write_ = true;
   stream_active_.store(true, std::memory_order_release);
   return true;
 }
 
-void SensorClient::EstablishGRPCStream() {
-  while (EstablishGRPCStreamSingle());
-  CLOG(INFO) << "Signal service client terminating.";
-}
-
-void SensorClient::Start() {
-  thread_.Start([this] { EstablishGRPCStream(); });
-}
-
-void SensorClient::Stop() {
-  stream_interrupted_.notify_one();
-  thread_.Stop();
-  context_->TryCancel();
-  context_.reset();
-}
-
 SignalHandler::Result SensorClient::SendMsg(const sensor::MsgFromCollector& msg) {
   if (!stream_active_.load(std::memory_order_acquire)) {
     CLOG_THROTTLED(ERROR, std::chrono::seconds(10))
@@ -65,16 +30,15 @@ SignalHandler::Result SensorClient::SendMsg(const sensor::MsgFromCollector& msg)
     return SignalHandler::NEEDS_REFRESH;
   }
 
-  if (!writer_->Write(msg)) {
+  auto res = writer_->Write(msg);
+  if (!res) {
     auto status = writer_->FinishNow();
     if (!status.ok()) {
       CLOG(ERROR) << "GRPC writes failed: (" << status.error_code() << ") " << status.error_message();
     }
     writer_.reset();
 
     stream_active_.store(false, std::memory_order_release);
-    CLOG(ERROR) << "GRPC stream interrupted";
-    stream_interrupted_.notify_one();
     return SignalHandler::ERROR;
   }
 

collector/lib/SensorClient.h

@@ -26,8 +26,7 @@ class ISensorClient {
   ISensorClient& operator=(ISensorClient&&) = delete;
   virtual ~ISensorClient() = default;
 
-  virtual void Start() = 0;
-  virtual void Stop() = 0;
+  virtual bool Refresh() = 0;
 
   virtual SignalHandler::Result SendMsg(const sensor::MsgFromCollector& msg) = 0;
 };
@@ -36,24 +35,26 @@ class SensorClient : public ISensorClient {
  public:
   using Service = sensor::CollectorService;
 
+  SensorClient(const SensorClient&) = delete;
+  SensorClient(SensorClient&&) = delete;
+  SensorClient& operator=(const SensorClient&) = delete;
+  SensorClient& operator=(SensorClient&&) = delete;
+  ~SensorClient() override {
+    context_->TryCancel();
+  }
+
   explicit SensorClient(std::shared_ptr<grpc::Channel> channel)
       : channel_(std::move(channel)) {
   }
 
-  void Start() override;
-  void Stop() override;
+  bool Refresh() override;
 
   SignalHandler::Result SendMsg(const sensor::MsgFromCollector& msg) override;
 
  private:
-  void EstablishGRPCStream();
-  bool EstablishGRPCStreamSingle();
-
   std::shared_ptr<grpc::Channel> channel_;
 
-  StoppableThread thread_;
   std::atomic<bool> stream_active_ = false;
-  std::condition_variable stream_interrupted_;
 
   // This needs to have the same lifetime as the class.
   std::unique_ptr<grpc::ClientContext> context_;
@@ -63,8 +64,7 @@ class SensorClient : public ISensorClient {
 };
 
 class SensorClientStdout : public ISensorClient {
-  void Start() override {}
-  void Stop() override {}
+  bool Refresh() override { return true; }
 
   SignalHandler::Result SendMsg(const sensor::MsgFromCollector& msg) override {
     LogProtobufMessage(msg);