benchmark:Measure vsock host to VM transfer rate

Bug: 236123069
Bug: 244408006
Test: atest MicrodroidBenchmarks
Change-Id: I3731c7134efa6b0fddc24dc11b4e7a813f679be1
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
index 2111620..817df7d 100644
--- a/tests/benchmark/Android.bp
+++ b/tests/benchmark/Android.bp
@@ -16,7 +16,10 @@
         "truth-prebuilt",
     ],
     libs: ["android.system.virtualmachine"],
-    jni_libs: ["MicrodroidBenchmarkNativeLib"],
+    jni_libs: [
+        "MicrodroidBenchmarkNativeLib",
+        "libiovsock_host_jni",
+    ],
     platform_apis: true,
     use_embedded_native_libs: true,
     compile_multilib: "64",
@@ -25,6 +28,7 @@
 cc_library_shared {
     name: "MicrodroidBenchmarkNativeLib",
     srcs: ["src/native/benchmarkbinary.cpp"],
+    static_libs: ["libiovsock_vm"],
     shared_libs: [
         "android.system.virtualmachineservice-ndk",
         "com.android.microdroid.testservice-ndk",
@@ -34,3 +38,12 @@
         "liblog",
     ],
 }
+
+cc_library {
+    name: "libiovsock_vm",
+    srcs: ["src/native/io_vsock.cpp"],
+    export_include_dirs: ["src/native/include"],
+    shared_libs: [
+        "libbase",
+    ],
+}
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/IoVsockHostNative.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/IoVsockHostNative.java
new file mode 100644
index 0000000..2957ec0
--- /dev/null
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/IoVsockHostNative.java
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2022 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.microdroid.benchmark;
+
+class IoVsockHostNative {
+    static {
+        System.loadLibrary("iovsock_host_jni");
+    }
+
+    /**
+     * Measures the rate for sending data from host to VM.
+     *
+     * @return The send rate in MB/s for sending data from host to VM.
+     */
+    static native double measureSendRate(int fd, int numBytesToSend);
+}
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index 5f52989..e19c72a 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -23,10 +23,12 @@
 
 import android.app.Instrumentation;
 import android.os.Bundle;
+import android.os.ParcelFileDescriptor;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineConfig;
 import android.system.virtualmachine.VirtualMachineConfig.DebugLevel;
 import android.system.virtualmachine.VirtualMachineException;
+import android.util.Log;
 
 import com.android.microdroid.test.MicrodroidDeviceTestBase;
 import com.android.microdroid.testservice.IBenchmarkService;
@@ -43,12 +45,13 @@
 import java.nio.file.Files;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 @RunWith(Parameterized.class)
 public class MicrodroidBenchmarks extends MicrodroidDeviceTestBase {
     private static final String TAG = "MicrodroidBenchmarks";
     private static final String METRIC_NAME_PREFIX = "avf_perf/microdroid/";
-    private static final int VIRTIO_BLK_TRIAL_COUNT = 5;
+    private static final int IO_TEST_TRIAL_COUNT = 5;
 
     @Rule public Timeout globalTimeout = Timeout.seconds(300);
 
@@ -178,6 +181,25 @@
     }
 
     @Test
+    public void testVsockTransferFromHostToVM() throws Exception {
+        VirtualMachineConfig config =
+                mInner.newVmConfigBuilder("assets/vm_config_io.json")
+                        .debugLevel(DebugLevel.FULL)
+                        .build();
+        List<Double> transferRates = new ArrayList<>();
+
+        for (int i = 0; i < IO_TEST_TRIAL_COUNT; ++i) {
+            int port = (mProtectedVm ? 5666 : 6666) + i;
+            String vmName = "test_vm_io_" + i;
+            mInner.forceCreateNewVirtualMachine(vmName, config);
+            VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
+            VsockVmEventListener listener = new VsockVmEventListener(transferRates, port);
+            listener.runToFinish(TAG, vm);
+        }
+        reportMetrics(transferRates, "vsock/transfer_host_to_vm_", "_mb_per_sec");
+    }
+
+    @Test
     public void testVirtioBlkSeqReadRate() throws Exception {
         testVirtioBlkReadRate(/*isRand=*/ false);
     }
