From 97df44c3b95e4db38f07cd646d4006ecc9c8ad82 Mon Sep 17 00:00:00 2001
From: Brian Goff <cpuguy83@gmail.com>
Date: Wed, 19 Jan 2022 00:23:52 +0000
Subject: [PATCH] shimv2: handle sigint/sigterm

This causes sigint/sigterm to trigger a shutdown of the shim.
It is needed because otherwise the v2 shim hangs system shutdown.

Signed-off-by: Brian Goff <cpuguy83@gmail.com>
(cherry picked from commit 3ffb6a61130298cd27636df2b64a53e2161cdced)
Signed-off-by: Varsha Teratipally <teratipally@google.com>
---
 runtime/v2/shim/shim.go         | 23 +++++++++++++----------
 runtime/v2/shim/shim_unix.go    | 20 +++++++++++++++++++-
 runtime/v2/shim/shim_windows.go |  5 ++++-
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/runtime/v2/shim/shim.go b/runtime/v2/shim/shim.go
index c14aacca99b8f..46558ec1f514f 100644
--- a/runtime/v2/shim/shim.go
+++ b/runtime/v2/shim/shim.go
@@ -40,9 +40,10 @@ import (
 
 // Client for a shim server
 type Client struct {
-	service shimapi.TaskService
-	context context.Context
-	signals chan os.Signal
+	service  shimapi.TaskService
+	context  context.Context
+	signals  chan os.Signal
+	shutdown context.CancelFunc
 }
 
 // Publisher for events
@@ -214,7 +215,7 @@ func run(id string, initFunc Init, config Config) error {
 			"pid":       os.Getpid(),
 			"namespace": namespaceFlag,
 		})
-		go handleSignals(ctx, logger, signals)
+		go reap(ctx, logger, signals)
 		response, err := service.Cleanup(ctx)
 		if err != nil {
 			return err
@@ -248,7 +249,7 @@ func run(id string, initFunc Init, config Config) error {
 				return err
 			}
 		}
-		client := NewShimClient(ctx, service, signals)
+		client := NewShimClient(ctx, service, signals, cancel)
 		if err := client.Serve(); err != nil {
 			if err != context.Canceled {
 				return err
@@ -271,11 +272,12 @@ func run(id string, initFunc Init, config Config) error {
 }
 
 // NewShimClient creates a new shim server client
-func NewShimClient(ctx context.Context, svc shimapi.TaskService, signals chan os.Signal) *Client {
+func NewShimClient(ctx context.Context, svc shimapi.TaskService, signals chan os.Signal, shutdown context.CancelFunc) *Client {
 	s := &Client{
-		service: svc,
-		context: ctx,
-		signals: signals,
+		service:  svc,
+		context:  ctx,
+		signals:  signals,
+		shutdown: shutdown,
 	}
 	return s
 }
@@ -310,7 +312,8 @@ func (s *Client) Serve() error {
 			dumpStacks(logger)
 		}
 	}()
-	return handleSignals(s.context, logger, s.signals)
+	go handleExitSignals(s.context, logger, s.shutdown)
+	return reap(s.context, logger, s.signals)
 }
 
 // serve serves the ttrpc API over a unix socket at the provided path
diff --git a/runtime/v2/shim/shim_unix.go b/runtime/v2/shim/shim_unix.go
index a61b642089217..ce599a501ddd9 100644
--- a/runtime/v2/shim/shim_unix.go
+++ b/runtime/v2/shim/shim_unix.go
@@ -70,7 +70,7 @@ func serveListener(path string) (net.Listener, error) {
 	return l, nil
 }
 
-func handleSignals(ctx context.Context, logger *logrus.Entry, signals chan os.Signal) error {
+func reap(ctx context.Context, logger *logrus.Entry, signals chan os.Signal) error {
 	logger.Info("starting signal loop")
 
 	for {
@@ -78,6 +78,8 @@ func handleSignals(ctx context.Context, logger *logrus.Entry, signals chan os.Si
 		case <-ctx.Done():
 			return ctx.Err()
 		case s := <-signals:
+			// Exit signals are handled separately from this loop
+			// They get registered with this channel so that we can ignore such signals for short-running actions (e.g. `delete`)
 			switch s {
 			case unix.SIGCHLD:
 				if err := reaper.Reap(); err != nil {
@@ -89,6 +91,22 @@ func handleSignals(ctx context.Context, logger *logrus.Entry, signals chan os.Si
 	}
 }
 
+func handleExitSignals(ctx context.Context, logger *logrus.Entry, cancel context.CancelFunc) {
+	ch := make(chan os.Signal, 32)
+	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+
+	for {
+		select {
+		case s := <-ch:
+			logger.WithField("signal", s).Debugf("Caught exit signal")
+			cancel()
+			return
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
 func openLog(ctx context.Context, _ string) (io.Writer, error) {
 	return fifo.OpenFifoDup2(ctx, "log", unix.O_WRONLY, 0700, int(os.Stderr.Fd()))
 }
diff --git a/runtime/v2/shim/shim_windows.go b/runtime/v2/shim/shim_windows.go
index 7339eb2a2e4c2..955e2aecabe4d 100644
--- a/runtime/v2/shim/shim_windows.go
+++ b/runtime/v2/shim/shim_windows.go
@@ -48,10 +48,13 @@ func serveListener(path string) (net.Listener, error) {
 	return nil, errors.New("not supported")
 }
 
-func handleSignals(ctx context.Context, logger *logrus.Entry, signals chan os.Signal) error {
+func reap(ctx context.Context, logger *logrus.Entry, signals chan os.Signal) error {
 	return errors.New("not supported")
 }
 
+func handleExitSignals(ctx context.Context, logger *logrus.Entry, cancel context.CancelFunc) {
+}
+
 func openLog(ctx context.Context, _ string) (io.Writer, error) {
 	return nil, errors.New("not supported")
 }