Expose setShouldUseHugepages to @SystemApi

Bug: 383347947
Test: atest MicrodroidTests
Change-Id: I5f525833e78a44ff254b7e0ec5ff570579e068bd
diff --git a/build/Android.bp b/build/Android.bp
index 2b97927..ef70fe4 100644
--- a/build/Android.bp
+++ b/build/Android.bp
@@ -75,3 +75,20 @@
     tools: ["dtc"],
     cmd: "FILES=($(in)) && $(location dtc) -@ -I dts -O dtb $${FILES[-1]} -o $(out)",
 }
+
+// This is a temporary workaround until b/343795511 is implemented.
+aconfig_declarations {
+    name: "avf_aconfig_flags",
+    package: "com.android.system.virtualmachine.flags",
+    container: "com.android.virt",
+    srcs: [
+        "avf_flags.aconfig",
+    ],
+}
+
+java_aconfig_library {
+    name: "avf_aconfig_flags_java",
+    aconfig_declarations: "avf_aconfig_flags",
+    sdk_version: "module_current",
+    apex_available: ["com.android.virt"],
+}
diff --git a/build/avf_flags.aconfig b/build/avf_flags.aconfig
new file mode 100644
index 0000000..9815c60
--- /dev/null
+++ b/build/avf_flags.aconfig
@@ -0,0 +1,11 @@
+package: "com.android.system.virtualmachine.flags"
+container: "com.android.virt"
+
+flag {
+  name: "promote_set_should_use_hugepages_to_system_api"
+  is_exported: true
+  namespace: "virtualization"
+  description: "Flag used in @FlaggedApi annotation for setShouldUseHugepages"
+  bug: "383347947"
+  is_fixed_read_only: true
+}
diff --git a/libs/framework-virtualization/Android.bp b/libs/framework-virtualization/Android.bp
index d5ac878..98fa53d 100644
--- a/libs/framework-virtualization/Android.bp
+++ b/libs/framework-virtualization/Android.bp
@@ -15,6 +15,7 @@
     ],
     static_libs: [
         "android.system.virtualizationservice-java",
+        "avf_aconfig_flags_java",
         // For android.sysprop.HypervisorProperties
         "PlatformProperties",
     ],
@@ -51,6 +52,9 @@
             "FlaggedApi",
         ],
     },
+    aconfig_declarations: [
+        "avf_aconfig_flags",
+    ],
 }
 
 gensrcs {
diff --git a/libs/framework-virtualization/api/system-current.txt b/libs/framework-virtualization/api/system-current.txt
index d9bafa1..233491f 100644
--- a/libs/framework-virtualization/api/system-current.txt
+++ b/libs/framework-virtualization/api/system-current.txt
@@ -66,6 +66,7 @@
     method public boolean isEncryptedStorageEnabled();
     method public boolean isProtectedVm();
     method public boolean isVmOutputCaptured();
+    method @FlaggedApi("com.android.system.virtualmachine.flags.promote_set_should_use_hugepages_to_system_api") public boolean shouldUseHugepages();
     field public static final int CPU_TOPOLOGY_MATCH_HOST = 1; // 0x1
     field public static final int CPU_TOPOLOGY_ONE_CPU = 0; // 0x0
     field public static final int DEBUG_LEVEL_FULL = 1; // 0x1
@@ -82,6 +83,7 @@
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setMemoryBytes(@IntRange(from=1) long);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setPayloadBinaryName(@NonNull String);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setProtectedVm(boolean);
+    method @FlaggedApi("com.android.system.virtualmachine.flags.promote_set_should_use_hugepages_to_system_api") @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setShouldUseHugepages(boolean);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setVmOutputCaptured(boolean);
   }
 
diff --git a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachineConfig.java b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachineConfig.java
index 3829f9f..fbc6c66 100644
--- a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -22,6 +22,7 @@
 
 import static java.util.Objects.requireNonNull;
 
+import android.annotation.FlaggedApi;
 import android.annotation.IntDef;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
@@ -49,6 +50,8 @@
 import android.text.TextUtils;
 import android.util.Log;
 