@@ -188,12 +210,13 @@
     }
 
     private void testVirtioBlkReadRate(boolean isRand) throws Exception {
-        VirtualMachineConfig.Builder builder =
-                mInner.newVmConfigBuilder("assets/vm_config_io.json");
-        VirtualMachineConfig config = builder.debugLevel(DebugLevel.FULL).build();
+        VirtualMachineConfig config =
+                mInner.newVmConfigBuilder("assets/vm_config_io.json")
+                        .debugLevel(DebugLevel.FULL)
+                        .build();
         List<Double> readRates = new ArrayList<>();
 
-        for (int i = 0; i < VIRTIO_BLK_TRIAL_COUNT + 1; ++i) {
+        for (int i = 0; i < IO_TEST_TRIAL_COUNT + 1; ++i) {
             if (i == 1) {
                 // Clear the first result because when the file was loaded the first time,
                 // the data also needs to be loaded from hard drive to host. This is
@@ -333,4 +356,50 @@
             forceStop(vm);
         }
     }
+
+    private static class VsockVmEventListener extends VmEventListener {
+        private static final int NUM_BYTES_TO_TRANSFER = 48 * 1024 * 1024;
+
+        private final List<Double> mReadRates;
+        private final int mPort;
+
+        VsockVmEventListener(List<Double> readRates, int port) {
+            mReadRates = readRates;
+            mPort = port;
+        }
+
+        @Override
+        public void onPayloadReady(VirtualMachine vm) {
+            try {
+                IBenchmarkService benchmarkService =
+                        IBenchmarkService.Stub.asInterface(
+                                vm.connectToVsockServer(IBenchmarkService.SERVICE_PORT).get());
+                assertThat(benchmarkService).isNotNull();
+                AtomicReference<Double> sendRate = new AtomicReference();
+
+                int serverFd = benchmarkService.initVsockServer(mPort);
+                new Thread(() -> sendRate.set(runVsockClientAndSendData(vm))).start();
+                benchmarkService.runVsockServerAndReceiveData(serverFd, NUM_BYTES_TO_TRANSFER);
+
+                mReadRates.add(sendRate.get());
+            } catch (Exception e) {
+                Log.e(TAG, "Test failed in VM:" + e);
+                throw new RuntimeException(e);
+            }
+            forceStop(vm);
+        }
+
+        private double runVsockClientAndSendData(VirtualMachine vm) {
+            try {
+                ParcelFileDescriptor fd = vm.connectVsock(mPort);
+                double sendRate =
+                        IoVsockHostNative.measureSendRate(fd.getFd(), NUM_BYTES_TO_TRANSFER);
+                fd.closeWithError("Cannot close socket file descriptor");
+                return sendRate;
+            } catch (Exception e) {
+                Log.e(TAG, "Error inside runVsockClientAndSendData():" + e);
+                throw new RuntimeException(e);
+            }
+        }
+    }
 }
diff --git a/tests/benchmark/src/jni/Android.bp b/tests/benchmark/src/jni/Android.bp
new file mode 100644
index 0000000..e1bc8b0
--- /dev/null
+++ b/tests/benchmark/src/jni/Android.bp
@@ -0,0 +1,10 @@
+package{
+    default_applicable_licenses : ["Android-Apache-2.0"],
+}
+
+cc_library_shared {
+    name: "libiovsock_host_jni",
+    srcs: [ "io_vsock_host_jni.cpp" ],
+    header_libs: ["jni_headers"],
+    shared_libs: ["libbase"],
+}
\ No newline at end of file
diff --git a/tests/benchmark/src/jni/io_vsock_host_jni.cpp b/tests/benchmark/src/jni/io_vsock_host_jni.cpp
new file mode 100644
index 0000000..7f3d655
--- /dev/null
+++ b/tests/benchmark/src/jni/io_vsock_host_jni.cpp
@@ -0,0 +1,54 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/result.h>
+#include <jni.h>
+#include <time.h>
+
+using android::base::Error;
+using android::base::Result;
+using android::base::WriteStringToFd;
+
+constexpr size_t kNumBytesPerMB = 1024 * 1024;
+
+Result<double> measure_send_rate(int fd, int num_bytes_to_send) {
+    std::string data;
+    data.assign(num_bytes_to_send, 'a');
+    clock_t start = clock();
+    if (!WriteStringToFd(data, fd)) {
+        return Error() << "Cannot send data to client";
+    }
+    clock_t end = clock();
+    double elapsed_seconds = (double)(end - start) / CLOCKS_PER_SEC;
+    LOG(INFO) << "Host:Finished sending data in " << elapsed_seconds << " seconds.";
+    double send_rate = num_bytes_to_send / kNumBytesPerMB / elapsed_seconds;
+    return {send_rate};
+}
+
+extern "C" JNIEXPORT jdouble JNICALL
+Java_com_android_microdroid_benchmark_IoVsockHostNative_measureSendRate(__unused JNIEnv *env,
+                                                                        __unused jclass clazz,
+                                                                        int fd,
+                                                                        int num_bytes_to_send) {
+    if (auto res = measure_send_rate(fd, num_bytes_to_send); res.ok()) {
+        return res.value();
+    } else {
+        LOG(ERROR) << "Cannot send data from host to VM: " << res.error();
+        abort();
+    }
+}
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 2558a7d..28799c8 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -16,22 +16,23 @@
 
 #include <aidl/android/system/virtualmachineservice/IVirtualMachineService.h>
 #include <aidl/com/android/microdroid/testservice/BnBenchmarkService.h>
+#include <android-base/logging.h>
+#include <android-base/parseint.h>
 #include <android-base/result.h>
+#include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 #include <fcntl.h>
 #include <linux/vm_sockets.h>
 #include <stdio.h>
+#include <time.h>
 #include <unistd.h>
 
 #include <binder_rpc_unstable.hpp>
-#include <chrono>
 #include <fstream>
 #include <random>
 #include <string>
 
-#include "android-base/logging.h"
-#include "android-base/parseint.h"
-#include "android-base/strings.h"
+#include "io_vsock.h"
 
 using aidl::android::system::virtualmachineservice::IVirtualMachineService;
 using android::base::ErrnoError;
@@ -74,6 +75,20 @@
         return ndk::ScopedAStatus::ok();
     }
 
