1#include "InetAddress.h"
2#include "TcpStream.h"
3
4#include <limits>
5#include <memory>
6#include <unistd.h>
7#include <stdlib.h>
8#include <stdio.h>
9#include <time.h>
10
11#include <vector>
12
13int64_t kNanos = 1e9;
14int64_t clock_diff(struct timespec x, struct timespec y)
15{
16  return (x.tv_sec - y.tv_sec) * kNanos + x.tv_nsec - y.tv_nsec;
17}
18
19struct timespec get_time()
20{
21  struct timespec now;
22  clock_gettime(CLOCK_MONOTONIC, &now);
23  return now;
24}
25
26struct Sample
27{
28  int index;
29  int64_t start_nano;
30  int64_t rtt_nano;
31};
32
33std::vector<Sample> run(const char* host, int delay_ms, int length_s, int batch, int payload_size, bool silent)
34{
35  const struct timespec start = get_time();
36  std::vector<Sample> rtts;
37
38  InetAddress addr;
39  if (!InetAddress::resolve(host, 3007, &addr))
40  {
41    printf("Unable to resolve %s\n", host);
42    return rtts;
43  }
44
45  // printf("connecting to %s\n", addr.toIpPort().c_str());
46  TcpStreamPtr stream(TcpStream::connect(addr));
47  if (!stream)
48  {
49    printf("Unable to connect %s\n", addr.toIpPort().c_str());
50    perror("");
51    return rtts;
52  }
53
54  std::unique_ptr<char[]> payload(new char[payload_size]);
55  int64_t count = 0, sum_rtt = 0;
56  int64_t min_rtt = std::numeric_limits<int64_t>::max(), max_rtt = 0;
57
58  while (true) {
59    const struct timespec batch_start = get_time();
60    double elapsed_s = (double)clock_diff(batch_start, start) / kNanos;
61    if (elapsed_s >= length_s) {
62      if (silent && count > 0) {
63        printf("count %ld, avg rtt %.2fus, min %.2fus, max %.2fus\n",
64               count, sum_rtt / 1e3 / count, min_rtt / 1e3, max_rtt / 1e3);
65      }
66      break;
67    }
68
69    for (int i = 0; i < batch; ++i) {
70      const struct timespec before = get_time();
71      int nw = stream->sendAll(payload.get(), payload_size);
72      if (nw != payload_size)
73        return rtts;
74      int nr = stream->receiveAll(payload.get(), payload_size);
75      if (nr != payload_size)
76        return rtts;
77      const struct timespec after = get_time();
78      int64_t rtt = clock_diff(after, before);
79      ++count;
80      sum_rtt += rtt;
81      if (rtt > max_rtt)
82        max_rtt = rtt;
83      if (rtt < min_rtt)
84        min_rtt = rtt;
85
86      Sample s = {
87        .index = i,
88        .rtt_nano = rtt,
89      };
90      if (i == 0)
91        s.start_nano = clock_diff(before, start);
92      else
93        s.start_nano = clock_diff(before, batch_start);
94      if (!silent)
95        rtts.push_back(s);
96    }
97
98    if (delay_ms > 0) {
99      ::usleep(delay_ms * 1000);
100    }
101  }
102  return rtts;
103}
104
105int main(int argc, char* argv[])
106{
107  int opt;
108  int delay = 0, length = 3, batch = 4, payload = 1;
109  bool silent = false;
110  while ((opt = getopt(argc, argv, "b:d:l:p:s")) != -1) {
111    switch (opt) {
112      case 'b':
113        batch = atoi(optarg);
114        break;
115      case 'd':
116        delay = atoi(optarg);
117        break;
118      case 'l':
119        length = atoi(optarg);
120        break;
121      case 'p':
122        payload = atoi(optarg);
123        break;
124      case 's':
125        silent = true;
126        break;
127      default:
128        ;
129    }
130  }
131  if (optind >= argc) {
132    fprintf(stderr, "Usage:\nroundtrip_tcp [-b batch_size] [-d delay_ms] [-l length_in_seconds] echo_server_host\n");
133    return 1;
134  }
135  if (batch < 1) {
136    batch = 1;
137  }
138  if (delay < 0) {
139    delay = 0;
140  }
141  if (payload < 1) {
142    payload = 1;
143  }
144
145  std::vector<Sample> rtts = run(argv[optind], delay, length, batch, payload, silent);
146  if (!silent) {
147    printf("index start rtt\n");
148    for (Sample s : rtts) {
149      printf("%d %ld %ld\n", s.index, s.start_nano, s.rtt_nano);
150    }
151  }
152}
153