Update the connected field to metrics for DoT

The "connected" field for DoT is supported now. It is beneficial
to further tell the difference from sending a query via a new connection
and sending a query via an existing connection.

Bug: 144881031
Test: atest
Test: Manually added some logs to check the result
Change-Id: I3cc97660e8d5851af39d91f74d63babaa88179f8
diff --git a/DnsTlsDispatcher.cpp b/DnsTlsDispatcher.cpp
index 7ca0464..e961b9f 100644
--- a/DnsTlsDispatcher.cpp
+++ b/DnsTlsDispatcher.cpp
@@ -101,14 +101,17 @@
     for (const auto& server : orderedServers) {
         DnsQueryEvent* dnsQueryEvent =
                 statp->event->mutable_dns_query_events()->add_dns_query_event();
+
+        bool connectTriggered = false;
         Stopwatch queryStopwatch;
-        code = this->query(server, statp->_mark, query, ans, resplen);
+        code = this->query(server, statp->_mark, query, ans, resplen, &connectTriggered);
 
         dnsQueryEvent->set_latency_micros(saturate_cast<int32_t>(queryStopwatch.timeTakenUs()));
         dnsQueryEvent->set_dns_server_index(serverCount++);
         dnsQueryEvent->set_ip_version(ipFamilyToIPVersion(server.ss.ss_family));
         dnsQueryEvent->set_protocol(PROTO_DOT);
         dnsQueryEvent->set_type(getQueryType(query.base(), query.size()));
+        dnsQueryEvent->set_connected(connectTriggered);
 
         switch (code) {
             // These response codes are valid responses and not expected to
@@ -141,8 +144,9 @@
 }
 
 DnsTlsTransport::Response DnsTlsDispatcher::query(const DnsTlsServer& server, unsigned mark,
-                                                  const Slice query,
-                                                  const Slice ans, int *resplen) {
+                                                  const Slice query, const Slice ans, int* resplen,
+                                                  bool* connectTriggered) {
+    uint32_t connectCounter;
     const Key key = std::make_pair(mark, server);
     Transport* xport;
     {
@@ -155,6 +159,7 @@
             xport = it->second.get();
         }
         ++xport->useCount;
+        connectCounter = xport->transport.getConnectCounter();
     }
 
     LOG(DEBUG) << "Sending query of length " << query.size();
@@ -178,6 +183,7 @@
     auto now = std::chrono::steady_clock::now();
     {
         std::lock_guard guard(sLock);
+        *connectTriggered = (xport->transport.getConnectCounter() > connectCounter);
         --xport->useCount;
         xport->lastUsed = now;
         cleanup(now);
diff --git a/DnsTlsDispatcher.h b/DnsTlsDispatcher.h
index b8e9968..9eb6dfe 100644
--- a/DnsTlsDispatcher.h
+++ b/DnsTlsDispatcher.h
@@ -54,11 +54,12 @@
                                     const netdutils::Slice ans, int* _Nonnull resplen);
 
     // Given a |query|, sends it to the server on the network indicated by |mark|,
-    // and writes the response into |ans|,  and indicates
-    // the number of bytes written in |resplen|.  Returns a success or error code.
+    // and writes the response into |ans|, and indicates the number of bytes written in |resplen|.
+    // If the whole procedure above triggers (or experiences) any new connection, |connectTriggered|
+    // is set. Returns a success or error code.
     DnsTlsTransport::Response query(const DnsTlsServer& server, unsigned mark,
                                     const netdutils::Slice query, const netdutils::Slice ans,
-                                    int* _Nonnull resplen);
+                                    int* _Nonnull resplen, bool* _Nonnull connectTriggered);
 
   private:
     // This lock is static so that it can be used to annotate the Transport struct.
diff --git a/DnsTlsQueryMap.h b/DnsTlsQueryMap.h
index c5ab023..a6227f4 100644
--- a/DnsTlsQueryMap.h
+++ b/DnsTlsQueryMap.h
@@ -74,6 +74,9 @@
     // Returns true if there are no pending queries.
     bool empty();
 
+    // The maximum number of times we will send a query before abandoning it.
+    static constexpr int kMaxTries = 3;
+
   private:
     std::mutex mLock;
 
@@ -87,9 +90,6 @@
         std::promise<Result> result;
     };
 
-    // The maximum number of times we will send a query before abandoning it.
-    static constexpr int kMaxTries = 3;
-
     // Outstanding queries by newId.
     std::map<uint16_t, QueryPromise> mQueries GUARDED_BY(mLock);
 
diff --git a/DnsTlsTransport.cpp b/DnsTlsTransport.cpp
index 3453765..09daded 100644
--- a/DnsTlsTransport.cpp
+++ b/DnsTlsTransport.cpp
@@ -51,6 +51,11 @@
     return std::move(record->result);
 }
 
