Merge pull request #554 from xhd2015/accelerate_go_list

Accelerating go-extractor by using 'go list -deps' instead of just 'go list'

extractor/extractor.go

@@ -102,13 +102,19 @@ func ExtractWithFlags(buildFlags []string, patterns []string) error {
 	extractUniverseScope()
 	log.Println("Done extracting universe scope.")
 
-	// a map of package path to package root directory (currently the module root or the source directory)
-	pkgRoots := make(map[string]string)
-	// a map of package path to source code directory
-	pkgDirs := make(map[string]string)
+	// a map of package path to source directory and module root directory
+	pkgInfos := make(map[string]*util.PkgInfo)
 	// root directories of packages that we want to extract
 	wantedRoots := make(map[string]bool)
 
+	log.Printf("Running go list to resolve package and module directories.")
+	// get all packages information
+	pkgInfos, err = util.GetPkgsInfo(patterns, true, modFlags...)
+	if err != nil {
+		log.Fatalf("Error getting dependency package or module directories: %v.", err)
+	}
+	log.Printf("Done running go list deps: resolved %d packages.", len(pkgInfos))
+
 	// recursively visit all packages in depth-first order;
 	// on the way down, associate each package scope with its corresponding package,
 	// and on the way up extract the package's scope
@@ -117,18 +123,6 @@ func ExtractWithFlags(buildFlags []string, patterns []string) error {
 	}, func(pkg *packages.Package) {
 		log.Printf("Processing package %s.", pkg.PkgPath)
 
-		if _, ok := pkgRoots[pkg.PkgPath]; !ok {
-			mdir := util.GetModDir(pkg.PkgPath, modFlags...)
-			pdir := util.GetPkgDir(pkg.PkgPath, modFlags...)
-			// GetModDir returns the empty string if the module directory cannot be determined, e.g. if the package
-			// is not using modules. If this is the case, fall back to the package directory
-			if mdir == "" {
-				mdir = pdir
-			}
-			pkgRoots[pkg.PkgPath] = mdir
-			pkgDirs[pkg.PkgPath] = pdir
-		}
-
 		log.Printf("Extracting types for package %s.", pkg.PkgPath)
 
 		tw, err := trap.NewWriter(pkg.PkgPath, pkg)
@@ -153,11 +147,21 @@ func ExtractWithFlags(buildFlags []string, patterns []string) error {
 	})
 
 	for _, pkg := range pkgs {
-		if pkgRoots[pkg.PkgPath] == "" {
+		info := pkgInfos[pkg.PkgPath]
+		if info == nil {
+			var err error
+			info, err = util.GetPkgInfo(pkg.PkgPath)
+			if err != nil {
+				log.Fatalf("Unable to get a source directory for input package %s: %s", pkg.PkgPath, err)
+			}
+		}
+		if info == nil || info.PkgDir == "" {
 			log.Fatalf("Unable to get a source directory for input package %s.", pkg.PkgPath)
 		}
-		wantedRoots[pkgRoots[pkg.PkgPath]] = true
-		wantedRoots[pkgDirs[pkg.PkgPath]] = true
+		wantedRoots[info.PkgDir] = true
+		if info.ModDir != "" {
+			wantedRoots[info.ModDir] = true
+		}
 	}
 
 	log.Println("Done processing dependencies.")
@@ -175,7 +179,9 @@ func ExtractWithFlags(buildFlags []string, patterns []string) error {
 		return true
 	}, func(pkg *packages.Package) {
 		for root, _ := range wantedRoots {
-			relDir, err := filepath.Rel(root, pkgDirs[pkg.PkgPath])
+			info := pkgInfos[pkg.PkgPath]
+
+			relDir, err := filepath.Rel(root, info.PkgDir)
 			if err != nil || noExtractRe.MatchString(relDir) {
 				// if the path can't be made relative or matches the noExtract regexp skip it
 				continue
@@ -183,8 +189,8 @@ func ExtractWithFlags(buildFlags []string, patterns []string) error {
 
 			extraction.extractPackage(pkg)
 
-			if pkgRoots[pkg.PkgPath] != "" {
-				modPath := filepath.Join(pkgRoots[pkg.PkgPath], "go.mod")
+			if info.ModDir != "" {
+				modPath := filepath.Join(info.ModDir, "go.mod")
 				if util.FileExists(modPath) {
 					log.Printf("Extracting %s", modPath)
 					start := time.Now()

extractor/util/util.go

@@ -1,7 +1,9 @@
 package util
 
 import (
+	"encoding/json"
 	"errors"
+	"io"
 	"log"
 	"os"
 	"os/exec"
@@ -12,9 +14,9 @@ import (
 
 var extractorPath string
 
-// Getenv retrieves the value of the environment variable named by the key.
-// If that variable is not present, it iterates over the given aliases until
-// it finds one that is. If none are present, the empty string is returned.
+// Getenv retrieves the value of the environment variable named by the key.  If that variable is not
+// present, it iterates over the given aliases until it finds one that is. If none are present, the
+// empty string is returned.
 func Getenv(key string, aliases ...string) string {
 	val := os.Getenv(key)
 	if val != "" {
@@ -30,16 +32,12 @@ func Getenv(key string, aliases ...string) string {
 }
 
 // runGoList is a helper function for running go list with format `format` and flags `flags` on
-// package `pkgpath`.
-func runGoList(format string, pkgpath string, flags ...string) (string, error) {
-	return runGoListWithEnv(format, pkgpath, nil, flags...)
-}
-
-func runGoListWithEnv(format string, pkgpath string, additionalEnv []string, flags ...string) (string, error) {
+// the packages defined by `patterns`.
+func runGoList(format string, patterns []string, flags ...string) (string, error) {
 	args := append([]string{"list", "-e", "-f", format}, flags...)
-	args = append(args, pkgpath)
+	args = append(args, patterns...)
 	cmd := exec.Command("go", args...)
-	cmd.Env = append(os.Environ(), additionalEnv...)
+	log.Printf("Running cmd: %s", cmd.String())
 	out, err := cmd.Output()
 
 	if err != nil {
@@ -54,50 +52,82 @@ func runGoListWithEnv(format string, pkgpath string, additionalEnv []string, fla
 	return strings.TrimSpace(string(out)), nil
 }
 
-// GetModDir gets the absolute directory of the module containing the package with path
-// `pkgpath`. It passes the `go list` the flags specified by `flags`.
-func GetModDir(pkgpath string, flags ...string) string {
-	// enable module mode so that we can find a module root if it exists, even if go module support is
-	// disabled by a build
-	mod, err := runGoListWithEnv("{{.Module}}", pkgpath, []string{"GO111MODULE=on"}, flags...)
-	if err != nil || mod == "<nil>" {
-		// if the command errors or modules aren't being used, return the empty string
-		return ""
-	}
-
-	modDir, err := runGoListWithEnv("{{.Module.Dir}}", pkgpath, []string{"GO111MODULE=on"}, flags...)
-	if err != nil {
-		return ""
-	}
-
-	abs, err := filepath.Abs(modDir)
-	if err != nil {
-		log.Printf("Warning: unable to make %s absolute: %s", modDir, err.Error())
-		return ""
-	}
-	return abs
+// PkgInfo holds package directory and module directory (if any) for a package
+type PkgInfo struct {
+	PkgDir string // the directory directly containing source code of this package
+	ModDir string // the module directory containing this package, empty if not a module
 }
 
-// GetPkgDir gets the absolute directory containing the package with path `pkgpath`. It passes the
-// `go list` command the flags specified by `flags`.
-func GetPkgDir(pkgpath string, flags ...string) string {
-	pkgDir, err := runGoList("{{.Dir}}", pkgpath, flags...)
-	if err != nil {
-		return ""
+// GetPkgsInfo gets the absolute module and package root directories for the packages matched by the
+// patterns `patterns`. It passes to `go list` the flags specified by `flags`.  If `includingDeps`
+// is true, all dependencies will also be included.
+func GetPkgsInfo(patterns []string, includingDeps bool, flags ...string) (map[string]*PkgInfo, error) {
+	// enable module mode so that we can find a module root if it exists, even if go module support is
+	// disabled by a build
+	if includingDeps {
+		// the flag `-deps` causes all dependencies to be retrieved
+		flags = append(flags, "-deps")
 	}
 
-	abs, err := filepath.Abs(pkgDir)
+	// using -json overrides -f format
+	output, err := runGoList("", patterns, append(flags, "-json")...)
 	if err != nil {
-		log.Printf("Warning: unable to make %s absolute: %s", pkgDir, err.Error())
-		return ""
+		return nil, err
 	}
-	return abs
+
+	// the output of `go list -json` is a stream of json object
+	type goListPkgInfo struct {
+		ImportPath string
+		Dir        string
+		Module     *struct {
+			Dir string
+		}
+	}
+	pkgInfoMapping := make(map[string]*PkgInfo)
+	streamDecoder := json.NewDecoder(strings.NewReader(output))
+	for {
+		var pkgInfo goListPkgInfo
+		decErr := streamDecoder.Decode(&pkgInfo)
+		if decErr == io.EOF {
+			break
+		}
+		if decErr != nil {
+			log.Printf("Error decoding output of go list -json: %s", err.Error())
+			return nil, decErr
+		}
+		pkgAbsDir, err := filepath.Abs(pkgInfo.Dir)
+		if err != nil {
+			log.Printf("Unable to make package dir %s absolute: %s", pkgInfo.Dir, err.Error())
+		}
+		var modAbsDir string
+		if pkgInfo.Module != nil {
+			modAbsDir, err = filepath.Abs(pkgInfo.Module.Dir)
+			if err != nil {
+				log.Printf("Unable to make module dir %s absolute: %s", pkgInfo.Module.Dir, err.Error())
+			}
+		}
+		pkgInfoMapping[pkgInfo.ImportPath] = &PkgInfo{
+			PkgDir: pkgAbsDir,
+			ModDir: modAbsDir,
+		}
+	}
+	return pkgInfoMapping, nil
+}
+
+// GetPkgsInfo gets the absolute module and package root directories for the package `pkg`, passing
+// the internal `go list` command the flags specified by `flags`.
+func GetPkgInfo(pkg string, flags ...string) (*PkgInfo, error) {
+	pkgInfos, err := GetPkgsInfo([]string{pkg}, false, flags...)
+	if err != nil {
+		return nil, err
+	}
+	return pkgInfos[pkg], nil
 }
 
 // DepErrors checks there are any errors resolving dependencies for `pkgpath`. It passes the `go
 // list` command the flags specified by `flags`.
 func DepErrors(pkgpath string, flags ...string) bool {
-	out, err := runGoList("{{if .DepsErrors}}{{else}}error{{end}}", pkgpath, flags...)
+	out, err := runGoList("{{if .DepsErrors}}{{else}}error{{end}}", []string{pkgpath}, flags...)
 	if err != nil {
 		// if go list failed, assume dependencies are broken
 		return false