sha256-client.cc

#include <unistd.h>
#include <sys/wait.h>
#include <errno.h>

#include <iostream>
#include <array>
#include <random>
#include <algorithm>
#include <memory>

#include <boost/asio.hpp>
#include <gflags/gflags.h>

namespace asio = boost::asio;
using asio::ip::tcp;

namespace {

DEFINE_string(host, "localhost", "Server hostname");
DEFINE_string(port, "9090", "Server port number");
DEFINE_uint64(proc, 1, "Number of concurrent clients");
DEFINE_uint64(req, 100000, "Number of requests (per client)");

void run(tcp::socket& sock, uint64_t req, boost::system::error_code& error) {
  std::random_device rd;
  std::mt19937 rng(rd());
  std::uniform_int_distribution<unsigned char> dist('a', 'z');

  std::array<unsigned char, 512> data;
  for (size_t i = 0; i < req; i++) {
    std::generate(begin(data), end(data), [&] { return dist(rng); });

    // std::cerr.write(reinterpret_cast<char*>(data.data()), data.size());
    // std::cerr << std::endl;

    asio::write(sock, asio::buffer(data), error);
    if (error)
      return;

    std::array<char, 64> sha256;
    size_t n = 0;
    while (n < sha256.size()) {
      auto len
          = sock.read_some(asio::buffer(sha256.data() + n, sha256.size() - n), error);
      if (error)
        return;
      n += len;
    }
    // std::cout.write(sha256.data(), sha256.size());
    // std::cout << std::endl;
  }
}

void client(const std::string& host, const std::string& port, uint64_t req) {
  try {
    asio::io_service service;

    tcp::resolver resolver(service);
    auto flags = tcp::resolver::query::numeric_service | tcp::resolver::query::passive
                 | tcp::resolver::query::address_configured;
    tcp::resolver::query query(host, port, flags);
    tcp::resolver::iterator endpoint_iterator = resolver.resolve(query);

    tcp::socket sock(service);
    boost::asio::connect(sock, endpoint_iterator);

    boost::system::error_code error;
    run(sock, req, error);
    if (error)
      throw boost::system::system_error(error);

  } catch (std::exception& e) {
    std::cerr << e.what() << std::endl;
  }
  exit(0);
}

}  // end namespace

int main(int argc, char* argv[]) {
  gflags::SetUsageMessage("network load generator");
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  std::cout << "creating " << FLAGS_proc << " clients" << std::endl;
  std::cout << "connecting to " << FLAGS_host << ":" << FLAGS_port << std::endl;
  std::vector<pid_t> pids;
  for (int i = 0; i < FLAGS_proc; i++) {
    auto pid = fork();
    if (pid == 0) {
      client(FLAGS_host, FLAGS_port, FLAGS_req);
    } else if (pid == -1) {
    } else {
      pids.push_back(pid);
    }
  }

  for (auto pid : pids) {
    int sts;
    if (waitpid(pid, &sts, 0) == -1) {
      perror("sha256-client");
    }
  }
  return 0;
}