+import com.android.system.virtualmachine.flags.Flags;
+
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
@@ -601,6 +604,18 @@
     }
 
     /**
+     * Returns whether this VM enabled the hint to use transparent huge pages.
+     *
+     * @see Builder#setShouldUseHugepages
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(Flags.FLAG_PROMOTE_SET_SHOULD_USE_HUGEPAGES_TO_SYSTEM_API)
+    public boolean shouldUseHugepages() {
+        return mShouldUseHugepages;
+    }
+
+    /**
      * Tests if this config is compatible with other config. Being compatible means that the configs
      * can be interchangeably used for the same virtual machine; they do not change the VM identity
      * or secrets. Such changes include varying the number of CPUs or the size of the RAM. Changes
@@ -1364,7 +1379,26 @@
             return this;
         }
 
-        /** @hide */
+        /**
+         * Hints whether the VM should make use of the transparent huge pages feature.
+         *
+         * <p>Note: this API just provides a hint, whether the VM will actually use transparent huge
+         * pages additionally depends on the following:
+         *
+         * <ul>
+         *   <li>{@code /sys/kernel/mm/transparent_hugepages/shmem_enabled} should be configured
+         *       with the value {@code 'advise'}.
+         *   <li>Android host kernel version should be at least {@code android15-5.15}
+         * </ul>
+         *
+         * @see https://docs.kernel.org/admin-guide/mm/transhuge.html
+         * @see
+         *     https://cs.android.com/android/platform/superproject/main/+/main:packages/modules/Virtualization/docs/hugepages.md
+         * @hide
+         */
+        @SystemApi
+        @FlaggedApi(Flags.FLAG_PROMOTE_SET_SHOULD_USE_HUGEPAGES_TO_SYSTEM_API)
+        @NonNull
         public Builder setShouldUseHugepages(boolean shouldUseHugepages) {
             mShouldUseHugepages = shouldUseHugepages;
             return this;
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 97a5e78..b2e65bc 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -175,7 +175,8 @@
 
     private static final String VM_SHARE_APP_PACKAGE_NAME = "com.android.microdroid.vmshare_app";
 
-    private void createAndConnectToVmHelper(int cpuTopology) throws Exception {
+    private void createAndConnectToVmHelper(int cpuTopology, boolean shouldUseHugepages)
+            throws Exception {
         assumeSupportedDevice();
 
         VirtualMachineConfig config =
@@ -183,6 +184,7 @@
                         .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .setCpuTopology(cpuTopology)
+                        .setShouldUseHugepages(shouldUseHugepages)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
 
@@ -210,13 +212,31 @@
     @Test
     @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-1"})
     public void createAndConnectToVm() throws Exception {
-        createAndConnectToVmHelper(CPU_TOPOLOGY_ONE_CPU);
+        createAndConnectToVmHelper(CPU_TOPOLOGY_ONE_CPU, /* shouldUseHugepages= */ false);
     }
 
     @Test
     @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-1"})
     public void createAndConnectToVm_HostCpuTopology() throws Exception {
-        createAndConnectToVmHelper(CPU_TOPOLOGY_MATCH_HOST);
+        createAndConnectToVmHelper(CPU_TOPOLOGY_MATCH_HOST, /* shouldUseHugepages= */ false);
+    }
+
+    @Test
+    public void createAndConnectToVm_WithHugepages() throws Exception {
+        // Note: setting shouldUseHugepages to true only hints that VM wants to use transparent huge
+        // pages. Whether it will actually be used depends on the value in the
+        // /sys/kernel/mm/transparent_hugepages/shmem_enabled.
+        // See packages/modules/Virtualization/docs/hugepages.md
+        createAndConnectToVmHelper(CPU_TOPOLOGY_ONE_CPU, /* shouldUseHugepages= */ true);
+    }
+
+    @Test
+    public void createAndConnectToVm_HostCpuTopology_WithHugepages() throws Exception {
+        // Note: setting shouldUseHugepages to true only hints that VM wants to use transparent huge
+        // pages. Whether it will actually be used depends on the value in the
+        // /sys/kernel/mm/transparent_hugepages/shmem_enabled.
+        // See packages/modules/Virtualization/docs/hugepages.md
+        createAndConnectToVmHelper(CPU_TOPOLOGY_MATCH_HOST, /* shouldUseHugepages= */ true);
     }
 
     @Test
@@ -579,6 +599,7 @@
         assertThat(minimal.getEncryptedStorageBytes()).isEqualTo(0);
         assertThat(minimal.isVmOutputCaptured()).isFalse();
         assertThat(minimal.getOs()).isEqualTo("microdroid");
+        assertThat(minimal.shouldUseHugepages()).isFalse();
 
         // Maximal has everything that can be set to some non-default value. (And has different
         // values than minimal for the required fields.)
@@ -594,7 +615,8 @@
                         .setCpuTopology(CPU_TOPOLOGY_MATCH_HOST)
                         .setEncryptedStorageBytes(1_000_000)
                         .setVmOutputCaptured(true)
-                        .setOs("microdroid_gki-android14-6.1");
+                        .setOs("microdroid_gki-android14-6.1")
+                        .setShouldUseHugepages(true);
         VirtualMachineConfig maximal = maximalBuilder.build();
 
         assertThat(maximal.getApkPath()).isEqualTo("/apk/path");
@@ -611,6 +633,7 @@
         assertThat(maximal.getEncryptedStorageBytes()).isEqualTo(1_000_000);
         assertThat(maximal.isVmOutputCaptured()).isTrue();
         assertThat(maximal.getOs()).isEqualTo("microdroid_gki-android14-6.1");
+        assertThat(maximal.shouldUseHugepages()).isTrue();
 
         assertThat(minimal.isCompatibleWith(maximal)).isFalse();
         assertThat(minimal.isCompatibleWith(minimal)).isTrue();
@@ -680,6 +703,7 @@
         assertConfigCompatible(
                         baseline, newBaselineBuilder().setCpuTopology(CPU_TOPOLOGY_MATCH_HOST))
                 .isTrue();
+        assertConfigCompatible(baseline, newBaselineBuilder().setShouldUseHugepages(true)).isTrue();
 
         // Changes that must be incompatible, since they must change the VM identity.
         assertConfigCompatible(baseline, newBaselineBuilder().addExtraApk("foo")).isFalse();