From 2c5f5a858743e3912d42662960e36151028b5ec6 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Mon, 4 Jun 2018 19:58:54 +0300
Subject: [PATCH] Improve PingConnection.

GitOrigin-RevId: be82fa18e2f9a154d3d276b5618d0ca72b9e4282
---
 td/mtproto/PingConnection.h           | 27 ++++++++++++++++++++++-----
 td/telegram/net/ConnectionCreator.cpp |  6 ++----
 test/mtproto.cpp                      |  6 ++++--
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/td/mtproto/PingConnection.h b/td/mtproto/PingConnection.h
index b726e48d..ca590d43 100644
--- a/td/mtproto/PingConnection.h
+++ b/td/mtproto/PingConnection.h
@@ -20,9 +20,11 @@
 
 namespace td {
 namespace mtproto {
+
 class PingConnection : private RawConnection::Callback {
  public:
-  explicit PingConnection(std::unique_ptr<RawConnection> raw_connection) : raw_connection_(std::move(raw_connection)) {
+  PingConnection(std::unique_ptr<RawConnection> raw_connection, size_t ping_count)
+      : raw_connection_(std::move(raw_connection)), ping_count_(ping_count) {
   }
 
   Fd &get_pollable() {
@@ -43,11 +45,17 @@ class PingConnection : private RawConnection::Callback {
       Random::secure_bytes(nonce.raw, sizeof(nonce));
       raw_connection_->send_no_crypto(PacketStorer<NoCryptoImpl>(1, create_storer(mtproto_api::req_pq_multi(nonce))));
       was_ping_ = true;
+      if (ping_count_ == 1) {
+        start_time_ = Time::now();
+      }
     }
     return raw_connection_->flush(AuthKey(), *this);
   }
   bool was_pong() const {
-    return was_pong_;
+    return finish_time_ > 0;
+  }
+  double rtt() const {
+    return finish_time_ - start_time_;
   }
 
   Status on_raw_packet(const PacketInfo &packet_info, BufferSlice packet) override {
@@ -56,14 +64,23 @@ class PingConnection : private RawConnection::Callback {
     }
     packet.confirm_read(12);
     // TODO: fetch_result
-    was_pong_ = true;
-    return Status::OK();
+
+    if (--ping_count_ > 0) {
+      was_ping_ = false;
+      return flush();
+    } else {
+      finish_time_ = Time::now();
+      return Status::OK();
+    }
   }
 
  private:
   std::unique_ptr<RawConnection> raw_connection_;
+  size_t ping_count_ = 1;
+  double start_time_ = 0.0;
+  double finish_time_ = 0.0;
   bool was_ping_ = false;
-  bool was_pong_ = false;
 };
+
 }  // namespace mtproto
 }  // namespace td
diff --git a/td/telegram/net/ConnectionCreator.cpp b/td/telegram/net/ConnectionCreator.cpp
index 17280684..f113c769 100644
--- a/td/telegram/net/ConnectionCreator.cpp
+++ b/td/telegram/net/ConnectionCreator.cpp
@@ -85,19 +85,17 @@ class PingActor : public Actor {
   PingActor(std::unique_ptr<mtproto::RawConnection> raw_connection,
             Promise<std::unique_ptr<mtproto::RawConnection>> promise, ActorShared<> parent)
       : promise_(std::move(promise)), parent_(std::move(parent)) {
-    ping_connection_ = std::make_unique<mtproto::PingConnection>(std::move(raw_connection));
+    ping_connection_ = std::make_unique<mtproto::PingConnection>(std::move(raw_connection), 2);
   }
 
  private:
   std::unique_ptr<mtproto::PingConnection> ping_connection_;
   Promise<std::unique_ptr<mtproto::RawConnection>> promise_;
   ActorShared<> parent_;
-  double start_at_;
 
   void start_up() override {
     ping_connection_->get_pollable().set_observer(this);
     subscribe(ping_connection_->get_pollable());
-    start_at_ = Time::now();
     set_timeout_in(10);
     yield();
   }
@@ -144,7 +142,7 @@ class PingActor : public Actor {
         raw_connection->close();
         promise_.set_error(std::move(status));
       } else {
-        raw_connection->rtt_ = Time::now() - start_at_;
+        raw_connection->rtt_ = ping_connection_->rtt();
         if (raw_connection->stats_callback()) {
           raw_connection->stats_callback()->on_pong();
         }
diff --git a/test/mtproto.cpp b/test/mtproto.cpp
index 2663d972..468e2be8 100644
--- a/test/mtproto.cpp
+++ b/test/mtproto.cpp
@@ -88,8 +88,10 @@ class TestPingActor : public Actor {
   Status *result_;
 
   void start_up() override {
-    ping_connection_ = std::make_unique<mtproto::PingConnection>(std::make_unique<mtproto::RawConnection>(
-        SocketFd::open(ip_address_).move_as_ok(), mtproto::TransportType{mtproto::TransportType::Tcp, 0, ""}, nullptr));
+    ping_connection_ = std::make_unique<mtproto::PingConnection>(
+        std::make_unique<mtproto::RawConnection>(SocketFd::open(ip_address_).move_as_ok(),
+                                                 mtproto::TransportType{mtproto::TransportType::Tcp, 0, ""}, nullptr),
+        3);
 
     ping_connection_->get_pollable().set_observer(this);
     subscribe(ping_connection_->get_pollable());