setCallback accepts Executor

The callbacks are executed on the threads provided by the executor, not
on the binder threads.

Bug: 204839392
Test: atest MicrodroidTestApp
Change-Id: I1bbdb0bac29147d9fbe8520367759513c2c5d0c5
diff --git a/demo/java/com/android/microdroid/demo/MainActivity.java b/demo/java/com/android/microdroid/demo/MainActivity.java
index 60e50bb..15d9046 100644
--- a/demo/java/com/android/microdroid/demo/MainActivity.java
+++ b/demo/java/com/android/microdroid/demo/MainActivity.java
@@ -285,7 +285,7 @@
                     mVirtualMachine = vmm.create("demo_vm", config);
                 }
                 mVirtualMachine.run();
-                mVirtualMachine.setCallback(callback);
+                mVirtualMachine.setCallback(Executors.newSingleThreadExecutor(), callback);
                 mStatus.postValue(mVirtualMachine.getStatus());
 
                 InputStream console = mVirtualMachine.getConsoleOutputStream();
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 63c9288..6556b87 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -19,6 +19,7 @@
 import static android.os.ParcelFileDescriptor.MODE_READ_ONLY;
 import static android.os.ParcelFileDescriptor.MODE_READ_WRITE;
 
+import android.annotation.CallbackExecutor;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
@@ -42,6 +43,7 @@
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.Files;
 import java.util.Optional;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -110,6 +112,9 @@
     /** The registered callback */
     private @Nullable VirtualMachineCallback mCallback;
 
+    /** The executor on which the callback will be executed */
+    private @NonNull Executor mCallbackExecutor;
+
     private @Nullable ParcelFileDescriptor mConsoleReader;
     private @Nullable ParcelFileDescriptor mConsoleWriter;
 
@@ -263,7 +268,10 @@
      * Registers the callback object to get events from the virtual machine. If a callback was
      * already registered, it is replaced with the new one.
      */
-    public void setCallback(@Nullable VirtualMachineCallback callback) {
+    public void setCallback(
+            @NonNull @CallbackExecutor Executor executor,
+            @Nullable VirtualMachineCallback callback) {
+        mCallbackExecutor = executor;
         mCallback = callback;
     }
 
@@ -328,7 +336,8 @@
                             if (cb == null) {
                                 return;
                             }
-                            cb.onPayloadStarted(VirtualMachine.this, stream);
+                            mCallbackExecutor.execute(
+                                    () -> cb.onPayloadStarted(VirtualMachine.this, stream));
                         }
 
                         @Override
@@ -337,7 +346,7 @@
                             if (cb == null) {
                                 return;
                             }
-                            cb.onPayloadReady(VirtualMachine.this);
+                            mCallbackExecutor.execute(() -> cb.onPayloadReady(VirtualMachine.this));
                         }
 
                         @Override
@@ -346,7 +355,8 @@
                             if (cb == null) {
                                 return;
                             }
-                            cb.onPayloadFinished(VirtualMachine.this, exitCode);
+                            mCallbackExecutor.execute(
+                                    () -> cb.onPayloadFinished(VirtualMachine.this, exitCode));
                         }
 
                         @Override
@@ -355,7 +365,7 @@
                             if (cb == null) {
                                 return;
                             }
-                            cb.onDied(VirtualMachine.this);
+                            mCallbackExecutor.execute(() -> cb.onDied(VirtualMachine.this));
                         }
                     });
             service.asBinder()
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 619044f..e0d6cc1 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -19,8 +19,6 @@
 import static org.junit.Assume.assumeNoException;
 
 import android.content.Context;
-import android.os.Handler;
-import android.os.Looper;
 import android.os.ParcelFileDescriptor;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineCallback;
@@ -38,6 +36,10 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
 @RunWith(JUnit4.class)
 public class MicrodroidTests {
     @Rule public Timeout globalTimeout = Timeout.seconds(300);
@@ -81,59 +83,24 @@
     }
 
     private abstract static class VmEventListener implements VirtualMachineCallback {
-        private final Handler mHandler;
+        private ExecutorService mExecutorService = Executors.newSingleThreadExecutor();
 
-        VmEventListener() {
-            Looper.prepare();
-            mHandler = new Handler(Looper.myLooper());
-        }
-
-        void runToFinish(VirtualMachine vm) throws VirtualMachineException {
-            vm.setCallback(mCallback);
+        void runToFinish(VirtualMachine vm) throws VirtualMachineException, InterruptedException {
+            vm.setCallback(mExecutorService, this);
             vm.run();
-            Looper.loop();
+            mExecutorService.awaitTermination(300, TimeUnit.SECONDS);
         }
 
         void forceStop(VirtualMachine vm) {
             try {
                 vm.stop();
                 this.onDied(vm);
-                Looper.myLooper().quitSafely();
+                mExecutorService.shutdown();
             } catch (VirtualMachineException e) {
                 throw new RuntimeException(e);
             }
         }
 
-        // This is the actual listener that is registered. Since the listener is executed in another
-        // thread, post a runnable to the current thread to call the corresponding mHandler method
-        // in the current thread.
-        private final VirtualMachineCallback mCallback =
-                new VirtualMachineCallback() {
-                    @Override
-                    public void onPayloadStarted(VirtualMachine vm, ParcelFileDescriptor stream) {
-                        mHandler.post(() -> VmEventListener.this.onPayloadStarted(vm, stream));
-                    }
-
-                    @Override
-                    public void onPayloadReady(VirtualMachine vm) {
-                        mHandler.post(() -> VmEventListener.this.onPayloadReady(vm));
-                    }
-
-                    @Override
-                    public void onPayloadFinished(VirtualMachine vm, int exitCode) {
-                        mHandler.post(() -> VmEventListener.this.onPayloadFinished(vm, exitCode));
-                    }
-
-                    @Override
-                    public void onDied(VirtualMachine vm) {
-                        mHandler.post(
-                                () -> {
-                                    VmEventListener.this.onDied(vm);
-                                    Looper.myLooper().quitSafely();
-                                });
-                    }
-                };
-
         @Override
         public void onPayloadStarted(VirtualMachine vm, ParcelFileDescriptor stream) {}