10 files changed, 171 insertions, 9 deletions

diff --git a/src/dnsperf.1.in b/src/dnsperf.1.in
index 99c0785..787f89f 100644
--- a/src/dnsperf.1.in
+++ b/src/dnsperf.1.in
@@ -413,6 +413,24 @@ Following type are available.
 \fBsendfailed\fR: Suppress messages about failure to send packets or if only parts of the packet were sent
 .br
 \fBsockready\fR: Suppress messages about socket readiness
+.br
+\fBunexpected\fR: Suppress messages about answers with an unexpected message ID
+.RE
+
+\fBnum-queries-per-conn=\fINUMBER\fR
+.br
+.RS
+This will limit the number of queries sent over a connection before
+triggering a re-connection. Once re-connected it will reset the counter and
+continue sending queries until the limit is reached again, triggering
+another re-connection and so on.
+Using this option will also enable counting number of responses received
+for each connection and once the limit is reached for sending queries it
+will wait until the same amount of responses has been received before
+re-connecting.
+Waiting for responses may timeout and the timeout used for this is the
+same as specified by \fB-t\fR.
+Note that this option is only useful for connection oriented protocols.
 .RE
 .SH "SEE ALSO"
 \fBresperf\fR(1)
diff --git a/src/dnsperf.c b/src/dnsperf.c
index c3ff297..40cedb4 100644
--- a/src/dnsperf.c
+++ b/src/dnsperf.c
@@ -89,6 +89,7 @@ typedef struct {
     bool               verbose;
     enum perf_net_mode mode;
     perf_suppress_t    suppress;
+    size_t             num_queries_per_conn;
 } config_t;
 
 typedef struct {
@@ -491,6 +492,8 @@ setup(int argc, char** argv, config_t* config)
         "the HTTP method to use for DNS-over-HTTPS: GET or POST", DEFAULT_DOH_METHOD, &doh_method);
     perf_long_opt_add("suppress", perf_opt_string, "message[,message,...]",
         "suppress messages/warnings, see man-page for list of message types", NULL, &local_suppress);
+    perf_long_opt_add("num-queries-per-conn", perf_opt_uint, "queries",
+        "Number of queries to send per connection", NULL, &config->num_queries_per_conn);
 
     bool log_stdout = false;
     perf_opt_add('W', perf_opt_boolean, NULL, "log warnings and errors to stdout instead of stderr", NULL, &log_stdout);
@@ -1000,7 +1003,7 @@ do_recv(void* arg)
                 perf_log_warning("received short response");
                 continue;
             }
