guest: debian: Implement storage_balloon_agent

Add storage_balloon_agent daemon in Debian and its client in
the TerminalApp for enabling storage ballooning.
This feature is hidden behind a feature flag 'terminalStorageBalloon'.

Since we still use non-sparse disks, the balloon shouldn't affect
the guest's disk space yet.

Bug: 382174138
Test: Run a VM and check logs

Change-Id: I75d926bb8aa8a02bf635e94e35715e5aa23c8090
diff --git a/android/TerminalApp/Android.bp b/android/TerminalApp/Android.bp
index 2bac412..e1e236a 100644
--- a/android/TerminalApp/Android.bp
+++ b/android/TerminalApp/Android.bp
@@ -15,6 +15,7 @@
         "android.system.virtualizationservice_internal-java",
         "androidx-constraintlayout_constraintlayout",
         "androidx.window_window",
+        "androidx.work_work-runtime",
         "apache-commons-compress",
         "avf_aconfig_flags_java",
         "com.google.android.material_material",
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/DebianServiceImpl.kt b/android/TerminalApp/java/com/android/virtualization/terminal/DebianServiceImpl.kt
index e035ad4..e81be7f 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/DebianServiceImpl.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/DebianServiceImpl.kt
@@ -18,7 +18,8 @@
 import android.content.Context
 import android.util.Log
 import androidx.annotation.Keep
-import com.android.virtualization.terminal.DebianServiceImpl.ForwarderHostCallback
+import com.android.internal.annotations.GuardedBy
+import com.android.system.virtualmachine.flags.Flags
 import com.android.virtualization.terminal.MainActivity.Companion.TAG
 import com.android.virtualization.terminal.PortsStateManager.Companion.getInstance
 import com.android.virtualization.terminal.proto.DebianServiceGrpc.DebianServiceImplBase
@@ -28,6 +29,8 @@
 import com.android.virtualization.terminal.proto.ReportVmActivePortsResponse
 import com.android.virtualization.terminal.proto.ShutdownQueueOpeningRequest
 import com.android.virtualization.terminal.proto.ShutdownRequestItem
+import com.android.virtualization.terminal.proto.StorageBalloonQueueOpeningRequest
+import com.android.virtualization.terminal.proto.StorageBalloonRequestItem
 import io.grpc.stub.ServerCallStreamObserver
 import io.grpc.stub.StreamObserver
 
@@ -35,6 +38,8 @@
     private val portsStateManager: PortsStateManager = getInstance(context)
     private var portsStateListener: PortsStateManager.Listener? = null
     private var shutdownRunnable: Runnable? = null
