Basic implementation using the new collector internal service

Author: Molter73

collector/collector.cpp

@@ -12,6 +12,7 @@
 #include <sys/wait.h>
 
 #include "ConfigLoader.h"
+#include "SensorClient.h"
 
 extern "C" {
 #include <assert.h>

collector/lib/CollectorService.cpp

@@ -119,7 +119,13 @@ void CollectorService::RunForever() {
     server.addHandler(collector_config_inspector->kBaseRoute, collector_config_inspector.get());
   }
 
-  system_inspector_.Init(config_, conn_tracker);
+  if (config_.grpc_channel != nullptr) {
+    client_ = std::make_unique<SensorClient>(config_.grpc_channel);
+  } else {
+    client_ = std::make_unique<SensorClientStdout>();
+  }
+
+  system_inspector_.Init(config_, conn_tracker, client_.get());
   system_inspector_.Start();
 
   ControlValue cv;

collector/lib/CollectorService.h

@@ -24,6 +24,7 @@ class CollectorService {
   const std::atomic<int>& signum_;
 
   system_inspector::Service system_inspector_;
+  std::unique_ptr<ISensorClient> client_;
 };
 
 bool SetupKernelDriver(CollectorService& collector, const CollectorConfig& config);

collector/lib/ProcessSignalFormatter.cpp

@@ -4,6 +4,7 @@
 
 #include <google/protobuf/util/time_util.h>
 
+#include "internalapi/sensor/collector_iservice.pb.h"
 #include "internalapi/sensor/signal_iservice.pb.h"
 
 #include "CollectorStats.h"
@@ -65,9 +66,9 @@ ProcessSignalFormatter::ProcessSignalFormatter(
   event_extractor_->Init(inspector);
 }
 
-ProcessSignalFormatter::~ProcessSignalFormatter() {}
+ProcessSignalFormatter::~ProcessSignalFormatter() = default;
 
-const SignalStreamMessage* ProcessSignalFormatter::ToProtoMessage(sinsp_evt* event) {
+const sensor::MsgFromCollector* ProcessSignalFormatter::ToProtoMessage(sinsp_evt* event) {
   if (process_signals[event->get_type()] == ProcessSignalType::UNKNOWN_PROCESS_TYPE) {
     return nullptr;
   }
@@ -80,38 +81,34 @@ const SignalStreamMessage* ProcessSignalFormatter::ToProtoMessage(sinsp_evt* eve
   }
 
   ProcessSignal* process_signal = CreateProcessSignal(event);
-  if (!process_signal) {
+  if (process_signal == nullptr) {
     return nullptr;
   }
 
-  Signal* signal = Allocate<Signal>();
-  signal->set_allocated_process_signal(process_signal);
-
-  SignalStreamMessage* signal_stream_message = AllocateRoot();
-  signal_stream_message->clear_collector_register_request();
-  signal_stream_message->set_allocated_signal(signal);
-  return signal_stream_message;
+  auto* msg = AllocateRoot();
+  msg->clear_info();
+  msg->clear_register_();
+  msg->set_allocated_process_signal(process_signal);
+  return msg;
 }
 
-const SignalStreamMessage* ProcessSignalFormatter::ToProtoMessage(sinsp_threadinfo* tinfo) {
+const sensor::MsgFromCollector* ProcessSignalFormatter::ToProtoMessage(sinsp_threadinfo* tinfo) {
   Reset();
   if (!ValidateProcessDetails(tinfo)) {
     CLOG(INFO) << "Dropping process event: " << tinfo;
     return nullptr;
   }
 
-  ProcessSignal* process_signal = CreateProcessSignal(tinfo);
-  if (!process_signal) {
+  ProcessSignal* signal = CreateProcessSignal(tinfo);
+  if (signal == nullptr) {
     return nullptr;
   }
 
-  Signal* signal = Allocate<Signal>();
-  signal->set_allocated_process_signal(process_signal);
-
-  SignalStreamMessage* signal_stream_message = AllocateRoot();
-  signal_stream_message->clear_collector_register_request();
-  signal_stream_message->set_allocated_signal(signal);
-  return signal_stream_message;
+  auto* msg = AllocateRoot();
+  msg->clear_register_();
+  msg->clear_info();
+  msg->set_allocated_process_signal(signal);
+  return msg;
 }
 
 ProcessSignal* ProcessSignalFormatter::CreateProcessSignal(sinsp_evt* event) {
@@ -173,7 +170,7 @@ ProcessSignal* ProcessSignalFormatter::CreateProcessSignal(sinsp_evt* event) {
   // set time
   auto timestamp = Allocate<Timestamp>();
   *timestamp = TimeUtil::NanosecondsToTimestamp(event->get_ts());
-  signal->set_allocated_time(timestamp);
+  signal->set_allocated_creation_time(timestamp);
 
   // set container_id
   if (const std::string* container_id = event_extractor_->get_container_id(event)) {
@@ -238,7 +235,7 @@ ProcessSignal* ProcessSignalFormatter::CreateProcessSignal(sinsp_threadinfo* tin
   // set time
   auto timestamp = Allocate<Timestamp>();
   *timestamp = TimeUtil::NanosecondsToTimestamp(tinfo->m_clone_ts);
-  signal->set_allocated_time(timestamp);
+  signal->set_allocated_creation_time(timestamp);
 
   // set container_id
   signal->set_container_id(tinfo->m_container_id);
@@ -315,20 +312,20 @@ void ProcessSignalFormatter::CountLineage(const std::vector<LineageInfo>& lineag
 
 void ProcessSignalFormatter::GetProcessLineage(sinsp_threadinfo* tinfo,
                                                std::vector<LineageInfo>& lineage) {
-  if (tinfo == NULL) {
+  if (tinfo == nullptr) {
     return;
   }
-  sinsp_threadinfo* mt = NULL;
+  sinsp_threadinfo* mt = nullptr;
   if (tinfo->is_main_thread()) {
     mt = tinfo;
   } else {
     mt = tinfo->get_main_thread();
-    if (mt == NULL) {
+    if (mt == nullptr) {
       return;
     }
   }
-  sinsp_threadinfo::visitor_func_t visitor = [this, &lineage](sinsp_threadinfo* pt) {
-    if (pt == NULL) {
+  sinsp_threadinfo::visitor_func_t visitor = [&lineage](sinsp_threadinfo* pt) {
+    if (pt == nullptr) {
       return false;
     }
     if (pt->m_pid == 0) {

collector/lib/ProcessSignalFormatter.h

@@ -6,37 +6,35 @@
 #include <gtest/gtest_prod.h>
 
 #include "api/v1/signal.pb.h"
-#include "internalapi/sensor/signal_iservice.pb.h"
-#include "storage/process_indicator.pb.h"
+#include "internalapi/sensor/collector_iservice.pb.h"
 
 #include "CollectorConfig.h"
-#include "CollectorStats.h"
 #include "ContainerMetadata.h"
 #include "EventNames.h"
 #include "ProtoSignalFormatter.h"
 
 // forward definitions
 class sinsp;
 class sinsp_threadinfo;
-namespace collector {
-namespace system_inspector {
+
+namespace collector::system_inspector {
 class EventExtractor;
 }
-}  // namespace collector
 
 namespace collector {
 
-class ProcessSignalFormatter : public ProtoSignalFormatter<sensor::SignalStreamMessage> {
+class ProcessSignalFormatter : public ProtoSignalFormatter<sensor::MsgFromCollector> {
  public:
   ProcessSignalFormatter(sinsp* inspector, const CollectorConfig& config);
   ~ProcessSignalFormatter();
 
   using Signal = v1::Signal;
-  using ProcessSignal = storage::ProcessSignal;
-  using LineageInfo = storage::ProcessSignal_LineageInfo;
+  using ProcessSignal = sensor::ProcessSignal;
+  using LineageInfo = sensor::ProcessSignal_LineageInfo;
+  using MsgFromCollector = sensor::MsgFromCollector;
 
-  const sensor::SignalStreamMessage* ToProtoMessage(sinsp_evt* event) override;
-  const sensor::SignalStreamMessage* ToProtoMessage(sinsp_threadinfo* tinfo);
+  const MsgFromCollector* ToProtoMessage(sinsp_evt* event) override;
+  const MsgFromCollector* ToProtoMessage(sinsp_threadinfo* tinfo);
 
   void GetProcessLineage(sinsp_threadinfo* tinfo, std::vector<LineageInfo>& lineage);
 

collector/lib/ProcessSignalHandler.cpp

@@ -13,7 +13,7 @@
 
 namespace collector {
 
-std::string compute_process_key(const ::storage::ProcessSignal& s) {
+std::string compute_process_key(const ::sensor::ProcessSignal& s) {
   std::stringstream ss;
   ss << s.container_id() << " " << s.name() << " ";
   if (s.args().length() <= 256) {
@@ -39,21 +39,21 @@ bool ProcessSignalHandler::Stop() {
 SignalHandler::Result ProcessSignalHandler::HandleSignal(sinsp_evt* evt) {
   const auto* signal_msg = formatter_.ToProtoMessage(evt);
 
-  if (!signal_msg) {
+  if (signal_msg == nullptr) {
     ++(stats_->nProcessResolutionFailuresByEvt);
     return IGNORED;
   }
 
-  const char* name = signal_msg->signal().process_signal().name().c_str();
-  const int pid = signal_msg->signal().process_signal().pid();
+  const char* name = signal_msg->process_signal().name().c_str();
+  const uint32_t pid = signal_msg->process_signal().pid();
   DTRACE_PROBE2(collector, process_signal_handler, name, pid);
 
-  if (!rate_limiter_.Allow(compute_process_key(signal_msg->signal().process_signal()))) {
+  if (!rate_limiter_.Allow(compute_process_key(signal_msg->process_signal()))) {
     ++(stats_->nProcessRateLimitCount);
     return IGNORED;
   }
 
-  auto result = client_->PushSignals(*signal_msg);
+  auto result = client_->SendMsg(*signal_msg);
   if (result == SignalHandler::PROCESSED) {
     ++(stats_->nProcessSent);
   } else if (result == SignalHandler::ERROR) {
@@ -65,17 +65,17 @@ SignalHandler::Result ProcessSignalHandler::HandleSignal(sinsp_evt* evt) {
 
 SignalHandler::Result ProcessSignalHandler::HandleExistingProcess(sinsp_threadinfo* tinfo) {
   const auto* signal_msg = formatter_.ToProtoMessage(tinfo);
-  if (!signal_msg) {
+  if (signal_msg == nullptr) {
     ++(stats_->nProcessResolutionFailuresByTinfo);
     return IGNORED;
   }
 
-  if (!rate_limiter_.Allow(compute_process_key(signal_msg->signal().process_signal()))) {
+  if (!rate_limiter_.Allow(compute_process_key(signal_msg->process_signal()))) {
     ++(stats_->nProcessRateLimitCount);
     return IGNORED;
   }
 
-  auto result = client_->PushSignals(*signal_msg);
+  auto result = client_->SendMsg(*signal_msg);
   if (result == SignalHandler::PROCESSED) {
     ++(stats_->nProcessSent);
   } else if (result == SignalHandler::ERROR) {

collector/lib/ProcessSignalHandler.h

@@ -8,6 +8,7 @@
 #include "CollectorConfig.h"
 #include "ProcessSignalFormatter.h"
 #include "RateLimit.h"
+#include "SensorClient.h"
 #include "SignalHandler.h"
 #include "system-inspector/Service.h"
 
@@ -22,7 +23,7 @@ class ProcessSignalHandler : public SignalHandler {
  public:
   ProcessSignalHandler(
       sinsp* inspector,
-      ISignalServiceClient* client,
+      ISensorClient* client,
       system_inspector::Stats* stats,
       const CollectorConfig& config)
       : client_(client),
@@ -38,7 +39,7 @@ class ProcessSignalHandler : public SignalHandler {
   std::vector<std::string> GetRelevantEvents() override;
 
  private:
-  ISignalServiceClient* client_;
+  ISensorClient* client_;
   ProcessSignalFormatter formatter_;
   system_inspector::Stats* stats_;
   RateLimitCache rate_limiter_;

collector/lib/SensorClient.cpp

@@ -0,0 +1,83 @@
+#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
+  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;
+  }
+  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))
+        << "GRPC stream is not established";
+    return SignalHandler::ERROR;
+  }
+
+  if (first_write_) {
+    first_write_ = false;
+    return SignalHandler::NEEDS_REFRESH;
+  }
+
+  if (!writer_->Write(msg)) {
+    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;
+  }
+
+  return SignalHandler::PROCESSED;
+}
+}  // namespace collector