1 files changed, 25 insertions, 42 deletions

diff --git a/tests/dns_metrics_listener/dns_metrics_listener.cpp b/tests/dns_metrics_listener/dns_metrics_listener.cpp
index 1715118d..c1d4883d 100644
--- a/tests/dns_metrics_listener/dns_metrics_listener.cpp
+++ b/tests/dns_metrics_listener/dns_metrics_listener.cpp
@@ -18,17 +18,13 @@
 
 #include <thread>
 
-#include <android-base/chrono_utils.h>
 #include <android-base/format.h>
 
-namespace android {
-namespace net {
-namespace metrics {
+namespace android::net::metrics {
 
 using android::base::ScopedLockAssertion;
 using std::chrono::milliseconds;
 
-constexpr milliseconds kRetryIntervalMs{20};
 constexpr milliseconds kEventTimeoutMs{5000};
 
 bool DnsMetricsListener::DnsEvent::operator==(const DnsMetricsListener::DnsEvent& o) const {
@@ -49,6 +45,7 @@ std::ostream& operator<<(std::ostream& os, const DnsMetricsListener::DnsEvent& d
     std::lock_guard lock(mMutex);
     mUnexpectedNat64PrefixUpdates++;
     if (netId == mNetId) mNat64Prefix = added ? prefixString : "";
+    mCv.notify_all();
     return ::ndk::ScopedAStatus::ok();
 }
 
@@ -60,7 +57,7 @@ std::ostream& operator<<(std::ostream& os, const DnsMetricsListener::DnsEvent& d
         // keep updating the server to have latest validation status.
         mValidationRecords.insert_or_assign({netId, ipAddress}, validated);
     }
-    mCv.notify_one();
+    mCv.notify_all();
     return ::ndk::ScopedAStatus::ok();
 }
 
@@ -74,39 +71,32 @@ std::ostream& operator<<(std::ostream& os, const DnsMetricsListener::DnsEvent& d
         mDnsEventRecords.push(
                 {netId, eventType, returnCode, hostname, ipAddresses, ipAddressesCount});
     }
+    mCv.notify_all();
     return ::ndk::ScopedAStatus::ok();
 }
 
 bool DnsMetricsListener::waitForNat64Prefix(ExpectNat64PrefixStatus status, milliseconds timeout) {
-    android::base::Timer t;
-    while (t.duration() < timeout) {
-        {
-            std::lock_guard lock(mMutex);
-            if ((status == EXPECT_FOUND && !mNat64Prefix.empty()) ||
-                (status == EXPECT_NOT_FOUND && mNat64Prefix.empty())) {
-                mUnexpectedNat64PrefixUpdates--;
-                return true;
-            }
-        }
-        std::this_thread::sleep_for(kRetryIntervalMs);
+    std::unique_lock lock(mMutex);
+    ScopedLockAssertion assume_lock(mMutex);
+
+    if (mCv.wait_for(lock, timeout, [&]() REQUIRES(mMutex) {
+            return (status == EXPECT_FOUND && !mNat64Prefix.empty()) ||
+                   (status == EXPECT_NOT_FOUND && mNat64Prefix.empty());
+        })) {
+        mUnexpectedNat64PrefixUpdates--;
+        return true;
     }
+
+    // Timeout.
     return false;
 }
 
 bool DnsMetricsListener::waitForPrivateDnsValidation(const std::string& serverAddr,
                                                      const bool validated) {
-    const auto now = std::chrono::steady_clock::now();
-
     std::unique_lock lock(mMutex);
-    ScopedLockAssertion assume_lock(mMutex);
-
-    // onPrivateDnsValidationEvent() might already be invoked. Search for the record first.
-    do {
-        if (findAndRemoveValidationRecord({mNetId, serverAddr}, validated)) return true;
-    } while (mCv.wait_until(lock, now + kEventTimeoutMs) != std::cv_status::timeout);
-
-    // Timeout.
-    return false;
+    return mCv.wait_for(lock, kEventTimeoutMs, [&]() REQUIRES(mMutex) {
+        return findAndRemoveValidationRecord({mNetId, serverAddr}, validated);
+    });
 }
 
 bool DnsMetricsListener::findAndRemoveValidationRecord(const ServerKey& key, const bool value) {
@@ -119,24 +109,17 @@ bool DnsMetricsListener::findAndRemoveValidationRecord(const ServerKey& key, con
 }
 
 std::optional<DnsMetricsListener::DnsEvent> DnsMetricsListener::popDnsEvent() {
-    // Wait until the queue is not empty or timeout.
-    android::base::Timer t;
-    while (t.duration() < milliseconds{1000}) {
-        {
-            std::lock_guard lock(mMutex);
-            if (!mDnsEventRecords.empty()) break;
-        }
-        std::this_thread::sleep_for(kRetryIntervalMs);
-    }
+    std::unique_lock lock(mMutex);
+    ScopedLockAssertion assume_lock(mMutex);
 
-    std::lock_guard lock(mMutex);
-    if (mDnsEventRecords.empty()) return std::nullopt;
+    if (!mCv.wait_for(lock, kEventTimeoutMs,
+                      [&]() REQUIRES(mMutex) { return !mDnsEventRecords.empty(); })) {
+        return std::nullopt;
+    }
 
     auto ret = mDnsEventRecords.front();
     mDnsEventRecords.pop();
     return ret;
 }
 
-}  // namespace metrics
-}  // namespace net
-}  // namespace android
+}  // namespace android::net::metrics