+    private val mLock = Object()
+    @GuardedBy("mLock") private var storageBalloonCallback: StorageBalloonCallback? = null
 
     override fun reportVmActivePorts(
         request: ReportVmActivePortsRequest,
@@ -80,10 +85,9 @@
         request: ShutdownQueueOpeningRequest?,
         responseObserver: StreamObserver<ShutdownRequestItem?>,
     ) {
-        val serverCallStreamObserver = responseObserver as ServerCallStreamObserver<ShutdownRequestItem?>
-        serverCallStreamObserver.setOnCancelHandler {
-            shutdownRunnable = null
-        }
+        val serverCallStreamObserver =
+            responseObserver as ServerCallStreamObserver<ShutdownRequestItem?>
+        serverCallStreamObserver.setOnCancelHandler { shutdownRunnable = null }
         Log.d(TAG, "openShutdownRequestQueue")
         shutdownRunnable = Runnable {
             if (serverCallStreamObserver.isCancelled()) {
@@ -95,6 +99,60 @@
         }
     }
 
+    private class StorageBalloonCallback(
+        private val responseObserver: StreamObserver<StorageBalloonRequestItem?>
+    ) {
+        fun setAvailableStorageBytes(availableBytes: Long) {
+            Log.d(TAG, "send setStorageBalloon: $availableBytes")
+            val item =
+                StorageBalloonRequestItem.newBuilder().setAvailableBytes(availableBytes).build()
+            responseObserver.onNext(item)
+        }
+
+        fun closeConnection() {
+            Log.d(TAG, "close StorageBalloonQueue")
+            responseObserver.onCompleted()
+        }
+    }
+
+    fun setAvailableStorageBytes(availableBytes: Long): Boolean {
+        synchronized(mLock) {
+            if (storageBalloonCallback == null) {
+                Log.d(TAG, "storageBalloonCallback is not ready.")
+                return false
+            }
+            storageBalloonCallback!!.setAvailableStorageBytes(availableBytes)
+        }
+        return true
+    }
+
+    override fun openStorageBalloonRequestQueue(
+        request: StorageBalloonQueueOpeningRequest?,
+        responseObserver: StreamObserver<StorageBalloonRequestItem?>,
+    ) {
+        if (!Flags.terminalStorageBalloon()) {
+            return
+        }
+        Log.d(TAG, "openStorageRequestQueue")
+        synchronized(mLock) {
+            if (storageBalloonCallback != null) {
+                Log.d(TAG, "RequestQueue already exists. Closing connection.")
+                storageBalloonCallback!!.closeConnection()
+            }
+            storageBalloonCallback = StorageBalloonCallback(responseObserver)
+        }
+    }
+
+    fun closeStorageBalloonRequestQueue() {
+        Log.d(TAG, "Stopping storage balloon queue")
+        synchronized(mLock) {
+            if (storageBalloonCallback != null) {
+                storageBalloonCallback!!.closeConnection()
+                storageBalloonCallback = null
+            }
+        }
+    }
+
     @Keep
     private class ForwarderHostCallback(
         private val responseObserver: StreamObserver<ForwardingRequestItem?>
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/StorageBalloonWorker.kt b/android/TerminalApp/java/com/android/virtualization/terminal/StorageBalloonWorker.kt
new file mode 100644
index 0000000..345bf75
--- /dev/null
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/StorageBalloonWorker.kt
@@ -0,0 +1,111 @@
+/*
+ * Copyright (C) 2025 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.virtualization.terminal
+
+import android.content.Context
+import android.os.storage.StorageManager
+import android.os.storage.StorageManager.UUID_DEFAULT
+import android.util.Log
+import androidx.work.WorkManager
+import androidx.work.Worker
+import androidx.work.WorkerParameters
+import com.android.virtualization.terminal.MainActivity.Companion.TAG
+import java.util.concurrent.TimeUnit
+
+class StorageBalloonWorker(appContext: Context, workerParams: WorkerParameters) :
+    Worker(appContext, workerParams) {
+
+    override fun doWork(): Result {
+        Log.d(TAG, "StorageBalloonWorker.doWork() called")
+
+        var storageManager =
+            applicationContext.getSystemService(Context.STORAGE_SERVICE) as StorageManager
+        val hostAllocatableBytes = storageManager.getAllocatableBytes(UUID_DEFAULT)
+
+        val guestAvailableBytes = calculateGuestAvailableStorageSize(hostAllocatableBytes)
+        // debianService must be set when this function is called.
+        debianService!!.setAvailableStorageBytes(guestAvailableBytes)
+
+        val delaySeconds = calculateDelaySeconds(hostAllocatableBytes)
+        scheduleNextTask(delaySeconds)
+
+        return Result.success()
+    }
+
+    private fun calculateGuestAvailableStorageSize(hostAllocatableBytes: Long): Long {
+        return hostAllocatableBytes - HOST_RESERVED_BYTES
+    }
+
+    private fun calculateDelaySeconds(hostAvailableBytes: Long): Long {
+        return when {
+            hostAvailableBytes < CRITICAL_STORAGE_THRESHOLD_BYTES -> CRITICAL_DELAY_SECONDS
+            hostAvailableBytes < LOW_STORAGE_THRESHOLD_BYTES -> LOW_STORAGE_DELAY_SECONDS
+            hostAvailableBytes < MODERATE_STORAGE_THRESHOLD_BYTES -> MODERATE_STORAGE_DELAY_SECONDS
+            else -> NORMAL_DELAY_SECONDS
+        }
+    }
+
+    private fun scheduleNextTask(delaySeconds: Long) {
+        val storageBalloonTaskRequest =
+            androidx.work.OneTimeWorkRequest.Builder(StorageBalloonWorker::class.java)
+                .setInitialDelay(delaySeconds, TimeUnit.SECONDS)
+                .build()
+        androidx.work.WorkManager.getInstance(applicationContext)
+            .enqueueUniqueWork(
+                "storageBalloonTask",
+                androidx.work.ExistingWorkPolicy.REPLACE,
+                storageBalloonTaskRequest,
+            )
+        Log.d(TAG, "next storage balloon task is scheduled in $delaySeconds seconds")
+    }
+
+    companion object {
+        private var debianService: DebianServiceImpl? = null
+
+        // Reserve 1GB as host-only region.
+        private const val HOST_RESERVED_BYTES = 1024L * 1024 * 1024
+
+        // Thresholds for deciding time period to report storage information to the guest.
+        // Less storage is available on the host, more frequently the host will report storage
+        // information to the guest.
+        //
+        // Critical: (host storage < 1GB) => report every 5 seconds
+        private const val CRITICAL_STORAGE_THRESHOLD_BYTES = 1L * 1024 * 1024 * 1024
+        private const val CRITICAL_DELAY_SECONDS = 5L
+        // Low: (1GB <= storage < 5GB) => report every 60 seconds
+        private const val LOW_STORAGE_THRESHOLD_BYTES = 5L * 1024 * 1024 * 1024
+        private const val LOW_STORAGE_DELAY_SECONDS = 60L
+        // Moderate: (5GB <= storage < 10GB) => report every 15 minutes
+        private const val MODERATE_STORAGE_THRESHOLD_BYTES = 10L * 1024 * 1024 * 1024
+        private const val MODERATE_STORAGE_DELAY_SECONDS = 15L * 60
+        // Normal: report every 60 minutes
+        private const val NORMAL_DELAY_SECONDS = 60L * 60
+
+        internal fun start(ctx: Context, ds: DebianServiceImpl) {
+            debianService = ds
+            val storageBalloonTaskRequest =
+                androidx.work.OneTimeWorkRequest.Builder(StorageBalloonWorker::class.java)
+                    .setInitialDelay(1, TimeUnit.SECONDS)
+                    .build()
+            androidx.work.WorkManager.getInstance(ctx)
+                .enqueueUniqueWork(
+                    "storageBalloonTask",
+                    androidx.work.ExistingWorkPolicy.REPLACE,
+                    storageBalloonTaskRequest,
+                )
+        }
+    }
+}
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt b/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
index 1857175..0a1f0ee 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
@@ -38,10 +38,9 @@
 import android.system.virtualmachine.VirtualMachineException
 import android.util.Log
 import android.widget.Toast
-import com.android.system.virtualmachine.flags.Flags.terminalGuiSupport
+import com.android.system.virtualmachine.flags.Flags
 import com.android.virtualization.terminal.MainActivity.Companion.TAG
 import com.android.virtualization.terminal.Runner.Companion.create
-import com.android.virtualization.terminal.VmLauncherService.VmLauncherServiceCallback
 import io.grpc.Grpc
 import io.grpc.InsecureServerCredentials
 import io.grpc.Metadata
@@ -54,7 +53,6 @@
 import java.io.File
 import java.io.FileOutputStream
 import java.io.IOException
-import java.lang.RuntimeException
 import java.net.InetSocketAddress
 import java.net.SocketAddress
 import java.nio.file.Files
@@ -285,7 +283,7 @@
 
         // Set the initial display size
         // TODO(jeongik): set up the display size on demand
-        if (terminalGuiSupport() && displayInfo != null) {
+        if (Flags.terminalGuiSupport() && displayInfo != null) {
             builder
                 .setDisplayConfig(
                     VirtualMachineCustomImageConfig.DisplayConfig.Builder()
@@ -360,6 +358,10 @@
                 }
             }
         )
+
+        if (Flags.terminalStorageBalloon()) {
+            StorageBalloonWorker.start(this, debianService!!)
+        }
     }
 
     override fun onDestroy() {
@@ -383,6 +385,7 @@
 
     private fun stopDebianServer() {
         debianService?.killForwarderHost()
+        debianService?.closeStorageBalloonRequestQueue()
         server?.shutdown()
     }