Replace Fladle with gcloud CLI for macrobenchmarks

auto-benchmark-plugin/plugin/build.gradle.kts

@@ -38,7 +38,6 @@ dependencies {
     compileOnly(kotlin("stdlib"))
     compileOnly("com.android.tools.build:gradle:8.0.0")
     compileOnly("org.jetbrains.kotlin:kotlin-gradle-plugin:1.8.10")
-    implementation("com.osacky.flank.gradle:fladle:0.17.4")
     testImplementation(platform("org.junit:junit-bom:5.9.1"))
     testImplementation("org.junit.jupiter:junit-jupiter")
 }

auto-benchmark-plugin/plugin/src/main/kotlin/io.github.sagar.auto_benchmark_plugin/AutoBenchmarkPlugin.kt

@@ -25,28 +25,18 @@
 package io.github.sagar.auto_benchmark_plugin
 
 import com.android.build.api.variant.AndroidComponentsExtension
-import com.osacky.flank.gradle.FlankGradleExtension
 import org.gradle.api.Plugin
 import org.gradle.api.Project
 import java.io.File
 
 class AutoBenchmarkPlugin : Plugin<Project> {
     companion object {
         private const val TASK_NAME = "runBenchmarkAndVerifyProfile"
-        private const val FLADLE_CONFIG_NAME = "autoBenchmark"
-        private const val LOCAL_RESULT_DIR = "benchmark_result"
-        private const val FLADLE_TASK_NAME = "runFlankAutoBenchmark"
-        private const val ADDITIONAL_TEST_OUTPUT_DIR = "/sdcard/Download/"
-        private const val ENV_ADDITIONAL_TEST_OUTPUT_KEY = "additionalTestOutputDir"
-        private const val ENV_ADDITIONAL_NO_ISOLATION_KEY = "no-isolated-storage"
-        private const val ENV_ADDITIONAL_NO_ISOLATION_VALUE = "true"
+        private const val BENCHMARK_TASK_NAME = "runBenchmarkOnGcloud"
         private const val PLUGIN_APPLY_ERROR_MESSAGE = "This plugin is only applicable for Android modules"
     }
 
     override fun apply(project: Project) {
-        // Apply fladle plugin as dependency
-        project.pluginManager.apply("com.osacky.fladle")
-
         // Fail if plugin is not applied to android app module
         runCatching {
             project.extensions.getByType(AndroidComponentsExtension::class.java)
@@ -55,68 +45,38 @@ class AutoBenchmarkPlugin : Plugin<Project> {
         // Get the extensions
         val extension = AutoBenchmarkExtension.create(project)
 
-        // Register Fladle configuration to run tests on firebase test lab and download benchmark result
-        configureFladle(project, extension)
         // Register task to verify downloaded benchmark result
         setupMacroBenchmarkVerificationTask(project, extension)
     }
 
     /**
-     * Adds a fladle configuration to be executed for uploading apks to Firebase test lab.
-     * Also, configures custom test output directory and download directory from G-Cloud.
+     * Register a task for benchmark result verification.
+     * This task depends on fladle task created above to run tests on firebase test lab and download result.
      *
      * @param project An instance of gradle project
      * @param extension An instance of [AutoBenchmarkExtension]
      */
-    private fun configureFladle(project: Project, extension: AutoBenchmarkExtension) {
-        project.extensions.configure(FlankGradleExtension::class.java) {
-            configs.register(FLADLE_CONFIG_NAME) {
-                with(this@configure) {
-                    flakyTestAttempts.set(1)
-                    localResultsDir.set(LOCAL_RESULT_DIR)
-                    performanceMetrics.set(false)
-                    disableSharding.set(true)
-                    devices.set(
-                        listOf(
-                            extension.physicalDevices.get()
-                        )
-                    )
-                    serviceAccountCredentials.set(File(extension.serviceAccountJsonFilePath.get()))
-                }
+    private fun setupMacroBenchmarkVerificationTask(project: Project, extension: AutoBenchmarkExtension) {
+        val benchmarkTask = project.tasks.register(BENCHMARK_TASK_NAME, RunBenchmarkOnGcloudTask::class.java) {
+            appApk.set(project.layout.file(project.provider { File(project.rootDir, extension.appApkFilePath.get()) }))
+            benchmarkApk.set(project.layout.file(project.provider { File(project.rootDir, extension.benchmarkApkFilePath.get()) }))
 
-                apply {
-                    filesToDownload.set(listOf(".*$ADDITIONAL_TEST_OUTPUT_DIR.*"))
-                    directoriesToPull.set(listOf(ADDITIONAL_TEST_OUTPUT_DIR))
-                    debugApk.set(project.provider { "${project.rootDir.path}${extension.appApkFilePath.get()}" })
-                    instrumentationApk.set(project.provider {
-                        "${project.rootDir.path}${extension.benchmarkApkFilePath.get()}"
-                    })
-                    environmentVariables.set(
-                        mapOf(
-                            ENV_ADDITIONAL_TEST_OUTPUT_KEY to ADDITIONAL_TEST_OUTPUT_DIR,
-                            ENV_ADDITIONAL_NO_ISOLATION_KEY to ENV_ADDITIONAL_NO_ISOLATION_VALUE
-                        )
-                    )
-                }
+            if (extension.serviceAccountJsonFilePath.isPresent) {
+                serviceAccountJson.set(project.layout.file(extension.serviceAccountJsonFilePath.map { File(project.rootDir, it) }))
             }
+
+            physicalDevices.set(extension.physicalDevices)
+            benchmarkResultDir.set(project.layout.buildDirectory.dir("benchmark_results"))
         }
-    }
 
-    /**
-     * Register a task for benchmark result verification.
-     * This task depends on fladle task created above to run tests on firebase test lab and download result.
-     *
-     * @param project An instance of gradle project
-     * @param extension An instance of [AutoBenchmarkExtension]
-     */
-    private fun setupMacroBenchmarkVerificationTask(project: Project, extension: AutoBenchmarkExtension) {
         project.tasks.register(
             TASK_NAME,
             BenchmarkJsonParserTask::class.java
         ) {
-            buildDirectory.set(project.buildDir)
+            // Set the build directory to the output of the benchmark task
+            buildDirectory.set(benchmarkTask.flatMap { it.benchmarkResultDir })
             tolerancePercentage.set(extension.tolerancePercentage)
-            dependsOn(FLADLE_TASK_NAME)
+            dependsOn(benchmarkTask)
         }
     }
 }

auto-benchmark-plugin/plugin/src/main/kotlin/io.github.sagar.auto_benchmark_plugin/BenchmarkJsonParserTask.kt

@@ -45,7 +45,7 @@ abstract class BenchmarkJsonParserTask : DefaultTask() {
     @TaskAction
     @Suppress("UNCHECKED_CAST")
     fun execute() {
-        val jsonFile = buildDirectory.dir("fladle/benchmark_result").get().asFileTree.filter { file ->
+        val jsonFile = buildDirectory.get().asFileTree.filter { file ->
             file.name.contains("benchmarkData")
         }.singleFile
         val json: Map<String, Any> = JsonSlurper().parse(jsonFile) as Map<String, Any>

auto-benchmark-plugin/plugin/src/main/kotlin/io.github.sagar.auto_benchmark_plugin/RunBenchmarkOnGcloudTask.kt

@@ -0,0 +1,134 @@
+package io.github.sagar.auto_benchmark_plugin
+
+import groovy.json.JsonSlurper
+import org.gradle.api.DefaultTask
+import org.gradle.api.GradleException
+import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.file.RegularFileProperty
+import org.gradle.api.provider.MapProperty
+import org.gradle.api.tasks.Input
+import org.gradle.api.tasks.InputFile
+import org.gradle.api.tasks.Optional
+import org.gradle.api.tasks.OutputDirectory
+import org.gradle.api.tasks.PathSensitive
+import org.gradle.api.tasks.PathSensitivity
+import org.gradle.api.tasks.TaskAction
+import java.io.ByteArrayOutputStream
+import java.io.File
+
+abstract class RunBenchmarkOnGcloudTask : DefaultTask() {
+
+    @get:InputFile
+    @get:PathSensitive(PathSensitivity.RELATIVE)
+    abstract val appApk: RegularFileProperty
+
+    @get:InputFile
+    @get:PathSensitive(PathSensitivity.RELATIVE)
+    abstract val benchmarkApk: RegularFileProperty
+
+    @get:InputFile
+    @get:PathSensitive(PathSensitivity.NONE)
+    @get:Optional
+    abstract val serviceAccountJson: RegularFileProperty
+
+    @get:Input
+    abstract val physicalDevices: MapProperty<String, String>
+
+    @get:OutputDirectory
+    abstract val benchmarkResultDir: DirectoryProperty
+
+    @TaskAction
+    fun execute() {
+        val appApkFile = appApk.get().asFile
+        val benchmarkApkFile = benchmarkApk.get().asFile
+
+        if (serviceAccountJson.isPresent) {
+             val serviceAccountJsonFile = serviceAccountJson.get().asFile
+             if (serviceAccountJsonFile.exists()) {
+                 authenticate(serviceAccountJsonFile)
+             } else {
+                 project.logger.warn("Service account JSON not found at: ${serviceAccountJsonFile.absolutePath}. Assuming gcloud is already authenticated.")
+             }
+        }
+
+        val deviceString = physicalDevices.get().entries.joinToString(",") { "${it.key}=${it.value}" }
+
+        val gcloudCommand = mutableListOf(
+            "gcloud", "firebase", "test", "android", "run",
+            "--type", "instrumentation",
+            "--app", appApkFile.absolutePath,
+            "--test", benchmarkApkFile.absolutePath,
+            "--device", deviceString,
+            "--directories-to-pull", "/sdcard/Download",
+            "--environment-variables", "additionalTestOutputDir=/sdcard/Download,no-isolated-storage=true",
+            "--format", "json"
+        )
+
+        project.logger.lifecycle("Running gcloud command: ${gcloudCommand.joinToString(" ")}")
+
+        val output = runCommand(gcloudCommand)
+        val gcsPath = parseGcsPath(output)
+
+        project.logger.lifecycle("Benchmark results stored at: $gcsPath")
+
+        downloadResults(gcsPath)
+    }
+
+    private fun authenticate(serviceAccountJson: File) {
+        val command = listOf(
+            "gcloud", "auth", "activate-service-account",
+            "--key-file", serviceAccountJson.absolutePath
+        )
+        runCommand(command)
+    }
+
+    private fun runCommand(command: List<String>): String {
+        val stdout = ByteArrayOutputStream()
+        val stderr = ByteArrayOutputStream()
+        val result = project.exec {
+            commandLine = command
+            standardOutput = stdout
+            errorOutput = stderr
+            isIgnoreExitValue = true
+        }
+        if (result.exitValue != 0) {
+            throw GradleException("Command failed: ${command.joinToString(" ")}\nError output: ${stderr.toString().trim()}")
+        }
+        return stdout.toString().trim()
+    }
+
+    private fun parseGcsPath(jsonOutput: String): String {
+         try {
+             val parsed = JsonSlurper().parseText(jsonOutput)
+             // It usually is a List of objects, one per execution (matrix)
+             if (parsed is List<*>) {
+                  val first = parsed.firstOrNull() as? Map<*, *>
+                  val resultStorage = first?.get("resultStorage") as? Map<*, *>
+                  val googleCloudStorage = resultStorage?.get("googleCloudStorage") as? Map<*, *>
+                  return googleCloudStorage?.get("gcsPath") as? String
+                      ?: throw GradleException("Could not find gcsPath in output")
+             } else if (parsed is Map<*, *>) {
+                  val resultStorage = parsed["resultStorage"] as? Map<*, *>
+                  val googleCloudStorage = resultStorage?.get("googleCloudStorage") as? Map<*, *>
+                  return googleCloudStorage?.get("gcsPath") as? String
+                      ?: throw GradleException("Could not find gcsPath in output")
+             }
+             throw GradleException("Unexpected JSON output format: ${parsed::class.java}")
+         } catch (e: Exception) {
+             throw GradleException("Failed to parse gcloud output: ${e.message}. Output: $jsonOutput", e)
+         }
+    }
+
+    private fun downloadResults(gcsPath: String) {
+        val destDir = benchmarkResultDir.get().asFile
+        destDir.mkdirs()
+
+        // Use gcloud storage cp to copy.
+        val command = listOf(
+            "gcloud", "storage", "cp", "-r", gcsPath, destDir.absolutePath
+        )
+
+        project.logger.lifecycle("Downloading results: ${command.joinToString(" ")}")
+        runCommand(command)
+    }
+}