Set OnClose shim function

crosbymichael · crosbymichael · commit 87d15a5ffcea · 2018-02-06T10:45:50.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
@@ -80,7 +80,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")
 		}
@@ -149,13 +149,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
`@@ -80,7 +80,7 @@ func WithStart(binary, address, daemonAddress, cgroup string, debug bool, exitHa`
`80`	`80`	`if err = sys.SetOOMScore(cmd.Process.Pid, sys.OOMScoreMaxKillable); err != nil {`
`81`	`81`	`return nil, nil, errors.Wrap(err, "failed to set OOM Score on shim")`
`82`	`82`	`}`
`83`		`- c, clo, err := WithConnect(address)(ctx, config)`
	`83`	`+ c, clo, err := WithConnect(address, func() {})(ctx, config)`
`84`	`84`	`if err != nil {`
`85`	`85`	`return nil, nil, errors.Wrap(err, "failed to connect")`
`86`	`86`	`}`
`@@ -149,13 +149,15 @@ func annonDialer(address string, timeout time.Duration) (net.Conn, error) {`
`149`	`149`	`}`
`150`	`150`
`151`	`151`	`// WithConnect connects to an existing shim`
`152`		`-func WithConnect(address string) Opt {`
	`152`	`+func WithConnect(address string, onClose func()) Opt {`
`153`	`153`	`return func(ctx context.Context, config shim.Config) (shimapi.ShimService, io.Closer, error) {`
`154`	`154`	`conn, err := connect(address, annonDialer)`
`155`	`155`	`if err != nil {`
`156`	`156`	`return nil, nil, err`
`157`	`157`	`}`
`158`		`- return shimapi.NewShimClient(ttrpc.NewClient(conn)), conn, nil`
	`158`	`+ client := ttrpc.NewClient(conn)`
	`159`	`+ client.OnClose(onClose)`
	`160`	`+ return shimapi.NewShimClient(client), conn, nil`
`159`	`161`	`}`
`160`	`162`	`}`
`161`	`163`