+    ndk::ScopedAStatus initVsockServer(int32_t port, int32_t* out) override {
+        auto res = io_vsock::init_vsock_server(port);
+        if (res.ok()) {
+            *out = res.value();
+        }
+        return resultStatus(res);
+    }
+
+    ndk::ScopedAStatus runVsockServerAndReceiveData(int32_t server_fd,
+                                                    int32_t num_bytes_to_receive) override {
+        auto res = io_vsock::run_vsock_server_and_receive_data(server_fd, num_bytes_to_receive);
+        return resultStatus(res);
+    }
+
 private:
     /** Returns the elapsed seconds for reading the file. */
     Result<double> read_file(const std::string& filename, int64_t fileSizeBytes, bool is_rand) {
@@ -170,12 +185,10 @@
             sleep(1000);
         }
     } else if (strcmp(argv[1], "io") == 0) {
-        if (auto res = run_io_benchmark_tests(); res.ok()) {
-            return 0;
-        } else {
+        if (auto res = run_io_benchmark_tests(); !res.ok()) {
             LOG(ERROR) << "IO benchmark test failed: " << res.error() << "\n";
-            return 1;
+            return EXIT_FAILURE;
         }
     }
-    return 0;
+    return EXIT_SUCCESS;
 }
diff --git a/tests/benchmark/src/native/include/io_vsock.h b/tests/benchmark/src/native/include/io_vsock.h
new file mode 100644
index 0000000..3cd6339
--- /dev/null
+++ b/tests/benchmark/src/native/include/io_vsock.h
@@ -0,0 +1,25 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+#pragma once
+
+#include <android-base/result.h>
+
+namespace io_vsock {
+using android::base::Result;
+Result<int> init_vsock_server(unsigned int port);
+Result<void> run_vsock_server_and_receive_data(int server_fd, int num_bytes_to_receive);
+} // namespace io_vsock
diff --git a/tests/benchmark/src/native/io_vsock.cpp b/tests/benchmark/src/native/io_vsock.cpp
new file mode 100644
index 0000000..8edc7c8
--- /dev/null
+++ b/tests/benchmark/src/native/io_vsock.cpp
@@ -0,0 +1,73 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+#include "io_vsock.h"
+
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/result.h>
+#include <android-base/unique_fd.h>
+#include <linux/vm_sockets.h>
+#include <sys/socket.h>
+
+using namespace android::base;
+
+namespace io_vsock {
+Result<int> init_vsock_server(unsigned int port) {
+    int server_fd(TEMP_FAILURE_RETRY(socket(AF_VSOCK, SOCK_STREAM | SOCK_CLOEXEC, 0)));
+    if (server_fd < 0) {
+        return Error() << "VM:cannot create socket";
+    }
+    struct sockaddr_vm server_sa = (struct sockaddr_vm){
+            .svm_family = AF_VSOCK,
+            .svm_port = port,
+            .svm_cid = VMADDR_CID_ANY,
+    };
+    LOG(INFO) << "VM:Connecting on port " << port << "...";
+    int ret = TEMP_FAILURE_RETRY(bind(server_fd, (struct sockaddr *)&server_sa, sizeof(server_sa)));
+    if (ret < 0) {
+        return Error() << "VM:cannot bind an address with the socket";
+    }
+    ret = TEMP_FAILURE_RETRY(listen(server_fd, /*backlog=*/1));
+    if (ret < 0) {
+        return Error() << "VM:cannot listen to port";
+    }
+    LOG(INFO) << "Server now listening";
+    return server_fd;
+}
+
+Result<void> run_vsock_server_and_receive_data(int server_fd, int num_bytes_to_receive) {
+    LOG(INFO) << "Accepting connection...";
+    struct sockaddr_vm client_sa;
+    socklen_t client_sa_len = sizeof(client_sa);
+    unique_fd client_fd(TEMP_FAILURE_RETRY(
+            accept4(server_fd, (struct sockaddr *)&client_sa, &client_sa_len, SOCK_CLOEXEC)));
+    if (client_fd < 0) {
+        return Error() << "Cannot retrieve connect requests";
+    }
+    LOG(INFO) << "VM:Connection from CID " << client_sa.svm_cid << " on port "
+              << client_sa.svm_port;
+    std::string data;
+    if (!ReadFdToString(client_fd, &data)) {
+        return Error() << "Cannot get data from the host.";
+    }
+    if (data.length() != num_bytes_to_receive) {
+        return Error() << "Received data length(" << data.length() << ") is not equal to "
+                       << num_bytes_to_receive;
+    }
+    LOG(INFO) << "VM:Finished reading data.";
+    return {};
+}
+} // namespace io_vsock
\ No newline at end of file