tlbmc/feed_client_grpc.cc

#include "tlbmc/feed_client_grpc.h"

#include <climits>
#include <memory>

#include "tlbmc/feed_client.pb.h"
#include "agent_config_service.grpc.pb.h"
#include "data_sink_service.grpc.pb.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "grpc/grpc.h"
#include "grpcpp/client_context.h"
#include "grpcpp/security/credentials.h"
#include "grpcpp/support/channel_arguments.h"
#include "grpcpp/support/status.h"

namespace platforms_syshealth::collection::feed {

void FeedClientGrpc::InitStubs() {
  ::grpc::ChannelArguments args;
  args.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS, INT_MAX);
  args.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS, 20 * 1000 /*20 sec*/);
  args.SetInt(GRPC_ARG_KEEPALIVE_PERMIT_WITHOUT_CALLS, 0);
  std::shared_ptr<::grpc::Channel> channel =
      ::grpc::CreateCustomChannel(config_.target_address(), credentials_, args);
  stub_ = DataSinkService::NewStub(channel);
  agent_config_stub_ = AgentConfigService::NewStub(channel);
  dry_run_ = false;
}

FeedClientGrpc::FeedClientGrpc(
    const FeedClientConfig& config,
    const std::shared_ptr<::grpc::ChannelCredentials>& credentials)
    : config_(config), credentials_(credentials) {
  if (config_.target_address().empty()) {
    dry_run_ = true;
    LOG(WARNING) << "FeedClientGrpc created in dry run mode because "
                    "target_address is not set.";
  } else {
    InitStubs();
  }
}

void FeedClientGrpc::SetTargetAddress(absl::string_view target_address) {
  if (!dry_run_ && target_address == config_.target_address()) {
    return;
  }
  config_.set_target_address(target_address);
  InitStubs();
}

absl::StatusOr<WriteMetricsResponse> FeedClientGrpc::WriteMetrics(
    const WriteMetricsRequest& request) {
  if (dry_run_) {
    LOG(WARNING) << "[Dry Run] Would send WriteMetricsRequest: " << request;
    return WriteMetricsResponse();
  }
  if (!stub_) {
    return absl::InternalError("gRPC stub not initialized");
  }
  ::grpc::ClientContext context;
  context.set_deadline(absl::ToChronoTime(
      absl::Now() + absl::Seconds(config_.write_metrics_timeout_seconds())));
  WriteMetricsResponse response;
  ::grpc::Status status = stub_->WriteMetrics(&context, request, &response);
  if (!status.ok()) {
    return absl::InternalError(status.error_message());
  }
  return response;
}

absl::StatusOr<GetPoliciesResponse> FeedClientGrpc::GetPolicies(
    const GetPoliciesRequest& request) {
  if (dry_run_) {
    LOG(WARNING) << "[Dry Run] Would send GetPoliciesRequest: " << request;
    return GetPoliciesResponse();
  }
  if (!agent_config_stub_) {
    return absl::InternalError("gRPC config stub not initialized");
  }
  ::grpc::ClientContext context;
  context.set_deadline(absl::ToChronoTime(
      absl::Now() + absl::Seconds(config_.get_policies_timeout_seconds())));
  GetPoliciesResponse response;
  ::grpc::Status status =
      agent_config_stub_->GetPolicies(&context, request, &response);
  if (!status.ok()) {
    return absl::InternalError(status.error_message());
  }
  return response;
}

}  // namespace platforms_syshealth::collection::feed