From 99e08b7406c5161ba21b78f4904bfdc32d99877a Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Tue, 22 Jun 2021 23:22:13 +0300
Subject: [PATCH] Make network receive time metric to correctly include the
 time spent waiting for data from the client to INSERT #9958

---
 src/IO/ReadBufferFromPocoSocket.cpp             |  8 +++++++-
 ...network_receive_time_metric_insert.reference |  1 +
 .../01923_network_receive_time_metric_insert.sh | 17 +++++++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)
 create mode 100644 tests/queries/0_stateless/01923_network_receive_time_metric_insert.reference
 create mode 100755 tests/queries/0_stateless/01923_network_receive_time_metric_insert.sh

diff --git a/src/IO/ReadBufferFromPocoSocket.cpp b/src/IO/ReadBufferFromPocoSocket.cpp
index d1ceaaf6a35..e043764d280 100644
--- a/src/IO/ReadBufferFromPocoSocket.cpp
+++ b/src/IO/ReadBufferFromPocoSocket.cpp
@@ -76,7 +76,13 @@ ReadBufferFromPocoSocket::ReadBufferFromPocoSocket(Poco::Net::Socket & socket_,
 
 bool ReadBufferFromPocoSocket::poll(size_t timeout_microseconds) const
 {
-    return available() || socket.poll(timeout_microseconds, Poco::Net::Socket::SELECT_READ | Poco::Net::Socket::SELECT_ERROR);
+    if (available())
+        return true;
+
+    Stopwatch watch;
+    bool res = socket.poll(timeout_microseconds, Poco::Net::Socket::SELECT_READ | Poco::Net::Socket::SELECT_ERROR);
+    ProfileEvents::increment(ProfileEvents::NetworkReceiveElapsedMicroseconds, watch.elapsedMicroseconds());
+    return res;
 }
 
 }
diff --git a/tests/queries/0_stateless/01923_network_receive_time_metric_insert.reference b/tests/queries/0_stateless/01923_network_receive_time_metric_insert.reference
new file mode 100644
index 00000000000..d00491fd7e5
--- /dev/null
+++ b/tests/queries/0_stateless/01923_network_receive_time_metric_insert.reference
@@ -0,0 +1 @@
+1
diff --git a/tests/queries/0_stateless/01923_network_receive_time_metric_insert.sh b/tests/queries/0_stateless/01923_network_receive_time_metric_insert.sh
new file mode 100755
index 00000000000..cd3202e94c9
--- /dev/null
+++ b/tests/queries/0_stateless/01923_network_receive_time_metric_insert.sh
@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+# shellcheck source=../shell_config.sh
+. "$CURDIR"/../shell_config.sh
+
+${CLICKHOUSE_CLIENT} --multiquery --query "DROP TABLE IF EXISTS t; CREATE TABLE t (x UInt64) ENGINE = Memory;"
+
+# Rate limit is chosen for operation to spent about one second.
+seq 1 1000 | pv --quiet --rate-limit 3893 | ${CLICKHOUSE_CLIENT} --query "INSERT INTO t FORMAT TSV"
+
+# We check that the value of NetworkReceiveElapsedMicroseconds is correctly include the time spent waiting data from the client.
+${CLICKHOUSE_CLIENT} --multiquery --query "SYSTEM FLUSH LOGS;
+    SELECT ProfileEvents.Values[indexOf(ProfileEvents.Names, 'NetworkReceiveElapsedMicroseconds')] >= 1000000 FROM system.query_log
+        WHERE current_database = currentDatabase() AND query_kind = 'Insert' AND event_date >= yesterday() AND type = 2 ORDER BY event_time DESC LIMIT 1;"
+
+${CLICKHOUSE_CLIENT} --query "DROP TABLE t"