Improve shim shutdown logic

Shims no longer call `os.Exit` but close the context on shutdown so that
events and other resources have hit the `defer`s.

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>

runtime/v2/shim/shim.go

@@ -20,11 +20,13 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"runtime"
 	"runtime/debug"
 	"strings"
+	"sync"
 	"time"
 
 	v1 "github.com/containerd/containerd/api/services/ttrpc/events/v1"
@@ -46,8 +48,14 @@ type Client struct {
 	signals chan os.Signal
 }
 
+// Publisher for events
+type Publisher interface {
+	events.Publisher
+	io.Closer
+}
+
 // Init func for the creation of a shim server
-type Init func(context.Context, string, events.Publisher) (Shim, error)
+type Init func(context.Context, string, Publisher, func()) (Shim, error)
 
 // Shim server interface
 type Shim interface {
@@ -156,15 +164,18 @@ func run(id string, initFunc Init, config Config) error {
 			return err
 		}
 	}
-
-	publisher := &remoteEventsPublisher{
-		address: fmt.Sprintf("%s.ttrpc", addressFlag),
-	}
-	conn, err := connect(publisher.address, dialer)
+	address := fmt.Sprintf("%s.ttrpc", addressFlag)
+	conn, err := connect(address, dialer)
 	if err != nil {
 		return err
 	}
-	defer conn.Close()
+	publisher := &remoteEventsPublisher{
+		address: address,
+		conn:    conn,
+		closed:  make(chan struct{}),
+	}
+	defer publisher.Close()
+
 	publisher.client = v1.NewEventsClient(ttrpc.NewClient(conn))
 	if namespaceFlag == "" {
 		return fmt.Errorf("shim namespace cannot be empty")
@@ -172,8 +183,9 @@ func run(id string, initFunc Init, config Config) error {
 	ctx := namespaces.WithNamespace(context.Background(), namespaceFlag)
 	ctx = context.WithValue(ctx, OptsKey{}, Opts{BundlePath: bundlePath, Debug: debugFlag})
 	ctx = log.WithLogger(ctx, log.G(ctx).WithField("runtime", id))
+	ctx, cancel := context.WithCancel(ctx)
 
-	service, err := initFunc(ctx, idFlag, publisher)
+	service, err := initFunc(ctx, idFlag, publisher, cancel)
 	if err != nil {
 		return err
 	}
@@ -183,7 +195,7 @@ func run(id string, initFunc Init, config Config) error {
 			"pid":       os.Getpid(),
 			"namespace": namespaceFlag,
 		})
-		go handleSignals(logger, signals)
+		go handleSignals(ctx, logger, signals)
 		response, err := service.Cleanup(ctx)
 		if err != nil {
 			return err
@@ -210,7 +222,17 @@ func run(id string, initFunc Init, config Config) error {
 			return err
 		}
 		client := NewShimClient(ctx, service, signals)
-		return client.Serve()
+		if err := client.Serve(); err != nil {
+			if err != context.Canceled {
+				return err
+			}
+		}
+		select {
+		case <-publisher.Done():
+			return nil
+		case <-time.After(5 * time.Second):
+			return errors.New("publisher not closed")
+		}
 	}
 }
 
@@ -254,7 +276,7 @@ func (s *Client) Serve() error {
 			dumpStacks(logger)
 		}
 	}()
-	return handleSignals(logger, s.signals)
+	return handleSignals(s.context, logger, s.signals)
 }
 
 // serve serves the ttrpc API over a unix socket at the provided path
@@ -291,7 +313,22 @@ func dumpStacks(logger *logrus.Entry) {
 
 type remoteEventsPublisher struct {
 	address string
+	conn    net.Conn
 	client  v1.EventsService
+	closed  chan struct{}
+	closer  sync.Once
+}
+
+func (l *remoteEventsPublisher) Done() <-chan struct{} {
+	return l.closed
+}
+
+func (l *remoteEventsPublisher) Close() (err error) {
+	l.closer.Do(func() {
+		err = l.conn.Close()
+		close(l.closed)
+	})
+	return err
 }
 
 func (l *remoteEventsPublisher) Publish(ctx context.Context, topic string, event events.Event) error {