add integration test

Author: sreya

cli/docker.go

@@ -7,14 +7,12 @@ import (
 	"io"
 	"net/url"
 	"os"
-	"os/signal"
+	"os/exec"
 	"path"
 	"path/filepath"
 	"sort"
 	"strconv"
 	"strings"
-	"syscall"
-	"time"
 
 	dockertypes "github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
@@ -148,7 +146,7 @@ type flags struct {
 	ethlink       string
 }
 
-func dockerCmd() *cobra.Command {
+func dockerCmd(ch chan func() error) *cobra.Command {
 	var flags flags
 
 	cmd := &cobra.Command{
@@ -287,7 +285,7 @@ func dockerCmd() *cobra.Command {
 				return xerrors.Errorf("wait for dockerd: %w", err)
 			}
 
-			err = runDockerCVM(ctx, log, client, blog, flags)
+			err = runDockerCVM(ctx, log, client, blog, ch, flags)
 			if err != nil {
 				// It's possible we failed because we ran out of disk while
 				// pulling the image. We should restart the daemon and use
@@ -316,7 +314,7 @@ func dockerCmd() *cobra.Command {
 					}()
 
 					log.Debug(ctx, "reattempting container creation")
-					err = runDockerCVM(ctx, log, client, blog, flags)
+					err = runDockerCVM(ctx, log, client, blog, ch, flags)
 				}
 				if err != nil {
 					blog.Errorf("Failed to run envbox: %v", err)
@@ -359,7 +357,7 @@ func dockerCmd() *cobra.Command {
 	return cmd
 }
 
-func runDockerCVM(ctx context.Context, log slog.Logger, client dockerutil.DockerClient, blog buildlog.Logger, flags flags) error {
+func runDockerCVM(ctx context.Context, log slog.Logger, client dockerutil.DockerClient, blog buildlog.Logger, shutdownCh chan func() error, flags flags) error {
 	fs := xunix.GetFS(ctx)
 
 	// Set our OOM score to something really unfavorable to avoid getting killed
@@ -692,6 +690,7 @@ func runDockerCVM(ctx context.Context, log slog.Logger, client dockerutil.Docker
 	if err != nil {
 		return xerrors.Errorf("create exec: %w", err)
 	}
+
 	resp, err := client.ContainerExecAttach(ctx, bootstrapExec.ID, dockertypes.ExecStartCheck{})
 	if err != nil {
 		return xerrors.Errorf("attach exec: %w", err)
@@ -715,53 +714,31 @@ func runDockerCVM(ctx context.Context, log slog.Logger, client dockerutil.Docker
 		}
 	}()
 
-	inspect, err := client.ContainerExecInspect(ctx, bootstrapExec.ID)
+	// We can't just call ExecInspect because there's a race where the cmd
+	// hasn't been assigned a PID yet.
+	bootstrapPID, err := dockerutil.GetExecPID(ctx, client, bootstrapExec.ID)
 	if err != nil {
 		return xerrors.Errorf("exec inspect: %w", err)
 	}
-	sigs := make(chan os.Signal, 1)
-	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
-	go func() {
-		err := func() error {
-			log.Info(ctx, "waiting for signal")
-			sig := <-sigs
-			sigstr := "TERM"
-			if sig == syscall.SIGINT {
-				sigstr = "INT"
-			}
-			log.Debug(ctx, "received signal", slog.F("signal", sigstr))
 
-			killExec, err := client.ContainerExecCreate(ctx, containerID, dockertypes.ExecConfig{
-				User:         imgMeta.UID,
-				Cmd:          []string{"sh", "-c", fmt.Sprintf("kill -%s %d", sigstr, inspect.Pid)},
-				AttachStdout: true,
-				AttachStderr: true,
-			})
-			if err != nil {
-				return xerrors.Errorf("create kill exec: %w", err)
-			}
-
-			err = dockerutil.WaitForExit(ctx, client, killExec.ID)
-			if err != nil {
-				return xerrors.Errorf("wait for kill exec to complete: %w", err)
-			}
+	shutdownCh <- func() error {
+		log.Debug(ctx, "killing container", slog.F("bootstrap_pid", bootstrapPID))
 
-			err = dockerutil.WaitForExit(ctx, client, bootstrapExec.ID)
-			if err != nil {
-				return xerrors.Errorf("wait for exit: %w", err)
-			}
+		// The PID returned is the PID _outside_ the container...
+		out, err := exec.Command("kill", "-TERM", strconv.Itoa(bootstrapPID)).CombinedOutput()
+		if err != nil {
+			return xerrors.Errorf("kill bootstrap process (%s): %w", out, err)
+		}
 
-			return nil
-		}()
-		log.Info(ctx, "exiting envbox", slog.Error(err))
-		log.Sync()
+		log.Debug(ctx, "sent kill signal waiting for process to exit")
+		err = dockerutil.WaitForExit(ctx, client, bootstrapExec.ID)
 		if err != nil {
-			os.Exit(1)
+			return xerrors.Errorf("wait for exit: %w", err)
 		}
-		os.Exit(0)
-	}()
-	log.Info(ctx, "HELP")
-	time.Sleep(time.Second)
+
+		log.Debug(ctx, "bootstrap process successfully exited")
+		return nil
+	}
 
 	return nil
 }

cli/root.go

@@ -4,7 +4,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func Root() *cobra.Command {
+func Root(ch chan func() error) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:           "envbox",
 		SilenceErrors: true,
@@ -15,6 +15,6 @@ func Root() *cobra.Command {
 		},
 	}
 
-	cmd.AddCommand(dockerCmd())
+	cmd.AddCommand(dockerCmd(ch))
 	return cmd
 }

cmd/envbox/main.go

@@ -3,18 +3,37 @@ package main
 import (
 	"fmt"
 	"os"
+	"os/signal"
 	"runtime"
+	"syscall"
 
 	"github.com/coder/envbox/cli"
 )
 
 func main() {
-	_, err := cli.Root().ExecuteC()
+	ch := make(chan func() error, 1)
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGTERM, syscall.SIGINT, syscall.SIGWINCH)
+	go func() {
+		fmt.Println("waiting for signal")
+		<-sigs
+		fmt.Println("Got signal")
+		select {
+		case fn := <-ch:
+			fmt.Println("running shutdown function")
+			err := fn()
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "shutdown function failed: %v", err)
+				os.Exit(1)
+			}
+		default:
+			fmt.Println("no shutdown function")
+		}
+		os.Exit(0)
+	}()
+	_, err := cli.Root(ch).ExecuteC()
 	if err != nil {
-		_, _ = fmt.Fprintln(os.Stderr, err.Error())
 		os.Exit(1)
 	}
-
-	// We exit the main thread while keepin all the other procs goin strong.
 	runtime.Goexit()
 }

dockerutil/exec.go

@@ -95,6 +95,23 @@ func ExecContainer(ctx context.Context, client DockerClient, config ExecConfig)
 	return &buf, inspect.Pid, nil
 }
 
+func GetExecPID(ctx context.Context, client DockerClient, execID string) (int, error) {
+	for r := retry.New(time.Second, time.Second); r.Wait(ctx); {
+		inspect, err := client.ContainerExecInspect(ctx, execID)
+		if err != nil {
+			return 0, xerrors.Errorf("exec inspect: %w", err)
+		}
+
+		if inspect.Pid == 0 {
+			continue
+		}
+		return inspect.Pid, nil
+	}
+
+	return 0, ctx.Err()
+
+}
+
 func WaitForExit(ctx context.Context, client DockerClient, execID string) error {
 	for r := retry.New(time.Second, time.Second); r.Wait(ctx); {
 		inspect, err := client.ContainerExecInspect(ctx, execID)
@@ -112,6 +129,5 @@ func WaitForExit(ctx context.Context, client DockerClient, execID string) error
 
 		return nil
 	}
-
 	return ctx.Err()
 }

integration/docker_test.go

@@ -289,10 +289,12 @@ func TestDocker(t *testing.T) {
 
 		binds = append(binds, bindMount(homeDir, "/home/coder", false))
 
+		envs := []string{fmt.Sprintf("%s=%s:%s", cli.EnvMounts, "/home/coder", "/home/coder")}
 		// Run the envbox container.
 		resource := integrationtest.RunEnvbox(t, pool, &integrationtest.CreateDockerCVMConfig{
 			Image:           integrationtest.UbuntuImage,
 			Username:        "root",
+			Envs:            envs,
 			Binds:           binds,
 			BootstrapScript: sigtrapScript,
 		})
@@ -303,9 +305,8 @@ func TestDocker(t *testing.T) {
 		})
 		require.Error(t, err)
 
-		time.Sleep(time.Second * 5)
-		err = pool.Client.StopContainer(resource.Container.ID, 30)
-		require.NoError(t, err)
+		// Simulate a shutdown.
+		integrationtest.StopContainer(t, pool, resource.Container.ID, 30*time.Second)
 
 		err = resource.Close()
 		require.NoError(t, err)
@@ -315,6 +316,7 @@ func TestDocker(t *testing.T) {
 		resource = integrationtest.RunEnvbox(t, pool, &integrationtest.CreateDockerCVMConfig{
 			Image:           integrationtest.UbuntuImage,
 			Username:        "root",
+			Envs:            envs,
 			Binds:           binds,
 			BootstrapScript: sigtrapScript,
 		})
@@ -357,17 +359,13 @@ func bindMount(src, dest string, ro bool) string {
 }
 
 const sigtrapScript = `#!/bin/bash
-
-# Function to handle cleanup
 cleanup() {
-	touch /home/coder/foo
+	echo "HANDLING A SIGNAL!" && touch /home/coder/foo && echo "touched file"
     exit 0
 }
 
-# Trap SIGINT (Ctrl+C) and SIGTERM signals
 trap 'cleanup' INT TERM
 
-# Main loop or processing logic (replace with your script's logic)
 while true; do
     echo "Working..."
     sleep 1

integration/integrationtest/docker.go

@@ -33,7 +33,7 @@ const (
 	HelloWorldImage = "gcr.io/coder-dev-1/sreya/hello-world"
 	// UbuntuImage is just vanilla ubuntu (80MB) but the user is set to a non-root
 	// user .
-	UbuntuImage = "gcr.io/coder-dev-1/sreya/ubuntu-coder"
+	UbuntuImage = "gcr.io/coder-dev-1/sreya/ubuntu-coder:jon"
 )
 
 // TODO use df to determine if an environment is running in a docker container or not.
@@ -264,6 +264,31 @@ func ExecEnvbox(t *testing.T, pool *dockertest.Pool, conf ExecConfig) ([]byte, e
 	return buf.Bytes(), nil
 }
 
+func StopContainer(t *testing.T, pool *dockertest.Pool, id string, to time.Duration) {
+	t.Helper()
+
+	err := pool.Client.KillContainer(docker.KillContainerOptions{
+		ID:     id,
+		Signal: docker.SIGTERM,
+	})
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), to)
+	defer cancel()
+	for r := retry.New(time.Second, time.Second); r.Wait(ctx); {
+		cnt, err := pool.Client.InspectContainer(id)
+		require.NoError(t, err)
+
+		if cnt.State.Running {
+			continue
+		}
+
+		return
+	}
+
+	t.Fatalf("timed out waiting for container %s to stop", id)
+}
+
 // cmdLineEnvs returns args passed to the /envbox command
 // but using their env var alias.
 func cmdLineEnvs(c *CreateDockerCVMConfig) []string {