Set OnClose shim function

crosbymichael · crosbymichael · commit 04b5f3fbbf64 · 2018-02-06T14:02:33.000-05:00
When restoring a task make sure that dead shims are handled. Fixes containerd#2078 Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
diff --git a/linux/bundle.go b/linux/bundle.go
@@ -84,9 +84,9 @@ func ShimLocal(exchange *exchange.Exchange) ShimOpt {
 }
 
 // ShimConnect is a ShimOpt for connecting to an existing remote shim
-func ShimConnect() ShimOpt {
+func ShimConnect(onClose func()) ShimOpt {
 	return func(b *bundle, ns string, ropts *runctypes.RuncOptions) (shim.Config, client.Opt) {
-		return b.shimConfig(ns, ropts), client.WithConnect(b.shimAddress(ns))
+		return b.shimConfig(ns, ropts), client.WithConnect(b.shimAddress(ns), onClose)
 	}
 }
 
diff --git a/linux/runtime.go b/linux/runtime.go
@@ -383,7 +383,17 @@ func (r *Runtime) loadTasks(ctx context.Context, ns string) ([]*Task, error) {
 		)
 		ctx = namespaces.WithNamespace(ctx, ns)
 		pid, _ := runc.ReadPidFile(filepath.Join(bundle.path, proc.InitPidFile))
-		s, err := bundle.NewShimClient(ctx, ns, ShimConnect(), nil)
+		s, err := bundle.NewShimClient(ctx, ns, ShimConnect(func() {
+			log.G(ctx).WithError(err).WithFields(logrus.Fields{
+				"id":        id,
+				"namespace": ns,
+			}).Error("connecting to shim")
+			err := r.cleanupAfterDeadShim(ctx, bundle, ns, id, pid)
+			if err != nil {
+				log.G(ctx).WithError(err).WithField("bundle", bundle.path).
+					Error("cleaning up after dead shim")
+			}
+		}), nil)
 		if err != nil {
 			log.G(ctx).WithError(err).WithFields(logrus.Fields{
 				"id":        id,
diff --git a/linux/shim/client/client.go b/linux/shim/client/client.go
@@ -77,7 +77,7 @@ func WithStart(binary, address, daemonAddress, cgroup string, debug bool, exitHa
 		if err = sys.SetOOMScore(cmd.Process.Pid, sys.OOMScoreMaxKillable); err != nil {
 			return nil, nil, errors.Wrap(err, "failed to set OOM Score on shim")
 		}
-		c, clo, err := WithConnect(address)(ctx, config)
+		c, clo, err := WithConnect(address, func() {})(ctx, config)
 		if err != nil {
 			return nil, nil, errors.Wrap(err, "failed to connect")
 		}
@@ -146,13 +146,15 @@ func annonDialer(address string, timeout time.Duration) (net.Conn, error) {
 }
 
 // WithConnect connects to an existing shim
-func WithConnect(address string) Opt {
+func WithConnect(address string, onClose func()) Opt {
 	return func(ctx context.Context, config shim.Config) (shimapi.ShimService, io.Closer, error) {
 		conn, err := connect(address, annonDialer)
 		if err != nil {
 			return nil, nil, err
 		}
-		return shimapi.NewShimClient(ttrpc.NewClient(conn)), conn, nil
+		client := ttrpc.NewClient(conn)
+		client.OnClose(onClose)
+		return shimapi.NewShimClient(client), conn, nil
 	}
 }
 

Original file line number	Diff line number	Diff line change
`@@ -84,9 +84,9 @@ func ShimLocal(exchange *exchange.Exchange) ShimOpt {`
`84`	`84`	`}`
`85`	`85`
`86`	`86`	`// ShimConnect is a ShimOpt for connecting to an existing remote shim`
`87`		`-func ShimConnect() ShimOpt {`
	`87`	`+func ShimConnect(onClose func()) ShimOpt {`
`88`	`88`	`return func(b bundle, ns string, ropts runctypes.RuncOptions) (shim.Config, client.Opt) {`
`89`		`- return b.shimConfig(ns, ropts), client.WithConnect(b.shimAddress(ns))`
	`89`	`+ return b.shimConfig(ns, ropts), client.WithConnect(b.shimAddress(ns), onClose)`
`90`	`90`	`}`
`91`	`91`	`}`
`92`	`92`
Original file line number	Diff line number	Diff line change
`@@ -77,7 +77,7 @@ func WithStart(binary, address, daemonAddress, cgroup string, debug bool, exitHa`
`77`	`77`	`if err = sys.SetOOMScore(cmd.Process.Pid, sys.OOMScoreMaxKillable); err != nil {`
`78`	`78`	`return nil, nil, errors.Wrap(err, "failed to set OOM Score on shim")`
`79`	`79`	`}`
`80`		`- c, clo, err := WithConnect(address)(ctx, config)`
	`80`	`+ c, clo, err := WithConnect(address, func() {})(ctx, config)`
`81`	`81`	`if err != nil {`
`82`	`82`	`return nil, nil, errors.Wrap(err, "failed to connect")`
`83`	`83`	`}`
`@@ -146,13 +146,15 @@ func annonDialer(address string, timeout time.Duration) (net.Conn, error) {`
`146`	`146`	`}`
`147`	`147`
`148`	`148`	`// WithConnect connects to an existing shim`
`149`		`-func WithConnect(address string) Opt {`
	`149`	`+func WithConnect(address string, onClose func()) Opt {`
`150`	`150`	`return func(ctx context.Context, config shim.Config) (shimapi.ShimService, io.Closer, error) {`
`151`	`151`	`conn, err := connect(address, annonDialer)`
`152`	`152`	`if err != nil {`
`153`	`153`	`return nil, nil, err`
`154`	`154`	`}`
`155`		`- return shimapi.NewShimClient(ttrpc.NewClient(conn)), conn, nil`
	`155`	`+ client := ttrpc.NewClient(conn)`
	`156`	`+ client.OnClose(onClose)`
	`157`	`+ return shimapi.NewShimClient(client), conn, nil`
`156`	`158`	`}`
`157`	`159`	`}`
`158`	`160`