Merge "Move QueryTaskConfig to a standalone class" into main
diff --git a/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java b/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
index 0eebc61..161669b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
+++ b/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
@@ -16,6 +16,7 @@
 
 package com.android.server.connectivity.mdns;
 
+import android.annotation.NonNull;
 import android.util.ArraySet;
 
 import java.util.Set;
@@ -47,5 +48,17 @@
             }
             executor.shutdownNow();
         }
+        serviceTypeClientSchedulerExecutors.clear();
+    }
+
+    /**
+     * Shutdown one executor service and remove the executor service from the set.
+     * @param executorService the executorService to be shutdown
+     */
+    public void shutdownExecutorService(@NonNull ScheduledExecutorService executorService) {
+        if (!executorService.isShutdown()) {
+            executorService.shutdownNow();
+        }
+        serviceTypeClientSchedulerExecutors.remove(executorService);
     }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
index f386dd4..d55098c 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
@@ -204,6 +204,7 @@
                         if (serviceTypeClient == null) return;
                         // Notify all listeners that all services are removed from this socket.
                         serviceTypeClient.notifySocketDestroyed();
+                        executorProvider.shutdownExecutorService(serviceTypeClient.getExecutor());
                         perSocketServiceTypeClients.remove(serviceTypeClient);
                     }
                 });
@@ -238,6 +239,7 @@
             if (serviceTypeClient.stopSendAndReceive(listener)) {
                 // No listener is registered for the service type anymore, remove it from the list
                 // of the service type clients.
+                executorProvider.shutdownExecutorService(serviceTypeClient.getExecutor());
                 perSocketServiceTypeClients.remove(serviceTypeClient);
             }
         }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
index b2e6fa5..b5fd8a0 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -329,6 +329,13 @@
         }
     }
 
+    /**
+     * Get the executor service.
+     */
+    public ScheduledExecutorService getExecutor() {
+        return executor;
+    }
+
     private void removeScheduledTask() {
         dependencies.removeMessages(handler, EVENT_START_QUERYTASK);
         sharedLog.log("Remove EVENT_START_QUERYTASK"
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
new file mode 100644
index 0000000..cb01732
--- /dev/null
+++ b/tests/benchmark/Android.bp
@@ -0,0 +1,41 @@
+//
+// 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 {
+    // See: http://go/android-license-faq
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_test {
+    name: "ConnectivityBenchmarkTests",
+    defaults: [
+        "framework-connectivity-internal-test-defaults",
+    ],
+    platform_apis: true,
+    srcs: [
+        "src/**/*.kt",
+        "src/**/*.aidl",
+    ],
+    static_libs: [
+        "androidx.test.rules",
+        "net-tests-utils",
+        "service-connectivity-pre-jarjar",
+        "service-connectivity-tiramisu-pre-jarjar",
+    ],
+    test_suites: ["device-tests"],
+    jarjar_rules: ":connectivity-jarjar-rules",
+}
+
diff --git a/tests/benchmark/AndroidManifest.xml b/tests/benchmark/AndroidManifest.xml
new file mode 100644
index 0000000..bd2fce5
--- /dev/null
+++ b/tests/benchmark/AndroidManifest.xml
@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- 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.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+        package="com.android.server.connectivity.benchmarktests">
+    <application>
+        <uses-library android:name="android.test.runner" />
+    </application>
+
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+         android:targetPackage="com.android.server.connectivity.benchmarktests"
+         android:label="Connectivity Benchmark Tests" />
+</manifest>
diff --git a/tests/benchmark/OWNERS b/tests/benchmark/OWNERS
new file mode 100644
index 0000000..3101da5
--- /dev/null
+++ b/tests/benchmark/OWNERS
@@ -0,0 +1,2 @@
+# Bug template url: http://b/new?component=31808
+# TODO: move bug template config to common owners file once b/226427845 is resolved
\ No newline at end of file
diff --git a/tests/benchmark/res/raw/netstats-many-uids-zip b/tests/benchmark/res/raw/netstats-many-uids-zip
new file mode 100644
index 0000000..22e8254
--- /dev/null
+++ b/tests/benchmark/res/raw/netstats-many-uids-zip
Binary files differ
diff --git a/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt
new file mode 100644
index 0000000..177014f
--- /dev/null
+++ b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt
@@ -0,0 +1,87 @@
+/*
+ * 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/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
index 1a4ae5d..e869b91 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
@@ -53,6 +53,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
 
 /** Tests for {@link MdnsDiscoveryManager}. */
 @RunWith(DevSdkIgnoreRunner.class)
@@ -80,6 +81,7 @@
     private static final Pair<String, SocketKey> PER_SOCKET_SERVICE_TYPE_2_NETWORK_2 =
             Pair.create(SERVICE_TYPE_2, SOCKET_KEY_NETWORK_2);
     @Mock private ExecutorProvider executorProvider;
+    @Mock private ScheduledExecutorService mockExecutorService;
     @Mock private MdnsSocketClientBase socketClient;
     @Mock private MdnsServiceTypeClient mockServiceTypeClientType1NullNetwork;
     @Mock private MdnsServiceTypeClient mockServiceTypeClientType1Network1;
@@ -128,6 +130,7 @@
                         return null;
                     }
                 };
+        doReturn(mockExecutorService).when(mockServiceTypeClientType1NullNetwork).getExecutor();
     }
 
     @After
@@ -165,11 +168,25 @@
         when(mockServiceTypeClientType1NullNetwork.stopSendAndReceive(mockListenerOne))
                 .thenReturn(true);
         runOnHandler(() -> discoveryManager.unregisterListener(SERVICE_TYPE_1, mockListenerOne));
+        verify(executorProvider).shutdownExecutorService(mockExecutorService);
         verify(mockServiceTypeClientType1NullNetwork).stopSendAndReceive(mockListenerOne);
         verify(socketClient).stopDiscovery();
     }
 
     @Test
+    public void onSocketDestroy_shutdownExecutorService() throws IOException {
+        final MdnsSearchOptions options =
+                MdnsSearchOptions.newBuilder().setNetwork(null /* network */).build();
+        final SocketCreationCallback callback = expectSocketCreationCallback(
+                SERVICE_TYPE_1, mockListenerOne, options);
+        runOnHandler(() -> callback.onSocketCreated(SOCKET_KEY_NULL_NETWORK));
+        verify(mockServiceTypeClientType1NullNetwork).startSendAndReceive(mockListenerOne, options);
+
+        runOnHandler(() -> callback.onSocketDestroyed(SOCKET_KEY_NULL_NETWORK));
+        verify(executorProvider).shutdownExecutorService(mockExecutorService);
+    }
+
+    @Test
     public void registerMultipleListeners() throws IOException {
         final MdnsSearchOptions options =
                 MdnsSearchOptions.newBuilder().setNetwork(null /* network */).build();