From 5a5130e3f43e0b45d64e3bfa287340bd18c22cbc Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Wed, 18 Sep 2024 18:29:29 +0400
Subject: [PATCH] vetu run: support graceful termination (#59)

---
 cmd/vetu/main.go              |  3 ++-
 internal/command/run/run.go   | 26 +++++++++++++++++++++++++-
 internal/command/stop/stop.go |  1 +
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/cmd/vetu/main.go b/cmd/vetu/main.go
index c5d399f..0cae1f6 100644
--- a/cmd/vetu/main.go
+++ b/cmd/vetu/main.go
@@ -6,6 +6,7 @@ import (
 	"github.com/cirruslabs/vetu/internal/command"
 	"github.com/cirruslabs/vetu/internal/version"
 	"github.com/getsentry/sentry-go"
+	"golang.org/x/sys/unix"
 	"log"
 	"os"
 	"os/signal"
@@ -54,7 +55,7 @@ func main() {
 	})
 
 	// Set up a signal-interruptible context
-	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+	ctx, cancel := signal.NotifyContext(context.Background(), unix.SIGINT, unix.SIGTERM)
 
 	// Disable log timestamping
 	log.SetFlags(log.Flags() &^ (log.Ldate | log.Ltime))
diff --git a/internal/command/run/run.go b/internal/command/run/run.go
index 6df83a5..ce2d0ab 100644
--- a/internal/command/run/run.go
+++ b/internal/command/run/run.go
@@ -1,6 +1,8 @@
 package run
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"github.com/cirruslabs/vetu/internal/externalcommand/cloudhypervisor"
 	"github.com/cirruslabs/vetu/internal/filelock"
@@ -16,10 +18,12 @@ import (
 	"github.com/cirruslabs/vetu/internal/vmdirectory"
 	"github.com/samber/lo"
 	"github.com/spf13/cobra"
+	"golang.org/x/sys/unix"
 	"os"
 	"path/filepath"
 	"runtime"
 	"strings"
+	"time"
 )
 
 var netBridged string
@@ -182,5 +186,25 @@ func runRun(cmd *cobra.Command, args []string) error {
 	hv.Stderr = os.Stderr
 	hv.Stdin = os.Stdin
 
-	return hv.Run()
+	// Graceful Cloud Hypervisor termination[1]
+	//
+	// [1]: https://www.cloudhypervisor.org/blog/cloud-hypervisor-v0.11.0-released/#sigtermsigint-interrupt-signal-handling
+	hv.Cancel = func() error {
+		return hv.Process.Signal(unix.SIGTERM)
+	}
+
+	// If Cancel() fails to terminate the Cloud Hypervisor for some reason,
+	// ensure that it will eventually be killed after some time.
+	hv.WaitDelay = 30 * time.Second
+
+	if err := hv.Run(); err != nil {
+		// Context cancellation is not an error
+		if errors.Is(err, context.Canceled) {
+			return nil
+		}
+
+		return err
+	}
+
+	return nil
 }
diff --git a/internal/command/stop/stop.go b/internal/command/stop/stop.go
index c993523..0806371 100644
--- a/internal/command/stop/stop.go
+++ b/internal/command/stop/stop.go
@@ -79,6 +79,7 @@ func runStop(cmd *cobra.Command, args []string) error {
 
 		return fmt.Errorf("VM is still running")
 	}, retry.Context(gracefulTerminationCtx),
+		retry.Attempts(0),
 		retry.DelayType(retry.FixedDelay),
 		retry.Delay(100*time.Millisecond),
 	)