+uint32_t DnsTlsTransport::getConnectCounter() const {
+    std::lock_guard guard(mLock);
+    return mConnectCounter;
+}
+
 bool DnsTlsTransport::sendQuery(const DnsTlsQueryMap::Query q) {
     // Strip off the ID number and send the new ID instead.
     bool sent = mSocket->query(q.newId, netdutils::drop(q.query, 2));
@@ -63,6 +68,7 @@
 void DnsTlsTransport::doConnect() {
     LOG(DEBUG) << "Constructing new socket";
     mSocket = mFactory->createDnsTlsSocket(mServer, mMark, this, &mCache);
+    mConnectCounter++;
 
     if (mSocket) {
         auto queries = mQueries.getAll();
diff --git a/DnsTlsTransport.h b/DnsTlsTransport.h
index 3e43c7e..72c336a 100644
--- a/DnsTlsTransport.h
+++ b/DnsTlsTransport.h
@@ -57,12 +57,14 @@
     // on networks where it doesn't actually work.
     static bool validate(const DnsTlsServer& server, unsigned netid, uint32_t mark);
 
+    uint32_t getConnectCounter() const EXCLUDES(mLock);
+
     // Implement IDnsTlsSocketObserver
     void onResponse(std::vector<uint8_t> response) override;
     void onClosed() override EXCLUDES(mLock);
 
   private:
-    std::mutex mLock;
+    mutable std::mutex mLock;
 
     DnsTlsSessionCache mCache;
     DnsTlsQueryMap mQueries;
@@ -85,6 +87,9 @@
 
     // Send a query to the socket.
     bool sendQuery(const DnsTlsQueryMap::Query q) REQUIRES(mLock);
+
+    // The number of times an attempt to connect the nameserver.
+    uint32_t mConnectCounter GUARDED_BY(mLock) = 0;
 };
 
 }  // end of namespace net
diff --git a/resolv_tls_unit_test.cpp b/resolv_tls_unit_test.cpp
index 5d3ca26..41ede60 100644
--- a/resolv_tls_unit_test.cpp
+++ b/resolv_tls_unit_test.cpp
@@ -148,6 +148,7 @@
 
     EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
     EXPECT_EQ(QUERY, r.response);
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 // Fake Socket that echoes the observed query ID as the response body.
@@ -188,6 +189,7 @@
         // after each use.
         EXPECT_EQ(0, (r.response[2] << 8) | r.response[3]);
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 // These queries might be handled in serial or parallel as they race the
@@ -207,6 +209,7 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         EXPECT_EQ(QUERY, r.response);
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 // A server that waits until sDelay queries are queued before responding.
@@ -272,6 +275,7 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         EXPECT_EQ(QUERY, r.response);
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 TEST_F(TransportTest, ParallelColliding_Max) {
@@ -291,6 +295,7 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         EXPECT_EQ(QUERY, r.response);
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 TEST_F(TransportTest, ParallelUnique) {
@@ -310,6 +315,7 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         EXPECT_EQ(queries[i], r.response);
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 TEST_F(TransportTest, ParallelUnique_Max) {
@@ -331,6 +337,7 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         EXPECT_EQ(queries[i], r.response);
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 TEST_F(TransportTest, IdExhaustion) {
@@ -358,6 +365,7 @@
         EXPECT_EQ(std::future_status::timeout,
                 result.wait_for(std::chrono::duration<int>::zero()));
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 // Responses can come back from the server in any order.  This should have no
@@ -379,6 +387,7 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         EXPECT_EQ(queries[i], r.response);
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 TEST_F(TransportTest, ReverseOrder_Max) {
@@ -398,6 +407,7 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         EXPECT_EQ(queries[i], r.response);
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 // Returning null from the factory indicates a connection failure.
@@ -420,6 +430,7 @@
 
     EXPECT_EQ(DnsTlsTransport::Response::network_error, r.code);
     EXPECT_TRUE(r.response.empty());
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 // Simulate a socket that connects but then immediately receives a server
@@ -445,6 +456,9 @@
 
     EXPECT_EQ(DnsTlsTransport::Response::network_error, r.code);
     EXPECT_TRUE(r.response.empty());
+
+    // Reconnections are triggered since DnsTlsQueryMap is not empty.
+    EXPECT_EQ(transport.getConnectCounter(), static_cast<uint32_t>(DnsTlsQueryMap::kMaxTries));
 }
 
 // Simulate a server that occasionally closes the connection and silently
@@ -537,6 +551,9 @@
         EXPECT_EQ(DnsTlsTransport::Response::network_error, r.code);
         EXPECT_TRUE(r.response.empty());
     }
+
+    // Reconnections are triggered since DnsTlsQueryMap is not empty.
+    EXPECT_EQ(transport.getConnectCounter(), static_cast<uint32_t>(DnsTlsQueryMap::kMaxTries));
 }
 
 TEST_F(TransportTest, PartialDrop) {
@@ -561,6 +578,34 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         EXPECT_EQ(queries[i], r.response);
     }
+
+    // TODO: transport.getConnectCounter() seems not stable in this test. Find how to check the
+    // connect attempts for this test.
+}
+
+TEST_F(TransportTest, ConnectCounter) {
+    FakeSocketLimited::sLimit = 2;       // Close the socket after 2 queries.
+    FakeSocketLimited::sMaxSize = SIZE;  // No query drops.
+    FakeSocketFactory<FakeSocketLimited> factory;
+    DnsTlsTransport transport(SERVER1, MARK, &factory);
+
+    // Connecting on demand.
+    EXPECT_EQ(transport.getConnectCounter(), 0U);
+
+    const int num_queries = 10;
+    std::vector<std::future<DnsTlsTransport::Result>> results;
+    results.reserve(num_queries);
+    for (int i = 0; i < num_queries; i++) {
+        // Reconnections take place every two queries.
+        results.push_back(transport.query(makeSlice(QUERY)));
+    }
+    for (int i = 0; i < num_queries; i++) {
+        auto r = results[i].get();
+        EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
+    }
+
+    EXPECT_EQ(transport.getConnectCounter(),
+              static_cast<uint32_t>(num_queries / FakeSocketLimited::sLimit));
 }
 
 // Simulate a malfunctioning server that injects extra miscellaneous
@@ -604,6 +649,7 @@
         EXPECT_EQ(DnsTlsTransport::Response::success, r.code);
         // Don't check the response because this server is malfunctioning.
     }
+    EXPECT_EQ(transport.getConnectCounter(), 1U);
 }
 
 // Dispatcher tests