-            if (recvd[i].unexpected) {
+            if (recvd[i].unexpected && !tinfo->config->suppress.unexpected) {
                 perf_log_warning("received a response with an "
                                  "unexpected (maybe timed out) "
                                  "id: %u",
@@ -1222,6 +1225,9 @@ threadinfo_init(threadinfo_t* tinfo, const config_t* config,
         if (!tinfo->socks[i]) {
             perf_log_fatal("perf_net_opensocket(): no socket returned, out of memory?");
         }
+        if (config->num_queries_per_conn && tinfo->socks[i]->num_queries_per_conn) {
+            tinfo->socks[i]->num_queries_per_conn(tinfo->socks[i], config->num_queries_per_conn, config->timeout);
+        }
     }
     tinfo->current_sock = 0;
 
diff --git a/src/net.h b/src/net.h
index 8f44ab3..19fdb35 100644
--- a/src/net.h
+++ b/src/net.h
@@ -80,6 +80,8 @@ typedef enum perf_socket_event {
 /* Callback for socket events related to connection oriented protocols, for statistics */
 typedef void (*perf_net_event_cb_t)(struct perf_net_socket* sock, perf_socket_event_t event, uint64_t elapsed_time);
 
+typedef void (*perf_net_num_queries_per_conn_t)(struct perf_net_socket* sock, size_t num_queries_per_conn, size_t timeout);
+
 struct perf_net_socket {
     void* data; /* user data */
 
@@ -91,6 +93,8 @@ struct perf_net_socket {
     perf_net_sockready_t sockready;
     perf_net_have_more_t have_more;
 
+    perf_net_num_queries_per_conn_t num_queries_per_conn;
+
     /*
      * Not set by protocol, set by caller.
      * May be 0 if caller don't care.
diff --git a/src/net_doh.c b/src/net_doh.c
index 0920696..6bf673d 100644
--- a/src/net_doh.c
+++ b/src/net_doh.c
@@ -118,6 +118,10 @@ struct perf__doh_socket {
 
     char   recvbuf[TCP_RECV_BUF_SIZE];
     size_t recv_at;
+
+    size_t       num_queries_per_conn, nqpc_timeout;
+    unsigned int nqpc_sent, nqpc_recv;
+    uint64_t     nqpc_ts;
 };
 
 static pthread_mutex_t            _nghttp2_lock      = PTHREAD_MUTEX_INITIALIZER;
@@ -157,6 +161,10 @@ static void perf__doh_connect(struct perf_net_socket* sock)
 {
     int ret;
 
+    self->nqpc_sent = 0;
+    ck_pr_store_uint(&self->nqpc_recv, 0);
+    self->nqpc_ts = 0;
+
     nghttp2_session_del(self->http2.session);
     ret = nghttp2_session_client_new2(&self->http2.session, _nghttp2_callbacks, sock, _nghttp2_option);
     if (ret < 0) {
@@ -394,6 +402,9 @@ static ssize_t perf__doh_recv(struct perf_net_socket* sock, void* buf, size_t le
 
         // self->have_more = false; TODO
         PERF_UNLOCK(&self->lock);
+        if (self->num_queries_per_conn) {
+            ck_pr_inc_uint(&self->nqpc_recv);
+        }
         return len;
     }
 
@@ -586,6 +597,8 @@ static ssize_t perf__doh_sendto(struct perf_net_socket* sock, uint16_t qid, cons
     }
     PERF_UNLOCK(&self->lock);
 
+    self->nqpc_sent++;
+
     return len;
 }
 
@@ -628,6 +641,21 @@ static int perf__doh_sockready(struct perf_net_socket* sock, int pipe_fd, int64_
             }
             self->is_sending = false;
             sent             = true;
+            self->nqpc_sent++;
+        }
+        if (self->num_queries_per_conn && self->nqpc_sent >= self->num_queries_per_conn) {
+            if (!self->nqpc_ts) {
+                self->nqpc_ts = perf_get_time() + self->nqpc_timeout;
+            }
+            unsigned int r = ck_pr_load_uint(&self->nqpc_recv);
+            if (r >= self->nqpc_sent || perf_get_time() > self->nqpc_ts) {
+                self->do_reconnect = true;
+            }
+            PERF_UNLOCK(&self->lock);
+            if (sent && sock->sent) {
+                sock->sent(sock, self->qid);
+            }
+            return 0;
         }
         PERF_UNLOCK(&self->lock);
         if (sent && sock->sent) {
@@ -898,6 +926,12 @@ static int select_next_proto_cb(SSL* ssl, unsigned char** out,
 }
 #endif /* !OPENSSL_NO_NEXTPROTONEG */
 
+static void perf__doh_num_queries_per_conn(struct perf_net_socket* sock, size_t num_queries_per_conn, size_t timeout)
+{
+    self->num_queries_per_conn = num_queries_per_conn;
+    self->nqpc_timeout         = timeout;
+}
+
 struct perf_net_socket* perf_net_doh_opensocket(const perf_sockaddr_t* server, const perf_sockaddr_t* local, size_t bufsize, void* data, perf_net_sent_cb_t sent, perf_net_event_cb_t event)
 {
     struct perf__doh_socket* tmp  = calloc(1, sizeof(struct perf__doh_socket)); // clang scan-build
@@ -915,6 +949,8 @@ struct perf_net_socket* perf_net_doh_opensocket(const perf_sockaddr_t* server, c
     sock->sockready = perf__doh_sockready;
     sock->have_more = perf__doh_have_more;
 
+    sock->num_queries_per_conn = perf__doh_num_queries_per_conn;
+
     sock->data  = data;
     sock->sent  = sent;
     sock->event = event;
diff --git a/src/net_dot.c b/src/net_dot.c
index 6b53e35..a68a84f 100644
--- a/src/net_dot.c
+++ b/src/net_dot.c
@@ -55,12 +55,20 @@ struct perf__dot_socket {
 
     uint64_t            conn_ts;
     perf_socket_event_t conn_event, conning_event;
+
+    size_t       num_queries_per_conn, nqpc_timeout;
+    unsigned int nqpc_sent, nqpc_recv;
+    uint64_t     nqpc_ts;
 };
 
 static void perf__dot_connect(struct perf_net_socket* sock)
 {
     int ret;
 
+    self->nqpc_sent = 0;
+    ck_pr_store_uint(&self->nqpc_recv, 0);
+    self->nqpc_ts = 0;
+
     int fd = socket(self->server.sa.sa.sa_family, SOCK_STREAM, 0);
     if (fd == -1) {
         char __s[256];
@@ -204,6 +212,9 @@ static ssize_t perf__dot_recv(struct perf_net_socket* sock, void* buf, size_t le
     memcpy(buf, self->recvbuf + 2, len < dnslen ? len : dnslen);
     memmove(self->recvbuf, self->recvbuf + 2 + dnslen, self->at - 2 - dnslen);
     self->at -= 2 + dnslen;
+    if (self->num_queries_per_conn) {
+        ck_pr_inc_uint(&self->nqpc_recv);
+    }
 
     if (self->at > 2) {
         memcpy(&dnslen2, self->recvbuf, 2);
@@ -283,6 +294,8 @@ static ssize_t perf__dot_sendto(struct perf_net_socket* sock, uint16_t qid, cons
     }
     PERF_UNLOCK(&self->lock);
 
+    self->nqpc_sent++;
+
     return n - 2;
 }
 
@@ -342,9 +355,22 @@ static int perf__dot_sockready(struct perf_net_socket* sock, int pipe_fd, int64_
             if (sock->sent) {
                 sock->sent(sock, self->qid);
             }
-            return 1;
+            self->nqpc_sent++;
+        } else {
+            PERF_UNLOCK(&self->lock);
+        }
+        if (self->num_queries_per_conn && self->nqpc_sent >= self->num_queries_per_conn) {
+            if (!self->nqpc_ts) {
+                self->nqpc_ts = perf_get_time() + self->nqpc_timeout;
+            }
+            unsigned int r = ck_pr_load_uint(&self->nqpc_recv);
+            if (r >= self->nqpc_sent || perf_get_time() > self->nqpc_ts) {
+                PERF_LOCK(&self->lock);
+                perf__dot_reconnect(sock);
+                PERF_UNLOCK(&self->lock);
+            }
+            return 0;
         }
-        PERF_UNLOCK(&self->lock);
         return 1;
     }
 
@@ -417,6 +443,12 @@ static bool perf__dot_have_more(struct perf_net_socket* sock)
     return self->have_more;
 }
 
+static void perf__dot_num_queries_per_conn(struct perf_net_socket* sock, size_t num_queries_per_conn, size_t timeout)
+{
+    self->num_queries_per_conn = num_queries_per_conn;
+    self->nqpc_timeout         = timeout;
+}
+
 struct perf_net_socket* perf_net_dot_opensocket(const perf_sockaddr_t* server, const perf_sockaddr_t* local, size_t bufsize, void* data, perf_net_sent_cb_t sent, perf_net_event_cb_t event)
 {
     struct perf__dot_socket* tmp  = calloc(1, sizeof(struct perf__dot_socket)); // clang scan-build
@@ -434,6 +466,8 @@ struct perf_net_socket* perf_net_dot_opensocket(const perf_sockaddr_t* server, c
     sock->sockready = perf__dot_sockready;
     sock->have_more = perf__dot_have_more;
 
+    sock->num_queries_per_conn = perf__dot_num_queries_per_conn;
+
     sock->data  = data;
     sock->sent  = sent;
     sock->event = event;
diff --git a/src/net_tcp.c b/src/net_tcp.c
index 97d1f65..5473e4a 100644
--- a/src/net_tcp.c
+++ b/src/net_tcp.c
@@ -80,13 +80,20 @@ struct perf__tcp_socket {
 
     uint64_t            conn_ts;
     perf_socket_event_t conn_event, conning_event;
+
+    size_t       num_queries_per_conn, nqpc_timeout;
+    unsigned int nqpc_sent, nqpc_recv;
+    uint64_t     nqpc_ts;
 };
 
 static int perf__tcp_connect(struct perf_net_socket* sock)
 {
     int fd;
 
-    self->is_ready = true;
+    self->is_ready  = true;
+    self->nqpc_sent = 0;
+    ck_pr_store_uint(&self->nqpc_recv, 0);
+    self->nqpc_ts = 0;
 
     fd = socket(self->server.sa.sa.sa_family, SOCK_STREAM, 0);
     if (fd == -1) {
@@ -192,6 +199,9 @@ static ssize_t perf__tcp_recv(struct perf_net_socket* sock, void* buf, size_t le
     memcpy(buf, self->recvbuf + 2, len < dnslen ? len : dnslen);
     memmove(self->recvbuf, self->recvbuf + 2 + dnslen, self->at - 2 - dnslen);
     self->at -= 2 + dnslen;
+    if (self->num_queries_per_conn) {
+        ck_pr_inc_uint(&self->nqpc_recv);
+    }
 
     if (self->at > 2) {
         memcpy(&dnslen2, self->recvbuf, 2);
@@ -251,6 +261,7 @@ static ssize_t perf__tcp_sendto(struct perf_net_socket* sock, uint16_t qid, cons
         errno            = EINPROGRESS;
         return -1;
     }
+    self->nqpc_sent++;
 
     return n - 2;
 }
@@ -302,6 +313,17 @@ static int perf__tcp_sockready(struct perf_net_socket* sock, int pipe_fd, int64_
             if (sock->sent) {
                 sock->sent(sock, self->qid);
             }
+            self->nqpc_sent++;
+        }
+        if (self->num_queries_per_conn && self->nqpc_sent >= self->num_queries_per_conn) {
+            if (!self->nqpc_ts) {
+                self->nqpc_ts = perf_get_time() + self->nqpc_timeout;
+            }
+            unsigned int r = ck_pr_load_uint(&self->nqpc_recv);
+            if (r >= self->nqpc_sent || perf_get_time() > self->nqpc_ts) {
+                self->need_reconnect = true;
+            }
+            return 0;
         }
         return 1;
     }
@@ -352,6 +374,7 @@ conn_cont:
             if (sock->sent) {
                 sock->sent(sock, self->qid);
             }
+            self->nqpc_sent++;
         }
         return 1;
     }
@@ -367,6 +390,12 @@ static bool perf__tcp_have_more(struct perf_net_socket* sock)
     return self->have_more;
 }
 
+static void perf__tcp_num_queries_per_conn(struct perf_net_socket* sock, size_t num_queries_per_conn, size_t timeout)
+{
+    self->num_queries_per_conn = num_queries_per_conn;
+    self->nqpc_timeout         = timeout;
+}
+
 struct perf_net_socket* perf_net_tcp_opensocket(const perf_sockaddr_t* server, const perf_sockaddr_t* local, size_t bufsize, void* data, perf_net_sent_cb_t sent, perf_net_event_cb_t event)
 {
     struct perf__tcp_socket* tmp  = calloc(1, sizeof(struct perf__tcp_socket)); // clang scan-build
@@ -384,6 +413,8 @@ struct perf_net_socket* perf_net_tcp_opensocket(const perf_sockaddr_t* server, c
     sock->sockready = perf__tcp_sockready;
     sock->have_more = perf__tcp_have_more;
 
+    sock->num_queries_per_conn = perf__tcp_num_queries_per_conn;
+
     sock->data  = data;
     sock->sent  = sent;
     sock->event = event;
diff --git a/src/opt.c b/src/opt.c
index 01949e6..75830bd 100644
--- a/src/opt.c
+++ b/src/opt.c
@@ -369,7 +369,7 @@ void perf_opt_parse(int argc, char** argv)
 
 perf_suppress_t perf_opt_parse_suppress(const char* val)
 {
-    perf_suppress_t s = { false, false, false };
+    perf_suppress_t s = { false, false, false, false };
 
     while (val && *val) {
         const char* next = strchr(val, ',');
@@ -390,6 +390,8 @@ perf_suppress_t perf_opt_parse_suppress(const char* val)
             s.sendfailed = true;
         } else if (!strncmp(val, "sockready", len)) {
             s.sockready = true;
+        } else if (!strncmp(val, "unexpected", len)) {
+            s.unexpected = true;
         } else {
             fprintf(stderr, "unknown message type to suppress: %.*s\n", len, val);
             perf_opt_usage();
diff --git a/src/opt.h b/src/opt.h
index 7e8433f..0864e2b 100644
--- a/src/opt.h
+++ b/src/opt.h
@@ -37,6 +37,7 @@ typedef struct {
     bool congestion;
     bool sendfailed;
     bool sockready;
+    bool unexpected;
 } perf_suppress_t;
 
 void perf_opt_add(char c, perf_opttype_t type, const char* desc, const char* help, const char* defval, void* valp);
diff --git a/src/resperf.c b/src/resperf.c
index a62c819..ac41b38 100644
--- a/src/resperf.c
+++ b/src/resperf.c
@@ -255,6 +255,8 @@ static void setup(int argc, char** argv)
     const char*  doh_method      = DEFAULT_DOH_METHOD;
     const char*  local_suppress  = 0;
 
+    size_t num_queries_per_conn = 0;
+
     sock_family     = AF_UNSPEC;
     server_port     = 0;
     local_port      = DEFAULT_LOCAL_PORT;
@@ -337,6 +339,8 @@ static void setup(int argc, char** argv)
         "the HTTP method to use for DNS-over-HTTPS: GET or POST", DEFAULT_DOH_METHOD, &doh_method);
     perf_long_opt_add("suppress", perf_opt_string, "message[,message,...]",
         "suppress messages/warnings, see dnsperf(1) man-page for list of message types", NULL, &local_suppress);
+    perf_long_opt_add("num-queries-per-conn", perf_opt_uint, "queries",
+        "Number of queries to send per connection", NULL, &num_queries_per_conn);
 
     perf_opt_parse(argc, argv);
 
@@ -428,6 +432,9 @@ static void setup(int argc, char** argv)
         if (!socks[i]) {
             perf_log_fatal("perf_net_opensocket(): no socket returned, out of memory?");
         }
+        if (num_queries_per_conn && socks[i]->num_queries_per_conn) {
+            socks[i]->num_queries_per_conn(socks[i], num_queries_per_conn, query_timeout);
+        }
     }
 }
 
@@ -738,7 +745,9 @@ try_process_response(unsigned int sockindex)
 
     size_t idx = qid * nsocks + sockindex;
     if (idx >= max_outstanding || queries[idx].list != &outstanding_list) {
-        perf_log_warning("received a response with an unexpected id: %u", qid);
+        if (!suppress.unexpected) {
+            perf_log_warning("received a response with an unexpected id: %u", qid);
+        }
         return;
     }
     q = &queries[idx];
diff --git a/src/test/test5.sh b/src/test/test5.sh
index 515a5ff..4831eb4 100755
--- a/src/test/test5.sh
+++ b/src/test/test5.sh
@@ -3,21 +3,25 @@
 test "$TEST_DNSPERF_WITH_NETWORK" = "1" || exit 0
 
 dumdumd=`which dumdumd`
+dumdohd=`which dumdohd`
 
-if [ -n "$dumdumd" ]; then
-    pkill -9 dumdumd || true
+pkill -9 dumdumd || true
+pkill -9 dumdohd || true
 
+if [ -n "$dumdumd" ]; then
     $dumdumd 127.0.0.1 5353 -r -D 100 &
     pid="$!"
     sleep 2
     ../dnsperf -s 127.0.0.1 -p 5353 -d "$srcdir/datafile" -t 2 -l 2 -Q 10 -m tcp
     kill "$pid"
+    sleep 2
 
     $dumdumd 127.0.0.1 5353 -r -D 10 &
     pid="$!"
     sleep 2
     ../dnsperf -s 127.0.0.1 -p 5353 -d "$srcdir/datafile" -t 2 -l 10 -Q 100 -m tcp
     kill "$pid"
+    sleep 2
 
     rm -f key.pem cert.pem
     openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365 -nodes -subj "/C=AU/ST=Some-State/O=Internet Widgits Pty Ltd"
@@ -27,12 +31,29 @@ if [ -n "$dumdumd" ]; then
     sleep 2
     ../dnsperf -s 127.0.0.1 -p 5353 -d "$srcdir/datafile" -t 2 -l 2 -Q 10 -m dot
     kill "$pid"
+    sleep 2
 
     $dumdumd 127.0.0.1 5353 -r -T -D 10 &
     pid="$!"
     sleep 2
     ../dnsperf -s 127.0.0.1 -p 5353 -d "$srcdir/datafile" -t 2 -l 10 -Q 100 -m dot
     kill "$pid"
+fi
+
+if [ -n "$dumdohd" ]; then
+    rm -f key.pem cert.pem
+    openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365 -nodes -subj "/C=AU/ST=Some-State/O=Internet Widgits Pty Ltd"
 
-    pkill -9 dumdumd || true
+    $dumdohd 5353 key.pem cert.pem -D 100 &
+    pid="$!"
+    sleep 2
+    ../dnsperf -s 127.0.0.1 -p 5353 -d "$srcdir/datafile" -t 2 -l 2 -Q 10 -m doh
+    kill "$pid"
+    sleep 2
+
+    $dumdohd 5353 key.pem cert.pem -D 10 &
+    pid="$!"
+    sleep 2
+    ../dnsperf -s 127.0.0.1 -p 5353 -d "$srcdir/datafile" -t 2 -l 10 -Q 100 -m doh
+    kill "$pid"
 fi