3 files changed, 40 insertions, 91 deletions
diff --git a/test/runtimes/go/BUILD b/test/runtimes/go/BUILD
index c34f49ea6..ce971ee9d 100644
--- a/test/runtimes/go/BUILD
+++ b/test/runtimes/go/BUILD
@@ -5,4 +5,5 @@ package(licenses = ["notice"])
 go_binary(
     name = "proctor-go",
     srcs = ["proctor-go.go"],
+    deps = ["//test/runtimes/common"],
 )
diff --git a/test/runtimes/go/Dockerfile b/test/runtimes/go/Dockerfile
index cd55608cd..2d3477392 100644
--- a/test/runtimes/go/Dockerfile
+++ b/test/runtimes/go/Dockerfile
@@ -23,9 +23,13 @@ ENV LANG_DIR=${GOROOT}
 
 WORKDIR ${LANG_DIR}/src
 RUN ./make.bash
+# Pre-compile the tests for faster execution
+RUN ["/root/go/bin/go", "tool", "dist", "test", "-compile-only"]
 
 WORKDIR ${LANG_DIR}
 
-COPY proctor-go.go ${LANG_DIR}
+COPY common /root/go/src/gvisor.dev/gvisor/test/runtimes/common/common
+COPY go/proctor-go.go ${LANG_DIR}
+RUN ["/root/go/bin/go", "build", "-o", "/root/go/bin/proctor", "proctor-go.go"]
 
-ENTRYPOINT ["/root/go/bin/go", "run", "proctor-go.go"]
+ENTRYPOINT ["/root/go/bin/proctor"]
diff --git a/test/runtimes/go/proctor-go.go b/test/runtimes/go/proctor-go.go
index c5387e21d..3eb24576e 100644
--- a/test/runtimes/go/proctor-go.go
+++ b/test/runtimes/go/proctor-go.go
@@ -21,7 +21,6 @@
 package main
 
 import (
-	"flag"
 	"fmt"
 	"log"
 	"os"
@@ -29,133 +28,78 @@ import (
 	"path/filepath"
 	"regexp"
 	"strings"
+
+	"gvisor.dev/gvisor/test/runtimes/common"
 )
 
 var (
-	list    = flag.Bool("list", false, "list all available tests")
-	test    = flag.String("test", "", "run a single test from the list of available tests")
-	version = flag.Bool("v", false, "print out the version of node that is installed")
-
 	dir       = os.Getenv("LANG_DIR")
+	goBin     = filepath.Join(dir, "bin/go")
 	testDir   = filepath.Join(dir, "test")
 	testRegEx = regexp.MustCompile(`^.+\.go$`)
 
 	// Directories with .dir contain helper files for tests.
 	// Exclude benchmarks and stress tests.
-	exclDirs = regexp.MustCompile(`^.+\/(bench|stress)\/.+$|^.+\.dir.+$`)
+	dirFilter = regexp.MustCompile(`^(bench|stress)\/.+$|^.+\.dir.+$`)
 )
 
-func main() {
-	flag.Parse()
+type goRunner struct {
+}
 
-	if *list && *test != "" {
-		flag.PrintDefaults()
-		os.Exit(1)
-	}
-	if *list {
-		tests, err := listTests()
-		if err != nil {
-			log.Fatalf("Failed to list tests: %v", err)
-		}
-		for _, test := range tests {
-			fmt.Println(test)
-		}
-		return
-	}
-	if *version {
-		fmt.Println("Go version: ", os.Getenv("LANG_VER"), "  is installed.")
-		return
-	}
-	if *test != "" {
-		runTest(*test)
-		return
+func main() {
+	if err := common.LaunchFunc(goRunner{}); err != nil {
+		log.Fatalf("Failed to start: %v", err)
 	}
-	runAllTests()
 }
 
-func listTests() ([]string, error) {
+func (g goRunner) ListTests() ([]string, error) {
 	// Go tool dist test tests.
 	args := []string{"tool", "dist", "test", "-list"}
 	cmd := exec.Command(filepath.Join(dir, "bin/go"), args...)
 	cmd.Stderr = os.Stderr
 	out, err := cmd.Output()
 	if err != nil {
-		log.Fatalf("Failed to list: %v", err)
+		return nil, fmt.Errorf("failed to list: %v", err)
 	}
-	var testSlice []string
+	var toolSlice []string
 	for _, test := range strings.Split(string(out), "\n") {
-		testSlice = append(testSlice, test)
+		toolSlice = append(toolSlice, test)
 	}
 
 	// Go tests on disk.
-	if err := filepath.Walk(testDir, func(path string, info os.FileInfo, err error) error {
-		name := filepath.Base(path)
-
-		if info.IsDir() {
-			return nil
-		}
-
-		if !testRegEx.MatchString(name) {
-			return nil
-		}
-
-		if exclDirs.MatchString(path) {
-			return nil
+	diskSlice, err := common.Search(testDir, testRegEx)
+	if err != nil {
+		return nil, err
+	}
+	// Remove items from /bench/, /stress/ and .dir files
+	diskFiltered := diskSlice[:0]
+	for _, file := range diskSlice {
+		if !dirFilter.MatchString(file) {
+			diskFiltered = append(diskFiltered, file)
 		}
-
-		testSlice = append(testSlice, path)
-		return nil
-	}); err != nil {
-		return nil, fmt.Errorf("walking %q: %v", testDir, err)
 	}
 
-	return testSlice, nil
+	return append(toolSlice, diskFiltered...), nil
 }
 
-func runTest(test string) {
-	toolArgs := []string{
-		"tool",
-		"dist",
-		"test",
-	}
-	diskArgs := []string{
-		"run",
-		"run.go",
-		"-v",
-	}
+func (g goRunner) RunTest(test string) error {
 	// Check if test exists on disk by searching for file of the same name.
 	// This will determine whether or not it is a Go test on disk.
-	if _, err := os.Stat(test); err == nil {
-		relPath, err := filepath.Rel(testDir, test)
-		if err != nil {
-			log.Fatalf("Failed to get rel path: %v", err)
-		}
-		diskArgs = append(diskArgs, "--", relPath)
-		cmd := exec.Command(filepath.Join(dir, "bin/go"), diskArgs...)
+	if strings.HasSuffix(test, ".go") {
+		// Test has suffix ".go" which indicates a disk test, run it as such.
+		cmd := exec.Command(goBin, "run", "run.go", "-v", "--", test)
 		cmd.Dir = testDir
 		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
 		if err := cmd.Run(); err != nil {
-			log.Fatalf("Failed to run: %v", err)
+			return fmt.Errorf("failed to run test: %v", err)
 		}
-	} else if os.IsNotExist(err) {
-		// File was not found, try running as Go tool test.
-		toolArgs = append(toolArgs, "-run", test)
-		cmd := exec.Command(filepath.Join(dir, "bin/go"), toolArgs...)
+	} else {
+		// No ".go" suffix, run as a tool test.
+		cmd := exec.Command(goBin, "tool", "dist", "test", "-run", test)
 		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
 		if err := cmd.Run(); err != nil {
-			log.Fatalf("Failed to run: %v", err)
+			return fmt.Errorf("failed to run test: %v", err)
 		}
-	} else {
-		log.Fatalf("Error searching for test: %v", err)
-	}
-}
-
-func runAllTests() {
-	tests, err := listTests()
-	if err != nil {
-		log.Fatalf("Failed to list tests: %v", err)
-	}
-	for _, test := range tests {
-		runTest(test)
 	}
+	return nil
 }