From 6b8ecce2657d2edf19ae52f3544f7835ae7afb76 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Emmanuel=20Fr=C3=A9con?= <efrecon@gmail.com>
Date: Fri, 5 Apr 2024 14:38:48 +0200
Subject: [PATCH] Add support for multiple krunvm VMs

Add support for multiple krunvm-based VMs in the API. This uses a
temporary directory that will store the mappings between the name of the
VM and the PID of the underlying (top) process. The directory should be
removed through microvm_cleanup upon process exit.
---
 lib/microvm.sh  | 18 ++++++++++++++++--
 orchestrator.sh |  2 ++
 runner.sh       |  2 ++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/lib/microvm.sh b/lib/microvm.sh
index 5262210..34a9470 100644
--- a/lib/microvm.sh
+++ b/lib/microvm.sh
@@ -7,6 +7,9 @@
 # automatically selected.
 : "${KRUNVM_RUNNER_RUNTIME:=""}"
 
+# Directory used for VM->PID mapping when using krunvm
+: "${KRUNVM_RUNNER_STORAGE:=""}"
+
 # Run krunvm with the provided arguments, behind a buildah unshare.
 run_krunvm() {
   debug "Running krunvm $*"
@@ -50,6 +53,9 @@ microvm_runtime() {
     krunvm)
       check_command krunvm
       check_command buildah
+      if [ -z "$KRUNVM_RUNNER_STORAGE" ]; then
+        KRUNVM_RUNNER_STORAGE="$(mktemp -d)"
+      fi
       ;;
     *)
       error "Unknown microVM runtime: $KRUNVM_RUNNER_RUNTIME"
@@ -168,6 +174,7 @@ EOF
       run_krunvm start "$KRUNVM_RUNNER_NAME" "$RUNNER_ENTRYPOINT" -- "$@" </dev/null &
       KRUNVM_RUNNER_PID=$!
       eval "$optstate"; # Restore options
+      printf %d\\n "$KRUNVM_RUNNER_PID" > "${KRUNVM_RUNNER_STORAGE}/${KRUNVM_RUNNER_NAME}.pid"
       verbose "Started microVM '$KRUNVM_RUNNER_NAME' with PID $KRUNVM_RUNNER_PID"
       wait "$KRUNVM_RUNNER_PID"
       KRUNVM_RUNNER_PID=
@@ -190,6 +197,7 @@ microvm_wait() {
     podman*)
       podman wait "$1";;
     krunvm)
+      KRUNVM_RUNNER_PID=$(cat "${KRUNVM_RUNNER_STORAGE}/$1.pid")
       if [ -n "$KRUNVM_RUNNER_PID" ]; then
         # shellcheck disable=SC2046 # We want to wait for all children
         waitpid $(ps_tree "$KRUNVM_RUNNER_PID"|tac)
@@ -216,10 +224,12 @@ microvm_stop() {
       # TODO: Specify howlog to wait between TERM and KILL?
       podman stop "$1";;
     krunvm)
+      KRUNVM_RUNNER_PID=$(cat "${KRUNVM_RUNNER_STORAGE}/$1.pid")
       if [ -n "$KRUNVM_RUNNER_PID" ]; then
         kill_tree "$KRUNVM_RUNNER_PID"
         # shellcheck disable=SC2046 # We want to wait for all children
         microvm_wait "$1"
+        rm -f "${KRUNVM_RUNNER_STORAGE}/$1.pid" || true
       fi
       ;;
     *)
@@ -241,7 +251,8 @@ microvm_delete() {
       podman rm -f "$1";;
     krunvm)
       verbose "Removing microVM '$1'"
-      run_krunvm delete "$1";;
+      run_krunvm delete "$1"
+      ;;
     *)
       error "Unknown microVM runtime: $KRUNVM_RUNNER_RUNTIME"
       ;;
@@ -265,5 +276,8 @@ microvm_pull() {
       error "Unknown microVM runtime: $KRUNVM_RUNNER_RUNTIME"
       ;;
   esac
+}
 
-}
\ No newline at end of file
+microvm_cleanup() {
+  [ -n "$KRUNVM_RUNNER_STORAGE" ] && rm -rf "$KRUNVM_RUNNER_STORAGE"
+}
diff --git a/orchestrator.sh b/orchestrator.sh
index 25e1ed8..a5ffa23 100755
--- a/orchestrator.sh
+++ b/orchestrator.sh
@@ -119,6 +119,8 @@ EOF
     verbose "Removing isolation environment $ORCHESTRATOR_ENVIRONMENT"
     rm -rf "$ORCHESTRATOR_ENVIRONMENT"
   fi
+
+  microvm_cleanup
 }
 
 # Pass the runtime to the microvm script
diff --git a/runner.sh b/runner.sh
index cb1e6ab..7f8e282 100755
--- a/runner.sh
+++ b/runner.sh
@@ -325,6 +325,8 @@ cleanup() {
     vm_terminate "$RUNNER_ID"
     vm_delete "$RUNNER_ID"
   fi
+
+  microvm_cleanup
 }
 
 trap cleanup EXIT