executor: improve hosts cleanup

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

executor/containerdexecutor/executor.go

@@ -3,7 +3,6 @@ package containerdexecutor
 import (
 	"context"
 	"io"
-	"os"
 	"syscall"
 	"time"
 
@@ -40,12 +39,12 @@ func (w containerdExecutor) Exec(ctx context.Context, meta executor.Meta, root c
 		return err
 	}
 
-	hostsFile, err := oci.GetHostsFile(ctx, w.root, meta.ExtraHosts)
+	hostsFile, clean, err := oci.GetHostsFile(ctx, w.root, meta.ExtraHosts)
 	if err != nil {
 		return err
 	}
-	if len(meta.ExtraHosts) > 0 {
-		defer os.RemoveAll(hostsFile)
+	if clean != nil {
+		defer clean()
 	}
 
 	mountable, err := root.Mount(ctx, false)

executor/oci/hosts.go

@@ -17,51 +17,53 @@ const hostsContent = `
 ::1	localhost ip6-localhost ip6-loopback
 `
 
-func GetHostsFile(ctx context.Context, stateDir string, extraHosts []executor.HostIP) (string, error) {
+func GetHostsFile(ctx context.Context, stateDir string, extraHosts []executor.HostIP) (string, func(), error) {
 	if len(extraHosts) == 0 {
 		_, err := g.Do(ctx, stateDir, func(ctx context.Context) (interface{}, error) {
-			_, err := makeHostsFile(stateDir, nil)
+			_, _, err := makeHostsFile(stateDir, nil)
 			return nil, err
 		})
 		if err != nil {
-			return "", err
+			return "", nil, err
 		}
-		return filepath.Join(stateDir, "hosts"), nil
+		return filepath.Join(stateDir, "hosts"), func() {}, nil
 	}
 	return makeHostsFile(stateDir, extraHosts)
 }
 
-func makeHostsFile(stateDir string, extraHosts []executor.HostIP) (string, error) {
+func makeHostsFile(stateDir string, extraHosts []executor.HostIP) (string, func(), error) {
 	p := filepath.Join(stateDir, "hosts")
 	if len(extraHosts) != 0 {
 		p += "." + identity.NewID()
 	}
 	_, err := os.Stat(p)
 	if err == nil {
-		return "", nil
+		return "", func() {}, nil
 	}
 	if !os.IsNotExist(err) {
-		return "", err
+		return "", nil, err
 	}
 
 	b := &bytes.Buffer{}
 
 	if _, err := b.Write([]byte(hostsContent)); err != nil {
-		return "", err
+		return "", nil, err
 	}
 
 	for _, h := range extraHosts {
 		if _, err := b.Write([]byte(fmt.Sprintf("%s\t%s\n", h.IP.String(), h.Host))); err != nil {
-			return "", err
+			return "", nil, err
 		}
 	}
 
 	if err := ioutil.WriteFile(p+".tmp", b.Bytes(), 0644); err != nil {
-		return "", err
+		return "", nil, err
 	}
 
 	if err := os.Rename(p+".tmp", p); err != nil {
-		return "", err
+		return "", nil, err
 	}
-	return p, nil
+	return p, func() {
+		os.RemoveAll(p)
+	}, nil
 }

executor/runcexecutor/executor.go

@@ -103,12 +103,12 @@ func (w *runcExecutor) Exec(ctx context.Context, meta executor.Meta, root cache.
 		return err
 	}
 
-	hostsFile, err := oci.GetHostsFile(ctx, w.root, meta.ExtraHosts)
+	hostsFile, clean, err := oci.GetHostsFile(ctx, w.root, meta.ExtraHosts)
 	if err != nil {
 		return err
 	}
-	if len(meta.ExtraHosts) > 0 {
-		defer os.RemoveAll(hostsFile)
+	if clean != nil {
+		defer clean()
 	}
 
 	mountable, err := root.Mount(ctx, false)