diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 14ff111..65ce7ea 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -138,7 +138,7 @@
     private @Nullable VirtualMachineCallback mCallback;
 
     /** The executor on which the callback will be executed */
-    private @NonNull Executor mCallbackExecutor;
+    private @Nullable Executor mCallbackExecutor;
 
     private @Nullable ParcelFileDescriptor mConsoleReader;
     private @Nullable ParcelFileDescriptor mConsoleWriter;
@@ -298,11 +298,18 @@
      */
     public void setCallback(
             @NonNull @CallbackExecutor Executor executor,
-            @Nullable VirtualMachineCallback callback) {
+            @NonNull VirtualMachineCallback callback) {
         mCallbackExecutor = executor;
         mCallback = callback;
     }
 
+    /** Clears the currently registered callback. */
+    public void clearCallback() {
+        // TODO(b/220730550): synchronize with the callers of the callback
+        mCallback = null;
+        mCallbackExecutor = null;
+    }
+
     /** Returns the currently registered callback. */
     public @Nullable VirtualMachineCallback getCallback() {
         return mCallback;
@@ -369,6 +376,18 @@
             android.system.virtualizationservice.VirtualMachineConfig vmConfigParcel =
                     android.system.virtualizationservice.VirtualMachineConfig.appConfig(appConfig);
 
+            IBinder.DeathRecipient deathRecipient = new IBinder.DeathRecipient() {
+                @Override
+                public void binderDied() {
+                    final VirtualMachineCallback cb = mCallback;
+                    if (cb != null) {
+                        // TODO(b/220730550): don't call if the VM already died
+                        cb.onDied(VirtualMachine.this, VirtualMachineCallback
+                                .DEATH_REASON_VIRTUALIZATIONSERVICE_DIED);
+                    }
+                }
+            };
+
             mVirtualMachine = service.createVm(vmConfigParcel, mConsoleWriter, mLogWriter);
             mVirtualMachine.registerCallback(
                     new IVirtualMachineCallback.Stub() {
@@ -434,6 +453,7 @@
 
                         @Override
                         public void onDied(int cid, int reason) {
+                            service.asBinder().unlinkToDeath(deathRecipient, 0);
                             final VirtualMachineCallback cb = mCallback;
                             if (cb == null) {
                                 return;
@@ -447,19 +467,7 @@
                             }
                         }
                     });
-            service.asBinder()
-                    .linkToDeath(
-                            new IBinder.DeathRecipient() {
-                                @Override
-                                public void binderDied() {
-                                    final VirtualMachineCallback cb = mCallback;
-                                    if (cb != null) {
-                                        cb.onDied(VirtualMachine.this, VirtualMachineCallback
-                                                .DEATH_REASON_VIRTUALIZATIONSERVICE_DIED);
-                                    }
-                                }
-                            },
-                            0);
+            service.asBinder().linkToDeath(deathRecipient, 0);
             mVirtualMachine.start();
         } catch (IOException e) {
             throw new VirtualMachineException(e);
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 14ea402..cd9f284 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -143,6 +143,7 @@
 
         void forceStop(VirtualMachine vm) {
             try {
+                vm.clearCallback();
                 vm.stop();
                 mExecutorService.shutdown();
             } catch (VirtualMachineException e) {