diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/InstalledImage.kt b/android/TerminalApp/java/com/android/virtualization/terminal/InstalledImage.kt
index 59be4ae..50aaa33 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/InstalledImage.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/InstalledImage.kt
@@ -19,6 +19,7 @@
 import android.os.FileUtils
 import android.system.ErrnoException
 import android.system.Os
+import android.system.OsConstants
 import android.util.Log
 import com.android.virtualization.terminal.MainActivity.Companion.TAG
 import java.io.BufferedReader
@@ -127,13 +128,36 @@
         }
 
         if (roundedUpDesiredSize > curSize) {
-            allocateSpace(rootPartition, roundedUpDesiredSize)
+            if (!allocateSpace(roundedUpDesiredSize)) {
+                return curSize
+            }
         }
         resizeFilesystem(rootPartition, roundedUpDesiredSize)
         return getApparentSize()
     }
 
     @Throws(IOException::class)
+    private fun allocateSpace(sizeInBytes: Long): Boolean {
+        val curSizeInBytes = getApparentSize()
+        try {
+            RandomAccessFile(rootPartition.toFile(), "rw").use { raf ->
+                Os.posix_fallocate(raf.fd, 0, sizeInBytes)
+            }
+            Log.d(TAG, "Allocated space to: $sizeInBytes bytes")
+            return true
+        } catch (e: ErrnoException) {
+            Log.e(TAG, "Failed to allocate space", e)
+            if (e.errno == OsConstants.ENOSPC) {
+                Log.d(TAG, "Trying to truncate disk into the original size")
+                truncate(curSizeInBytes)
+                return false
+            } else {
+                throw IOException("Failed to allocate space", e)
+            }
+        }
+    }
+
+    @Throws(IOException::class)
     fun shrinkToMinimumSize(): Long {
         // Fix filesystem before resizing.
         runE2fsck(rootPartition)
@@ -153,8 +177,8 @@
             RandomAccessFile(rootPartition.toFile(), "rw").use { raf -> Os.ftruncate(raf.fd, size) }
             Log.d(TAG, "Truncated space to: $size bytes")
         } catch (e: ErrnoException) {
-            Log.e(TAG, "Failed to allocate space", e)
-            throw IOException("Failed to allocate space", e)
+            Log.e(TAG, "Failed to truncate space", e)
+            throw IOException("Failed to truncate space", e)
         }
     }
 
@@ -175,19 +199,6 @@
         }
 
         @Throws(IOException::class)