@@ -612,28 +658,38 @@
 TEST_F(DispatcherTest, Query) {
     bytevec ans(4096);
     int resplen = 0;
+    bool connectTriggered = false;
 
     auto factory = std::make_unique<FakeSocketFactory<FakeSocketEcho>>();
     DnsTlsDispatcher dispatcher(std::move(factory));
-    auto r = dispatcher.query(SERVER1, MARK, makeSlice(QUERY),
-                              makeSlice(ans), &resplen);
+    auto r = dispatcher.query(SERVER1, MARK, makeSlice(QUERY), makeSlice(ans), &resplen,
+                              &connectTriggered);
 
     EXPECT_EQ(DnsTlsTransport::Response::success, r);
     EXPECT_EQ(int(QUERY.size()), resplen);
+    EXPECT_TRUE(connectTriggered);
     ans.resize(resplen);
     EXPECT_EQ(QUERY, ans);
+
+    // Expect to reuse the connection.
+    r = dispatcher.query(SERVER1, MARK, makeSlice(QUERY), makeSlice(ans), &resplen,
+                         &connectTriggered);
+    EXPECT_EQ(DnsTlsTransport::Response::success, r);
+    EXPECT_FALSE(connectTriggered);
 }
 
 TEST_F(DispatcherTest, AnswerTooLarge) {
     bytevec ans(SIZE - 1);  // Too small to hold the answer
     int resplen = 0;
+    bool connectTriggered = false;
 
     auto factory = std::make_unique<FakeSocketFactory<FakeSocketEcho>>();
     DnsTlsDispatcher dispatcher(std::move(factory));
-    auto r = dispatcher.query(SERVER1, MARK, makeSlice(QUERY),
-                              makeSlice(ans), &resplen);
+    auto r = dispatcher.query(SERVER1, MARK, makeSlice(QUERY), makeSlice(ans), &resplen,
+                              &connectTriggered);
 
     EXPECT_EQ(DnsTlsTransport::Response::limit_error, r);
+    EXPECT_TRUE(connectTriggered);
 }
 
 template<class T>
@@ -678,10 +734,11 @@
             auto q = make_query(i, SIZE);
             bytevec ans(4096);
             int resplen = 0;
+            bool connectTriggered = false;
             unsigned mark = key.first;
             const DnsTlsServer& server = key.second;
-            auto r = dispatcher->query(server, mark, makeSlice(q),
-                                       makeSlice(ans), &resplen);
+            auto r = dispatcher->query(server, mark, makeSlice(q), makeSlice(ans), &resplen,
+                                       &connectTriggered);
             EXPECT_EQ(DnsTlsTransport::Response::success, r);
             EXPECT_EQ(int(q.size()), resplen);
             ans.resize(resplen);