Merge "[Cronet] skip presubmit for flaky tests" into main
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 83d435a..3384f7b 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -183,6 +183,9 @@
         },
         {
           "exclude-annotation": "com.android.testutils.ConnectivityModuleTest"
+        },
+        {
+          "exclude-annotation": "com.android.testutils.DnsResolverModuleTest"
         }
       ]
     },
@@ -198,6 +201,9 @@
         },
         {
           "exclude-annotation": "androidx.test.filters.RequiresDevice"
+        },
+        {
+          "exclude-annotation": "com.android.testutils.DnsResolverModuleTest"
         }
       ]
     },
diff --git a/remote_auth/OWNERS b/remote_auth/OWNERS
new file mode 100644
index 0000000..25a32b9
--- /dev/null
+++ b/remote_auth/OWNERS
@@ -0,0 +1,14 @@
+# Bug component: 1145231
+# Bug template url: http://b/new?component=1145231&template=1715387
+billyhuang@google.com
+boetger@google.com
+casbor@google.com
+derekjedral@google.com
+dlm@google.com
+igorzas@google.com
+jacobhobbie@google.com
+jasonsun@google.com
+jianbing@google.com
+jinjiechen@google.com
+justinmcclain@google.com
+salilr@google.com
diff --git a/remote_auth/README.md b/remote_auth/README.md
new file mode 100644
index 0000000..384fcf7
--- /dev/null
+++ b/remote_auth/README.md
@@ -0,0 +1,47 @@
+# RemoteAuth Mainline Module
+
+This directory contains code for the RemoteAuth module.
+
+## Directory Structure
+
+`framework`
+ - Contains client side APIs and AIDL files.
+
+`jni`
+ - JNI wrapper for invoking Android APIs from native code.
+
+`native`
+ - Native code implementation for RemoteAuth module services.
+
+`service`
+ - Server side implementation for RemoteAuth module services.
+
+`tests`
+ - Unit/Multi devices tests for RemoteAuth module (both Java and native code).
+
+## IDE setup
+
+### AIDEGen
+
+AIDEGen is deprecated, prefer ASfP [go/asfp](http://go/asfp)
+```sh
+$ source build/envsetup.sh && lunch <TARGET>
+$ cd packages/modules/Connectivity
+$ aidegen .
+# This will launch Intellij project for RemoteAuth module.
+```
+
+### ASfP
+
+See full instructions for ASfP at [go/asfp-getting-started](http://go/asfp-getting-started)
+
+## Build and Install
+
+```sh
+$ source build/envsetup.sh && lunch <TARGET>
+$ m com.google.android.tethering deapexer
+$ $ANDROID_BUILD_TOP/out/host/linux-x86/bin/deapexer decompress --input \
+    ${ANDROID_PRODUCT_OUT}/system/apex/com.google.android.tethering.capex \
+    --output /tmp/tethering.apex
+$ adb install -r /tmp/tethering.apex
+```
diff --git a/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
index 6c8e273..bcedbef 100644
--- a/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
+++ b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
@@ -16,6 +16,14 @@
 
 package com.android.metrics;
 
+import static com.android.metrics.NetworkNsdReported.Builder;
+
+import android.stats.connectivity.MdnsQueryResult;
+import android.stats.connectivity.NsdEventType;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.ConnectivityStatsLog;
+
 /**
  * Class to record the NetworkNsdReported into statsd. Each client should create this class to
  * report its data.
@@ -25,11 +33,93 @@
     private final boolean mIsLegacy;
     // The client id.
     private final int mClientId;
+    private final Dependencies mDependencies;
 
     public NetworkNsdReportedMetrics(boolean isLegacy, int clientId) {
-        mIsLegacy = isLegacy;
-        mClientId = clientId;
+        this(isLegacy, clientId, new Dependencies());
     }
 
-    // TODO: Report metrics data.
+    @VisibleForTesting
+    NetworkNsdReportedMetrics(boolean isLegacy, int clientId, Dependencies dependencies) {
+        mIsLegacy = isLegacy;
+        mClientId = clientId;
+        mDependencies = dependencies;
+    }
+
+    /**
+     * Dependencies of NetworkNsdReportedMetrics, for injection in tests.
+     */
+    public static class Dependencies {
+
+        /**
+         * @see ConnectivityStatsLog
+         */
+        public void statsWrite(NetworkNsdReported event) {
+            ConnectivityStatsLog.write(ConnectivityStatsLog.NETWORK_NSD_REPORTED,
+                    event.getIsLegacy(),
+                    event.getClientId(),
+                    event.getTransactionId(),
+                    event.getIsKnownService(),
+                    event.getType().getNumber(),
+                    event.getEventDurationMillisec(),
+                    event.getQueryResult().getNumber(),
+                    event.getFoundServiceCount(),
+                    event.getFoundCallbackCount(),
+                    event.getLostCallbackCount(),
+                    event.getRepliedRequestsCount());
+        }
+    }
+
+    private Builder makeReportedBuilder() {
+        final Builder builder = NetworkNsdReported.newBuilder();
+        builder.setIsLegacy(mIsLegacy);
+        builder.setClientId(mClientId);
+        return builder;
+    }
+
+    /**
+     * Report service registration succeeded metric data.
+     *
+     * @param transactionId The transaction id of service registration.
+     * @param durationMs The duration of service registration success.
+     */
+    public void reportServiceRegistrationSucceeded(int transactionId, long durationMs) {
+        final Builder builder = makeReportedBuilder();
+        builder.setTransactionId(transactionId);
+        builder.setType(NsdEventType.NET_REGISTER);
+        builder.setQueryResult(MdnsQueryResult.MQR_SERVICE_REGISTERED);
+        builder.setEventDurationMillisec(durationMs);
+        mDependencies.statsWrite(builder.build());
+    }
+
+    /**
+     * Report service registration failed metric data.
+     *
+     * @param transactionId The transaction id of service registration.
+     * @param durationMs The duration of service registration failed.
+     */
+    public void reportServiceRegistrationFailed(int transactionId, long durationMs) {
+        final Builder builder = makeReportedBuilder();
+        builder.setTransactionId(transactionId);
+        builder.setType(NsdEventType.NET_REGISTER);
+        builder.setQueryResult(MdnsQueryResult.MQR_SERVICE_REGISTRATION_FAILED);
+        builder.setEventDurationMillisec(durationMs);
+        mDependencies.statsWrite(builder.build());
+    }
+
+    /**
+     * Report service unregistration success metric data.
+     *
+     * @param transactionId The transaction id of service registration.
+     * @param durationMs The duration of service stayed registered.
+     */
+    public void reportServiceUnregistration(int transactionId, long durationMs) {
+        final Builder builder = makeReportedBuilder();
+        builder.setTransactionId(transactionId);
+        builder.setType(NsdEventType.NET_REGISTER);
+        builder.setQueryResult(MdnsQueryResult.MQR_SERVICE_UNREGISTERED);
+        builder.setEventDurationMillisec(durationMs);
+        // TODO: Report repliedRequestsCount
+        mDependencies.statsWrite(builder.build());
+    }
 }
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 95717c2..745c5bc 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -162,6 +162,7 @@
     public static final boolean DBG = Log.isLoggable(TAG, Log.DEBUG);
     private static final long CLEANUP_DELAY_MS = 10000;
     private static final int IFACE_IDX_ANY = 0;
+    private static final int NO_TRANSACTION = -1;
     private static final SharedLog LOGGER = new SharedLog("serviceDiscovery");
 
     private final Context mContext;
@@ -178,6 +179,8 @@
     private final MdnsSocketProvider mMdnsSocketProvider;
     @NonNull
     private final MdnsAdvertiser mAdvertiser;
+    @NonNull
+    private final Clock mClock;
     private final SharedLog mServiceLogs = LOGGER.forSubComponent(TAG);
     // WARNING : Accessing these values in any thread is not safe, it must only be changed in the
     // state machine thread. If change this outside state machine, it will need to introduce
@@ -530,8 +533,9 @@
                         try {
                             cb.asBinder().linkToDeath(arg.connector, 0);
                             final String tag = "Client" + arg.uid + "-" + mClientNumberId++;
-                            final NetworkNsdReportedMetrics metrics = new NetworkNsdReportedMetrics(
-                                    !arg.useJavaBackend, (int) new Clock().elapsedRealtime());
+                            final NetworkNsdReportedMetrics metrics =
+                                    mDeps.makeNetworkNsdReportedMetrics(
+                                            !arg.useJavaBackend, (int) mClock.elapsedRealtime());
                             cInfo = new ClientInfo(cb, arg.uid, arg.useJavaBackend,
                                     mServiceLogs.forSubComponent(tag), metrics);
                             mClients.put(arg.connector, cInfo);
@@ -569,7 +573,7 @@
                     case NsdManager.REGISTER_SERVICE:
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
-                            cInfo.onRegisterServiceFailed(
+                            cInfo.onRegisterServiceFailedImmediately(
                                     clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
                         break;
@@ -651,8 +655,8 @@
 
             private void storeLegacyRequestMap(int clientRequestId, int transactionId,
                     ClientInfo clientInfo, int what) {
-                clientInfo.mClientRequests.put(
-                        clientRequestId, new LegacyClientRequest(transactionId, what));
+                clientInfo.mClientRequests.put(clientRequestId, new LegacyClientRequest(
+                        transactionId, what, mClock, mClock.elapsedRealtime()));
                 mTransactionIdToClientInfoMap.put(transactionId, clientInfo);
                 // Remove the cleanup event because here comes a new request.
                 cancelStop();
@@ -660,8 +664,8 @@
 
             private void storeAdvertiserRequestMap(int clientRequestId, int transactionId,
                     ClientInfo clientInfo, @Nullable Network requestedNetwork) {
-                clientInfo.mClientRequests.put(clientRequestId,
-                        new AdvertiserClientRequest(transactionId, requestedNetwork));
+                clientInfo.mClientRequests.put(clientRequestId, new AdvertiserClientRequest(
+                        transactionId, requestedNetwork, mClock, mClock.elapsedRealtime()));
                 mTransactionIdToClientInfoMap.put(transactionId, clientInfo);
                 updateMulticastLock();
             }
@@ -684,8 +688,9 @@
             private void storeDiscoveryManagerRequestMap(int clientRequestId, int transactionId,
                     MdnsListener listener, ClientInfo clientInfo,
                     @Nullable Network requestedNetwork) {
-                clientInfo.mClientRequests.put(clientRequestId,
-                        new DiscoveryManagerRequest(transactionId, listener, requestedNetwork));
+                clientInfo.mClientRequests.put(clientRequestId, new DiscoveryManagerRequest(
+                        transactionId, listener, requestedNetwork, mClock,
+                        mClock.elapsedRealtime()));
                 mTransactionIdToClientInfoMap.put(transactionId, clientInfo);
                 updateMulticastLock();
             }
@@ -838,7 +843,7 @@
                         }
 
                         if (requestLimitReached(clientInfo)) {
-                            clientInfo.onRegisterServiceFailed(
+                            clientInfo.onRegisterServiceFailedImmediately(
                                     clientRequestId, NsdManager.FAILURE_MAX_LIMIT);
                             break;
                         }
@@ -854,8 +859,8 @@
                                 || useAdvertiserForType(registerServiceType)) {
                             if (registerServiceType == null) {
                                 Log.e(TAG, "Invalid service type: " + serviceType);
-                                clientInfo.onRegisterServiceFailed(clientRequestId,
-                                        NsdManager.FAILURE_INTERNAL_ERROR);
+                                clientInfo.onRegisterServiceFailedImmediately(
+                                        clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                                 break;
                             }
                             serviceInfo.setServiceType(registerServiceType);
@@ -882,7 +887,7 @@
                                 // Return success after mDns reports success
                             } else {
                                 unregisterService(transactionId);
-                                clientInfo.onRegisterServiceFailed(
+                                clientInfo.onRegisterServiceFailedImmediately(
                                         clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                             }
 
@@ -914,10 +919,12 @@
                         // instead of looking at the flag value.
                         if (request instanceof AdvertiserClientRequest) {
                             mAdvertiser.removeService(transactionId);
-                            clientInfo.onUnregisterServiceSucceeded(clientRequestId);
+                            clientInfo.onUnregisterServiceSucceeded(clientRequestId, transactionId,
+                                    request.calculateRequestDurationMs());
                         } else {
                             if (unregisterService(transactionId)) {
-                                clientInfo.onUnregisterServiceSucceeded(clientRequestId);
+                                clientInfo.onUnregisterServiceSucceeded(clientRequestId,
+                                        transactionId, request.calculateRequestDurationMs());
                             } else {
                                 clientInfo.onUnregisterServiceFailed(
                                         clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
@@ -1181,12 +1188,18 @@
                         final RegistrationInfo info = (RegistrationInfo) obj;
                         final String name = info.serviceName;
                         servInfo = new NsdServiceInfo(name, null /* serviceType */);
-                        clientInfo.onRegisterServiceSucceeded(clientRequestId, servInfo);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
+                        clientInfo.onRegisterServiceSucceeded(clientRequestId, servInfo,
+                                transactionId, request.calculateRequestDurationMs());
                         break;
                     }
                     case IMDnsEventListener.SERVICE_REGISTRATION_FAILED:
-                        clientInfo.onRegisterServiceFailed(
-                                clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
+                        clientInfo.onRegisterServiceFailed(clientRequestId,
+                                NsdManager.FAILURE_INTERNAL_ERROR, transactionId,
+                                request.calculateRequestDurationMs());
                         break;
                     case IMDnsEventListener.SERVICE_RESOLVED: {
                         final ResolutionInfo info = (ResolutionInfo) obj;
@@ -1566,6 +1579,7 @@
         handler.post(() -> mMdnsSocketClient.setCallback(mMdnsDiscoveryManager));
         mAdvertiser = deps.makeMdnsAdvertiser(handler.getLooper(), mMdnsSocketProvider,
                 new AdvertiserCallback(), LOGGER.forSubComponent("MdnsAdvertiser"));
+        mClock = deps.makeClock();
     }
 
     /**
@@ -1655,6 +1669,21 @@
         public int getCallingUid() {
             return Binder.getCallingUid();
         }
+
+        /**
+         * @see NetworkNsdReportedMetrics
+         */
+        public NetworkNsdReportedMetrics makeNetworkNsdReportedMetrics(
+                boolean isLegacy, int clientId) {
+            return new NetworkNsdReportedMetrics(isLegacy, clientId);
+        }
+
+        /**
+         * @see MdnsUtils.Clock
+         */
+        public Clock makeClock() {
+            return new Clock();
+        }
     }
 
     /**
@@ -1755,7 +1784,9 @@
             // onRegisterServiceSucceeded only has the service name in its info. This aligns with
             // historical behavior.
             final NsdServiceInfo cbInfo = new NsdServiceInfo(registeredInfo.getServiceName(), null);
-            clientInfo.onRegisterServiceSucceeded(clientRequestId, cbInfo);
+            final ClientRequest request = clientInfo.mClientRequests.get(clientRequestId);
+            clientInfo.onRegisterServiceSucceeded(
+                    clientRequestId, cbInfo, transactionId, request.calculateRequestDurationMs());
         }
 
         @Override
@@ -1765,8 +1796,9 @@
 
             final int clientRequestId = getClientRequestIdOrLog(clientInfo, transactionId);
             if (clientRequestId < 0) return;
-
-            clientInfo.onRegisterServiceFailed(clientRequestId, errorCode);
+            final ClientRequest request = clientInfo.mClientRequests.get(clientRequestId);
+            clientInfo.onRegisterServiceFailed(clientRequestId, errorCode, transactionId,
+                    request.calculateRequestDurationMs());
         }
 
         private ClientInfo getClientInfoOrLog(int transactionId) {
@@ -2026,17 +2058,27 @@
 
     private abstract static class ClientRequest {
         private final int mTransactionId;
+        private final Clock mClock;
+        private final long mStartTimeMs;
 
-        private ClientRequest(int transactionId) {
+        private ClientRequest(int transactionId, @NonNull Clock clock, long startTimeMs) {
             mTransactionId = transactionId;
+            mClock = clock;
+            mStartTimeMs = startTimeMs;
+        }
+
+        public long calculateRequestDurationMs() {
+            final long stopTimeMs = mClock.elapsedRealtime();
+            return stopTimeMs - mStartTimeMs;
         }
     }
 
     private static class LegacyClientRequest extends ClientRequest {
         private final int mRequestCode;
 
-        private LegacyClientRequest(int transactionId, int requestCode) {
-            super(transactionId);
+        private LegacyClientRequest(int transactionId, int requestCode, @NonNull Clock clock,
+                long startTimeMs) {
+            super(transactionId, clock, startTimeMs);
             mRequestCode = requestCode;
         }
     }
@@ -2045,8 +2087,9 @@
         @Nullable
         private final Network mRequestedNetwork;
 
-        private JavaBackendClientRequest(int transactionId, @Nullable Network requestedNetwork) {
-            super(transactionId);
+        private JavaBackendClientRequest(int transactionId, @Nullable Network requestedNetwork,
+                @NonNull Clock clock, long startTimeMs) {
+            super(transactionId, clock, startTimeMs);
             mRequestedNetwork = requestedNetwork;
         }
 
@@ -2057,8 +2100,9 @@
     }
 
     private static class AdvertiserClientRequest extends JavaBackendClientRequest {
-        private AdvertiserClientRequest(int transactionId, @Nullable Network requestedNetwork) {
-            super(transactionId, requestedNetwork);
+        private AdvertiserClientRequest(int transactionId, @Nullable Network requestedNetwork,
+                @NonNull Clock clock, long startTimeMs) {
+            super(transactionId, requestedNetwork, clock, startTimeMs);
         }
     }
 
@@ -2067,8 +2111,8 @@
         private final MdnsListener mListener;
 
         private DiscoveryManagerRequest(int transactionId, @NonNull MdnsListener listener,
-                @Nullable Network requestedNetwork) {
-            super(transactionId, requestedNetwork);
+                @Nullable Network requestedNetwork, @NonNull Clock clock, long startTimeMs) {
+            super(transactionId, requestedNetwork, clock, startTimeMs);
             mListener = listener;
         }
     }
@@ -2161,6 +2205,8 @@
 
                 if (request instanceof AdvertiserClientRequest) {
                     mAdvertiser.removeService(transactionId);
+                    mMetrics.reportServiceUnregistration(
+                            transactionId, request.calculateRequestDurationMs());
                     continue;
                 }
 
@@ -2177,6 +2223,8 @@
                         break;
                     case NsdManager.REGISTER_SERVICE:
                         unregisterService(transactionId);
+                        mMetrics.reportServiceUnregistration(
+                                transactionId, request.calculateRequestDurationMs());
                         break;
                     default:
                         break;
@@ -2268,7 +2316,13 @@
             }
         }
 
-        void onRegisterServiceFailed(int listenerKey, int error) {
+        void onRegisterServiceFailedImmediately(int listenerKey, int error) {
+            onRegisterServiceFailed(listenerKey, error, NO_TRANSACTION, 0 /* durationMs */);
+        }
+
+        void onRegisterServiceFailed(int listenerKey, int error, int transactionId,
+                long durationMs) {
+            mMetrics.reportServiceRegistrationFailed(transactionId, durationMs);
             try {
                 mCb.onRegisterServiceFailed(listenerKey, error);
             } catch (RemoteException e) {
@@ -2276,7 +2330,9 @@
             }
         }
 
-        void onRegisterServiceSucceeded(int listenerKey, NsdServiceInfo info) {
+        void onRegisterServiceSucceeded(int listenerKey, NsdServiceInfo info, int transactionId,
+                long durationMs) {
+            mMetrics.reportServiceRegistrationSucceeded(transactionId, durationMs);
             try {
                 mCb.onRegisterServiceSucceeded(listenerKey, info);
             } catch (RemoteException e) {
@@ -2292,7 +2348,8 @@
             }
         }
 
-        void onUnregisterServiceSucceeded(int listenerKey) {
+        void onUnregisterServiceSucceeded(int listenerKey, int transactionId, long durationMs) {
+            mMetrics.reportServiceUnregistration(transactionId, durationMs);
             try {
                 mCb.onUnregisterServiceSucceeded(listenerKey);
             } catch (RemoteException e) {
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
index cb01732..77383ad 100644
--- a/tests/benchmark/Android.bp
+++ b/tests/benchmark/Android.bp
@@ -31,6 +31,7 @@
     ],
     static_libs: [
         "androidx.test.rules",
+        "mockito-target-minus-junit4",
         "net-tests-utils",
         "service-connectivity-pre-jarjar",
         "service-connectivity-tiramisu-pre-jarjar",
diff --git a/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt
deleted file mode 100644
index 177014f..0000000
--- a/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt
+++ /dev/null
@@ -1,87 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.server.net.benchmarktests
-
-import android.net.NetworkStatsCollection
-import androidx.test.InstrumentationRegistry
-import com.android.internal.util.FileRotator.Reader
-import com.android.server.connectivity.benchmarktests.R
-import java.io.BufferedInputStream
-import java.io.DataInputStream
-import java.io.File
-import java.io.FileInputStream
-import java.io.FileOutputStream
-import java.nio.file.Files
-import java.util.concurrent.TimeUnit
-import java.util.zip.ZipInputStream
-import kotlin.test.assertTrue
-import org.junit.Test
-import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
-
-@RunWith(JUnit4::class)
-class NetworkStatsCollectionTest {
-    private val DEFAULT_BUFFER_SIZE = 8192
-    private val UID_COLLECTION_BUCKET_DURATION_MS = TimeUnit.HOURS.toMillis(2)
-
-    private val uidTestFiles: List<File> by lazy {
-        // These file generated by using real user dataset which has many uid records and agreed to
-        // share the dataset for testing purpose. These dataset can be extracted from rooted
-        // devices by using "adb pull /data/misc/apexdata/com.android.tethering/netstats" command.
-        val zipInputStream = ZipInputStream(getInputStreamForResource(R.raw.netstats_many_uids_zip))
-        getSortedListForPrefix(unzipToTempDir(zipInputStream), "uid")
-    }
-
-    @Test
-    fun testReadCollection_manyUids() {
-        val collection = NetworkStatsCollection(UID_COLLECTION_BUCKET_DURATION_MS)
-        for (file in uidTestFiles) {
-            readFile(file, collection)
-        }
-    }
-
-    private fun getInputStreamForResource(resourceId: Int): DataInputStream {
-        return DataInputStream(
-            InstrumentationRegistry.getContext()
-                .getResources().openRawResource(resourceId)
-        )
-    }
-
-    private fun unzipToTempDir(zis: ZipInputStream): File {
-        val statsDir =
-            Files.createTempDirectory(NetworkStatsCollectionTest::class.simpleName).toFile()
-        while (true) {
-            val entryName = zis.nextEntry?.name ?: break
-            val file = File(statsDir, entryName)
-            FileOutputStream(file).use { zis.copyTo(it, DEFAULT_BUFFER_SIZE) }
-        }
-        return statsDir
-    }
-
-    // List [xt|uid|uid_tag].<start>-<end> files under the given directory.
-    private fun getSortedListForPrefix(statsDir: File, prefix: String): List<File> {
-        assertTrue(statsDir.exists())
-        return (statsDir.list() ?: arrayOf()).mapNotNull {
-            if (it.startsWith("$prefix.")) File(statsDir, it) else null
-        }.sorted()
-    }
-
-    private fun readFile(file: File, reader: Reader) =
-        BufferedInputStream(FileInputStream(file)).use {
-            reader.read(it)
-        }
-}
diff --git a/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsTest.kt b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsTest.kt
new file mode 100644
index 0000000..8492268
--- /dev/null
+++ b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsTest.kt
@@ -0,0 +1,142 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.net.benchmarktests
+
+import android.net.NetworkStats.NonMonotonicObserver
+import android.net.NetworkStatsCollection
+import android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID
+import android.os.DropBoxManager
+import androidx.test.InstrumentationRegistry
+import com.android.internal.util.FileRotator
+import com.android.internal.util.FileRotator.Reader
+import com.android.server.connectivity.benchmarktests.R
+import com.android.server.net.NetworkStatsRecorder
+import java.io.BufferedInputStream
+import java.io.DataInputStream
+import java.io.File
+import java.io.FileInputStream
+import java.io.FileOutputStream
+import java.nio.file.Files
+import java.util.concurrent.TimeUnit
+import java.util.zip.ZipInputStream
+import kotlin.test.assertTrue
+import org.junit.BeforeClass
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.mockito.Mockito.mock
+
+@RunWith(JUnit4::class)
+class NetworkStatsTest {
+    companion object {
+        private val DEFAULT_BUFFER_SIZE = 8192
+        private val UID_COLLECTION_BUCKET_DURATION_MS = TimeUnit.HOURS.toMillis(2)
+        private val UID_RECORDER_ROTATE_AGE_MS = TimeUnit.DAYS.toMillis(15)
+        private val UID_RECORDER_DELETE_AGE_MS = TimeUnit.DAYS.toMillis(90)
+
+        private val testFilesDir by lazy {
+            // These file generated by using real user dataset which has many uid records
+            // and agreed to share the dataset for testing purpose. These dataset can be
+            // extracted from rooted devices by using
+            // "adb pull /data/misc/apexdata/com.android.tethering/netstats" command.
+            val zipInputStream =
+                ZipInputStream(getInputStreamForResource(R.raw.netstats_many_uids_zip))
+            unzipToTempDir(zipInputStream)
+        }
+
+        private val uidTestFiles: List<File> by lazy {
+            getSortedListForPrefix(testFilesDir, "uid")
+        }
+
+        // Test results shows the test cases who read the file first will take longer time to
+        // execute, and reading time getting shorter each time. Read files several times prior to
+        // tests to minimize the impact. This cannot live in setUp() since the time
+        // spent on the file reading will be attributed to the time spent on the individual
+        // test case.
+        @JvmStatic
+        @BeforeClass
+        fun setUpOnce() {
+            for (i in 1..10) {
+                val collection = NetworkStatsCollection(UID_COLLECTION_BUCKET_DURATION_MS)
+                for (file in uidTestFiles) {
+                    readFile(file, collection)
+                }
+            }
+        }
+
+        private fun getInputStreamForResource(resourceId: Int): DataInputStream {
+            return DataInputStream(
+                InstrumentationRegistry.getContext()
+                    .getResources().openRawResource(resourceId)
+            )
+        }
+
+        private fun unzipToTempDir(zis: ZipInputStream): File {
+            val statsDir =
+                Files.createTempDirectory(NetworkStatsTest::class.simpleName).toFile()
+            while (true) {
+                val entryName = zis.nextEntry?.name ?: break
+                val file = File(statsDir, entryName)
+                FileOutputStream(file).use { zis.copyTo(it, DEFAULT_BUFFER_SIZE) }
+            }
+            return statsDir
+        }
+
+        // List [xt|uid|uid_tag].<start>-<end> files under the given directory.
+        private fun getSortedListForPrefix(statsDir: File, prefix: String): List<File> {
+            assertTrue(statsDir.exists())
+            return (statsDir.list() ?: arrayOf()).mapNotNull {
+                if (it.startsWith("$prefix.")) File(statsDir, it) else null
+            }.sorted()
+        }
+
+        private fun readFile(file: File, reader: Reader) =
+            BufferedInputStream(FileInputStream(file)).use {
+                reader.read(it)
+            }
+    }
+
+    @Test
+    fun testReadCollection_manyUids() {
+        for (i in 1..10) {
+            val collection = NetworkStatsCollection(UID_COLLECTION_BUCKET_DURATION_MS)
+            for (file in uidTestFiles) {
+                readFile(file, collection)
+            }
+        }
+    }
+
+    @Test
+    fun testReadFromRecorder_manyUids() {
+        for (i in 1..10) {
+            val recorder = NetworkStatsRecorder(
+                FileRotator(
+                    testFilesDir, PREFIX_UID, UID_RECORDER_ROTATE_AGE_MS, UID_RECORDER_DELETE_AGE_MS
+                ),
+                mock<NonMonotonicObserver<String>>(),
+                mock(DropBoxManager::class.java),
+                PREFIX_UID,
+                UID_COLLECTION_BUCKET_DURATION_MS,
+                false /* includeTags */,
+                false /* wipeOnError */
+            )
+            recorder.orLoadCompleteLocked
+        }
+    }
+
+    inline fun <reified T : Any> mock(): T = mock(T::class.java)
+}
diff --git a/tests/cts/net/src/android/net/cts/DnsResolverTest.java b/tests/cts/net/src/android/net/cts/DnsResolverTest.java
index 3821cea..308aead 100644
--- a/tests/cts/net/src/android/net/cts/DnsResolverTest.java
+++ b/tests/cts/net/src/android/net/cts/DnsResolverTest.java
@@ -59,6 +59,7 @@
 import com.android.net.module.util.DnsPacket;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
+import com.android.testutils.DnsResolverModuleTest;
 import com.android.testutils.SkipPresubmit;
 
 import org.junit.After;
@@ -317,51 +318,61 @@
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQuery() throws Exception {
         doTestRawQuery(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryInline() throws Exception {
         doTestRawQuery(mExecutorInline);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryBlob() throws Exception {
         doTestRawQueryBlob(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryBlobInline() throws Exception {
         doTestRawQueryBlob(mExecutorInline);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryRoot() throws Exception {
         doTestRawQueryRoot(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryRootInline() throws Exception {
         doTestRawQueryRoot(mExecutorInline);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryNXDomain() throws Exception {
         doTestRawQueryNXDomain(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryNXDomainInline() throws Exception {
         doTestRawQueryNXDomain(mExecutorInline);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryNXDomainWithPrivateDns() throws Exception {
         doTestRawQueryNXDomainWithPrivateDns(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testRawQueryNXDomainInlineWithPrivateDns() throws Exception {
         doTestRawQueryNXDomainWithPrivateDns(mExecutorInline);
     }
@@ -610,41 +621,49 @@
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testQueryForInetAddress() throws Exception {
         doTestQueryForInetAddress(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testQueryForInetAddressInline() throws Exception {
         doTestQueryForInetAddress(mExecutorInline);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testQueryForInetAddressIpv4() throws Exception {
         doTestQueryForInetAddressIpv4(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testQueryForInetAddressIpv4Inline() throws Exception {
         doTestQueryForInetAddressIpv4(mExecutorInline);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testQueryForInetAddressIpv6() throws Exception {
         doTestQueryForInetAddressIpv6(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testQueryForInetAddressIpv6Inline() throws Exception {
         doTestQueryForInetAddressIpv6(mExecutorInline);
     }
 
     @Test
+    @DnsResolverModuleTest
     public void testContinuousQueries() throws Exception {
         doTestContinuousQueries(mExecutor);
     }
 
     @Test
+    @DnsResolverModuleTest
     @SkipPresubmit(reason = "Flaky: b/159762682; add to presubmit after fixing")
     public void testContinuousQueriesInline() throws Exception {
         doTestContinuousQueries(mExecutorInline);
diff --git a/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt b/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
new file mode 100644
index 0000000..961c422
--- /dev/null
+++ b/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License")
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.metrics
+
+import android.os.Build
+import android.stats.connectivity.MdnsQueryResult
+import android.stats.connectivity.NsdEventType
+import com.android.testutils.DevSdkIgnoreRule
+import com.android.testutils.DevSdkIgnoreRunner
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.verify
+
+@RunWith(DevSdkIgnoreRunner::class)
+@DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.TIRAMISU)
+class NetworkNsdReportedMetricsTest {
+    private val deps = mock(NetworkNsdReportedMetrics.Dependencies::class.java)
+
+    @Test
+    fun testReportServiceRegistrationSucceeded() {
+        val clientId = 99
+        val transactionId = 100
+        val durationMs = 10L
+        val metrics = NetworkNsdReportedMetrics(true /* isLegacy */, clientId, deps)
+        metrics.reportServiceRegistrationSucceeded(transactionId, durationMs)
+
+        val eventCaptor = ArgumentCaptor.forClass(NetworkNsdReported::class.java)
+        verify(deps).statsWrite(eventCaptor.capture())
+        eventCaptor.value.let {
+            assertTrue(it.isLegacy)
+            assertEquals(clientId, it.clientId)
+            assertEquals(transactionId, it.transactionId)
+            assertEquals(NsdEventType.NET_REGISTER, it.type)
+            assertEquals(MdnsQueryResult.MQR_SERVICE_REGISTERED, it.queryResult)
+            assertEquals(durationMs, it.eventDurationMillisec)
+        }
+    }
+
+    @Test
+    fun testReportServiceRegistrationFailed() {
+        val clientId = 99
+        val transactionId = 100
+        val durationMs = 10L
+        val metrics = NetworkNsdReportedMetrics(false /* isLegacy */, clientId, deps)
+        metrics.reportServiceRegistrationFailed(transactionId, durationMs)
+
+        val eventCaptor = ArgumentCaptor.forClass(NetworkNsdReported::class.java)
+        verify(deps).statsWrite(eventCaptor.capture())
+        eventCaptor.value.let {
+            assertFalse(it.isLegacy)
+            assertEquals(clientId, it.clientId)
+            assertEquals(transactionId, it.transactionId)
+            assertEquals(NsdEventType.NET_REGISTER, it.type)
+            assertEquals(MdnsQueryResult.MQR_SERVICE_REGISTRATION_FAILED, it.queryResult)
+            assertEquals(durationMs, it.eventDurationMillisec)
+        }
+    }
+
+    @Test
+    fun testReportServiceUnregistration() {
+        val clientId = 99
+        val transactionId = 100
+        val durationMs = 10L
+        val metrics = NetworkNsdReportedMetrics(true /* isLegacy */, clientId, deps)
+        metrics.reportServiceUnregistration(transactionId, durationMs)
+
+        val eventCaptor = ArgumentCaptor.forClass(NetworkNsdReported::class.java)
+        verify(deps).statsWrite(eventCaptor.capture())
+        eventCaptor.value.let {
+            assertTrue(it.isLegacy)
+            assertEquals(clientId, it.clientId)
+            assertEquals(transactionId, it.transactionId)
+            assertEquals(NsdEventType.NET_REGISTER, it.type)
+            assertEquals(MdnsQueryResult.MQR_SERVICE_UNREGISTERED, it.queryResult)
+            assertEquals(durationMs, it.eventDurationMillisec)
+        }
+    }
+}
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index f51b28d..55384b3 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -43,7 +43,9 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.any;
@@ -95,6 +97,7 @@
 import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 
+import com.android.metrics.NetworkNsdReportedMetrics;
 import com.android.server.NsdService.Dependencies;
 import com.android.server.connectivity.mdns.MdnsAdvertiser;
 import com.android.server.connectivity.mdns.MdnsDiscoveryManager;
@@ -104,6 +107,7 @@
 import com.android.server.connectivity.mdns.MdnsServiceInfo;
 import com.android.server.connectivity.mdns.MdnsSocketProvider;
 import com.android.server.connectivity.mdns.MdnsSocketProvider.SocketRequestMonitor;
+import com.android.server.connectivity.mdns.util.MdnsUtils;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
 import com.android.testutils.HandlerUtils;
@@ -138,6 +142,7 @@
     static final int PROTOCOL = NsdManager.PROTOCOL_DNS_SD;
     private static final long CLEANUP_DELAY_MS = 500;
     private static final long TIMEOUT_MS = 500;
+    private static final long TEST_TIME_MS = 123L;
     private static final String SERVICE_NAME = "a_name";
     private static final String SERVICE_TYPE = "_test._tcp";
     private static final String SERVICE_FULL_NAME = SERVICE_NAME + "." + SERVICE_TYPE;
@@ -164,6 +169,8 @@
     @Mock WifiManager mWifiManager;
     @Mock WifiManager.MulticastLock mMulticastLock;
     @Mock ActivityManager mActivityManager;
+    @Mock NetworkNsdReportedMetrics mMetrics;
+    @Mock MdnsUtils.Clock mClock;
     SocketRequestMonitor mSocketRequestMonitor;
     OnUidImportanceListener mUidImportanceListener;
     HandlerThread mThread;
@@ -210,6 +217,9 @@
         doReturn(DEFAULT_RUNNING_APP_ACTIVE_IMPORTANCE_CUTOFF).when(mDeps).getDeviceConfigInt(
                 eq(NsdService.MDNS_CONFIG_RUNNING_APP_ACTIVE_IMPORTANCE_CUTOFF), anyInt());
         doReturn(mAdvertiser).when(mDeps).makeMdnsAdvertiser(any(), any(), any(), any());
+        doReturn(mMetrics).when(mDeps).makeNetworkNsdReportedMetrics(anyBoolean(), anyInt());
+        doReturn(mClock).when(mDeps).makeClock();
+        doReturn(TEST_TIME_MS).when(mClock).elapsedRealtime();
         mService = makeService();
         final ArgumentCaptor<SocketRequestMonitor> cbMonitorCaptor =
                 ArgumentCaptor.forClass(SocketRequestMonitor.class);
@@ -512,14 +522,16 @@
                 eq(SERVICE_NAME), eq(SERVICE_TYPE), eq(PORT), any(), eq(IFACE_IDX_ANY));
 
         // Register service successfully.
+        final int regId = regIdCaptor.getValue();
         final RegistrationInfo registrationInfo = new RegistrationInfo(
-                regIdCaptor.getValue(),
+                regId,
                 IMDnsEventListener.SERVICE_REGISTERED,
                 SERVICE_NAME,
                 SERVICE_TYPE,
                 PORT,
                 new byte[0] /* txtRecord */,
                 IFACE_IDX_ANY);
+        doReturn(TEST_TIME_MS + 10L).when(mClock).elapsedRealtime();
         eventListener.onServiceRegistrationStatus(registrationInfo);
 
         final ArgumentCaptor<NsdServiceInfo> registeredInfoCaptor =
@@ -528,19 +540,22 @@
                 .onServiceRegistered(registeredInfoCaptor.capture());
         final NsdServiceInfo registeredInfo = registeredInfoCaptor.getValue();
         assertEquals(SERVICE_NAME, registeredInfo.getServiceName());
+        verify(mMetrics).reportServiceRegistrationSucceeded(regId, 10L /* durationMs */);
 
         // Fail to register service.
         final RegistrationInfo registrationFailedInfo = new RegistrationInfo(
-                regIdCaptor.getValue(),
+                regId,
                 IMDnsEventListener.SERVICE_REGISTRATION_FAILED,
                 null /* serviceName */,
                 null /* registrationType */,
                 0 /* port */,
                 new byte[0] /* txtRecord */,
                 IFACE_IDX_ANY);
+        doReturn(TEST_TIME_MS + 20L).when(mClock).elapsedRealtime();
         eventListener.onServiceRegistrationStatus(registrationFailedInfo);
         verify(regListener, timeout(TIMEOUT_MS))
                 .onRegistrationFailed(any(), eq(FAILURE_INTERNAL_ERROR));
+        verify(mMetrics).reportServiceRegistrationFailed(regId, 20L /* durationMs */);
     }
 
     @Test
@@ -1215,17 +1230,22 @@
 
         // Verify onServiceRegistered callback
         final MdnsAdvertiser.AdvertiserCallback cb = cbCaptor.getValue();
-        cb.onRegisterServiceSucceeded(idCaptor.getValue(), regInfo);
+        final int regId = idCaptor.getValue();
+        doReturn(TEST_TIME_MS + 10L).when(mClock).elapsedRealtime();
+        cb.onRegisterServiceSucceeded(regId, regInfo);
 
         verify(regListener, timeout(TIMEOUT_MS)).onServiceRegistered(argThat(info -> matches(info,
                 new NsdServiceInfo(regInfo.getServiceName(), null))));
+        verify(mMetrics).reportServiceRegistrationSucceeded(regId, 10L /* durationMs */);
 
+        doReturn(TEST_TIME_MS + 100L).when(mClock).elapsedRealtime();
         client.unregisterService(regListener);
         waitForIdle();
         verify(mAdvertiser).removeService(idCaptor.getValue());
         verify(regListener, timeout(TIMEOUT_MS)).onServiceUnregistered(
                 argThat(info -> matches(info, regInfo)));
         verify(mSocketProvider, timeout(TIMEOUT_MS)).requestStopWhenInactive();
+        verify(mMetrics).reportServiceUnregistration(regId, 100L /* durationMs */);
     }
 
     @Test
@@ -1251,6 +1271,7 @@
 
         verify(regListener, timeout(TIMEOUT_MS)).onRegistrationFailed(
                 argThat(info -> matches(info, regInfo)), eq(FAILURE_INTERNAL_ERROR));
+        verify(mMetrics).reportServiceRegistrationFailed(anyInt(), anyLong());
     }
 
     @Test
@@ -1280,10 +1301,13 @@
 
         // Verify onServiceRegistered callback
         final MdnsAdvertiser.AdvertiserCallback cb = cbCaptor.getValue();
-        cb.onRegisterServiceSucceeded(idCaptor.getValue(), regInfo);
+        final int regId = idCaptor.getValue();
+        doReturn(TEST_TIME_MS + 10L).when(mClock).elapsedRealtime();
+        cb.onRegisterServiceSucceeded(regId, regInfo);
 
         verify(regListener, timeout(TIMEOUT_MS)).onServiceRegistered(
                 argThat(info -> matches(info, new NsdServiceInfo(regInfo.getServiceName(), null))));
+        verify(mMetrics).reportServiceRegistrationSucceeded(regId, 10L /* durationMs */);
     }
 
     @Test