-        private fun allocateSpace(path: Path, sizeInBytes: Long) {
-            try {
-                RandomAccessFile(path.toFile(), "rw").use { raf ->
-                    Os.posix_fallocate(raf.fd, 0, sizeInBytes)
-                }
-                Log.d(TAG, "Allocated space to: $sizeInBytes bytes")
-            } catch (e: ErrnoException) {
-                Log.e(TAG, "Failed to allocate space", e)
-                throw IOException("Failed to allocate space", e)
-            }
-        }
-
-        @Throws(IOException::class)
         private fun runE2fsck(path: Path) {
             val p: String = path.toAbsolutePath().toString()
             runCommand("/system/bin/e2fsck", "-y", "-f", p)
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt b/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
index d43a8d1..84168e5 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
@@ -463,15 +463,32 @@
 
     @WorkerThread
     private fun doShutdown(resultReceiver: ResultReceiver?) {
-        stopForeground(STOP_FOREGROUND_REMOVE)
+        runner?.exitStatus?.thenAcceptAsync { success: Boolean ->
+            resultReceiver?.send(if (success) RESULT_STOP else RESULT_ERROR, null)
+            stopSelf()
+        }
         if (debianService != null && debianService!!.shutdownDebian()) {
             // During shutdown, change the notification content to indicate that it's closing
             val notification = createNotificationForTerminalClose()
             getSystemService<NotificationManager?>(NotificationManager::class.java)
                 .notify(this.hashCode(), notification)
-            runner?.exitStatus?.thenAcceptAsync { success: Boolean ->
-                resultReceiver?.send(if (success) RESULT_STOP else RESULT_ERROR, null)
-                stopSelf()
+
+            runner?.also { r ->
+                // For the case that shutdown from the guest agent fails.
+                // When timeout is set, the original CompletableFuture's every `thenAcceptAsync` is
+                // canceled as well. So add empty `thenAcceptAsync` to avoid interference.
+                r.exitStatus
+                    .thenAcceptAsync {}
+                    .orTimeout(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)
+                    .exceptionally {
+                        Log.e(
+                            TAG,
+                            "Stop the service directly because the VM instance isn't stopped with " +
+                                "graceful shutdown",
+                        )
+                        r.vm.stop()
+                        null
+                    }
             }
             runner = null
         } else {
@@ -498,6 +515,7 @@
         stopDebianServer()
         bgThreads.shutdownNow()
         mainWorkerThread.shutdown()
+        stopForeground(STOP_FOREGROUND_REMOVE)
         super.onDestroy()
     }
 
@@ -517,6 +535,8 @@
         private const val KEY_TERMINAL_IPADDRESS = "address"
         private const val KEY_TERMINAL_PORT = "port"
 
+        private const val SHUTDOWN_TIMEOUT_SECONDS = 3L
+
         private const val GUEST_SPARSE_DISK_SIZE_PERCENTAGE = 95
         private const val EXPECTED_PHYSICAL_SIZE_PERCENTAGE_FOR_NON_SPARSE = 90
 
diff --git a/build/apex/Android.bp b/build/apex/Android.bp
index 8934de0..f0eba7f 100644
--- a/build/apex/Android.bp
+++ b/build/apex/Android.bp
@@ -56,6 +56,7 @@
         "libvirtualizationservice_jni",
         "libvirtualmachine_jni",
     ],
+    native_shared_libs: ["libavf"],
     // TODO(b/295593640) Unfortunately these are added to the apex even though they are unused.
     // Once the build system is fixed, remove this.
     unwanted_transitive_deps: [
@@ -108,7 +109,6 @@
                 "rialto_bin",
                 "android_bootloader_crosvm_aarch64",
             ],
-            native_shared_libs: ["libavf"],
         },
         x86_64: {
             binaries: [
@@ -129,7 +129,6 @@
             prebuilts: [
                 "android_bootloader_crosvm_x86_64",
             ],
-            native_shared_libs: ["libavf"],
         },
     },
     binaries: [
diff --git a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
index ad63206..f6cd1cf 100644
--- a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
+++ b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
@@ -813,6 +813,10 @@
      */
     @GuardedBy("mLock")
     private void dropVm() {
+        if (mInputEventExecutor != null) {
+            mInputEventExecutor.shutdownNow();
+            mInputEventExecutor = null;
+        }
         if (mMemoryManagementCallbacks != null) {
             mContext.unregisterComponentCallbacks(mMemoryManagementCallbacks);
         }
diff --git a/libs/libavf/Android.bp b/libs/libavf/Android.bp
index aceb927..4469af5 100644
--- a/libs/libavf/Android.bp
+++ b/libs/libavf/Android.bp
@@ -40,30 +40,8 @@
     defaults: ["libavf.default"],
 }
 
-soong_config_module_type {
-    name: "virt_cc_defaults",
-    module_type: "cc_defaults",
-    config_namespace: "ANDROID",
-    bool_variables: [
-        "avf_enabled",
-    ],
-    properties: [
-        "apex_available",
-    ],
-}
-
-virt_cc_defaults {
-    name: "libavf_apex_available_defaults",
-    soong_config_variables: {
-        avf_enabled: {
-            apex_available: ["com.android.virt"],
-        },
-    },
-}
-
 cc_library {
     name: "libavf",
-    defaults: ["libavf_apex_available_defaults"],
     llndk: {
         symbol_file: "libavf.map.txt",
         moved_to_apex: true,
@@ -79,4 +57,5 @@
     stubs: {
         symbol_file: "libavf.map.txt",
     },
+    apex_available: ["com.android.virt"],
 }