samples/tcp_full_duplex_service/tcp_full_duplex_service.cpp

#include <userver/utest/using_namespace_userver.hpp>
 
#include <userver/clients/dns/component.hpp>
#include <userver/components/component.hpp>
#include <userver/components/minimal_server_component_list.hpp>
#include <userver/components/statistics_storage.hpp>
#include <userver/components/tcp_acceptor_base.hpp>
#include <userver/concurrent/queue.hpp>
#include <userver/server/handlers/server_monitor.hpp>
#include <userver/utils/daemon_run.hpp>
 
#include <userver/utils/statistics/metric_tag.hpp>
#include <userver/utils/statistics/metrics_storage.hpp>
 
#include <userver/clients/http/component.hpp>
#include <userver/server/handlers/tests_control.hpp>
#include <userver/testsuite/testsuite_support.hpp>
 
namespace samples::tcp::echo {
 
struct Stats;
 
class Echo final : public components::TcpAcceptorBase {
 public:
  static constexpr std::string_view kName = "tcp-echo";
 
  // Component is valid after construction and is able to accept requests
  Echo(const components::ComponentConfig& config,
       const components::ComponentContext& context);
 
  void ProcessSocket(engine::io::Socket&& sock) override;
 
 private:
  Stats& stats_;
};
 
}  // namespace samples::tcp::echo
 
 
namespace samples::tcp::echo {
 
struct Stats {
  std::atomic<std::uint64_t> opened_sockets{0};
  std::atomic<std::uint64_t> closed_sockets{0};
  std::atomic<std::uint64_t> bytes_read{0};
};
 
const utils::statistics::MetricTag<Stats> kTcpEchoTag{"tcp-echo"};
 
void DumpMetric(utils::statistics::Writer& writer, const Stats& stats) {
  writer["sockets"]["opened"] = stats.opened_sockets;
  writer["sockets"]["closed"] = stats.closed_sockets;
  writer["bytes"]["read"] = stats.bytes_read;
}
 
void ResetMetric(Stats& stats) {
  stats.opened_sockets = 0;
  stats.closed_sockets = 0;
  stats.bytes_read = 0;
}
 
Echo::Echo(const components::ComponentConfig& config,
           const components::ComponentContext& context)
    : TcpAcceptorBase(config, context),
      stats_(context.FindComponent<components::StatisticsStorage>()
                 .GetMetricsStorage()
                 ->GetMetric(kTcpEchoTag)) {}
 
namespace {
 
using Queue = concurrent::SpscQueue<std::string>;
 
void DoSend(engine::io::Socket& sock, Queue::Consumer consumer) {
  std::string data;
  while (consumer.Pop(data)) {
    const auto sent_bytes = sock.SendAll(data.data(), data.size(), {});
    if (sent_bytes != data.size()) {
      LOG_INFO() << "Failed to send all the data";
      return;
    }
  }
}
 
void DoRecv(engine::io::Socket& sock, Queue::Producer producer, Stats& stats) {
  std::array<char, 1024> buf;  // NOLINT(cppcoreguidelines-pro-type-member-init)
  while (!engine::current_task::ShouldCancel()) {
    const auto read_bytes = sock.ReadSome(buf.data(), buf.size(), {});
    if (!read_bytes) {
      LOG_INFO() << "Failed to read data";
      return;
    }
 
    stats.bytes_read += read_bytes;
    if (!producer.Push({buf.data(), read_bytes})) {
      return;
    }
  }
}
 
}  // anonymous namespace
 
void Echo::ProcessSocket(engine::io::Socket&& sock) {
  const auto sock_num = ++stats_.opened_sockets;
 
  tracing::Span span{fmt::format("sock_{}", sock_num)};
  span.AddTag("fd", std::to_string(sock.Fd()));
 
  utils::FastScopeGuard guard{[this]() noexcept {
    LOG_INFO() << "Closing socket";
    ++stats_.closed_sockets;
  }};
 
  auto queue = Queue::Create();
 
  auto send_task =
      utils::Async("send", DoSend, std::ref(sock), queue->GetConsumer());
 
  DoRecv(sock, queue->GetProducer(), stats_);
}
 
}  // namespace samples::tcp::echo
 
int main(int argc, const char* const argv[]) {
  const auto component_list = components::MinimalServerComponentList()
                                  .Append<server::handlers::ServerMonitor>()
                                  .Append<samples::tcp::echo::Echo>()
                                  // Testuite components:
                                  .Append<server::handlers::TestsControl>()
                                  .Append<components::TestsuiteSupport>()
                                  .Append<clients::dns::Component>()
                                  .Append<components::HttpClient>();
 
  return utils::DaemonMain(argc, argv, component_list);
}