x/tools/go/packages: fix Load's "argument list too long" error

The current implementation doesn't assume any limit on the amount of
patterns passed to Load which breaks it when called on large codebases.
Different platforms have implemented different values for ARG_MAX,
https://www.in-ulm.de/~mascheck/various/argmax/#results,
https://pubs.opengroup.org/onlinepubs/9699919799/basedefs/limits.h.html,
https://devblogs.microsoft.com/oldnewthing/20031210-00/?p=41553.

In this fix, we use a conservative estimation of ARG_MAX (32767)
that would suit all the supported platforms, subtracted by an approx.
length of env on a typical machine (16384). We use this value to split
the provided patterns into chunks to run the underlying driver on chunks
instead of the whole patterns slice.

Fixes golang/go#36909

Change-Id: I7c462c1e905b1f95013afa2b0ce2aa0a435125dc
Reviewed-on: https://go-review.googlesource.com/c/tools/+/550477
Reviewed-by: Michael Matloob <[email protected]>
Auto-Submit: Michael Matloob <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Matthew Dempsky <[email protected]>

Author: iharsuvorau

go/packages/packages.go

@@ -9,6 +9,7 @@ package packages
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"go/ast"
 	"go/parser"
@@ -24,6 +25,8 @@ import (
 	"sync"
 	"time"
 
+	"golang.org/x/sync/errgroup"
+
 	"golang.org/x/tools/go/gcexportdata"
 	"golang.org/x/tools/internal/gocommand"
 	"golang.org/x/tools/internal/packagesinternal"
@@ -255,8 +258,27 @@ func Load(cfg *Config, patterns ...string) ([]*Package, error) {
 // defaultDriver will fall back to the go list driver.
 // The boolean result indicates that an external driver handled the request.
 func defaultDriver(cfg *Config, patterns ...string) (*DriverResponse, bool, error) {
+	const (
+		// windowsArgMax specifies the maximum command line length for
+		// the Windows' CreateProcess function.
+		windowsArgMax = 32767
+		// maxEnvSize is a very rough estimation of the maximum environment
+		// size of a user.
+		maxEnvSize = 16384
+		// safeArgMax specifies the maximum safe command line length to use
+		// by the underlying driver excl. the environment. We choose the Windows'
+		// ARG_MAX as the starting point because it's one of the lowest ARG_MAX
+		// constants out of the different supported platforms,
+		// e.g., https://www.in-ulm.de/~mascheck/various/argmax/#results.
+		safeArgMax = windowsArgMax - maxEnvSize
+	)
+	chunks, err := splitIntoChunks(patterns, safeArgMax)
+	if err != nil {
+		return nil, false, err
+	}
+
 	if driver := findExternalDriver(cfg); driver != nil {
-		response, err := driver(cfg, patterns...)
+		response, err := callDriverOnChunks(driver, cfg, chunks)
 		if err != nil {
 			return nil, false, err
 		} else if !response.NotHandled {
@@ -265,11 +287,82 @@ func defaultDriver(cfg *Config, patterns ...string) (*DriverResponse, bool, erro
 		// (fall through)
 	}
 
-	response, err := goListDriver(cfg, patterns...)
+	response, err := callDriverOnChunks(goListDriver, cfg, chunks)
 	if err != nil {
 		return nil, false, err
 	}
-	return response, false, nil
+	return response, false, err
+}
+
+// splitIntoChunks chunks the slice so that the total number of characters
+// in a chunk is no longer than argMax.
+func splitIntoChunks(patterns []string, argMax int) ([][]string, error) {
+	if argMax <= 0 {
+		return nil, errors.New("failed to split patterns into chunks, negative safe argMax value")
+	}
+	var chunks [][]string
+	charsInChunk := 0
+	nextChunkStart := 0
+	for i, v := range patterns {
+		vChars := len(v)
+		if vChars > argMax {
+			// a single pattern is longer than the maximum safe ARG_MAX, hardly should happen
+			return nil, errors.New("failed to split patterns into chunks, a pattern is too long")
+		}
+		charsInChunk += vChars + 1 // +1 is for a whitespace between patterns that has to be counted too
+		if charsInChunk > argMax {
+			chunks = append(chunks, patterns[nextChunkStart:i])
+			nextChunkStart = i
+			charsInChunk = vChars
+		}
+	}
+	// add the last chunk
+	if nextChunkStart < len(patterns) {
+		chunks = append(chunks, patterns[nextChunkStart:])
+	}
+	return chunks, nil
+}
+
+func callDriverOnChunks(driver driver, cfg *Config, chunks [][]string) (*DriverResponse, error) {
+	if len(chunks) == 0 {
+		return driver(cfg)
+	}
+	responses := make([]*DriverResponse, len(chunks))
+	errNotHandled := errors.New("driver returned NotHandled")
+	var g errgroup.Group
+	for i, chunk := range chunks {
+		i := i
+		chunk := chunk
+		g.Go(func() (err error) {
+			responses[i], err = driver(cfg, chunk...)
+			if responses[i] != nil && responses[i].NotHandled {
+				err = errNotHandled
+			}
+			return err
+		})
+	}
+	if err := g.Wait(); err != nil {
+		if errors.Is(err, errNotHandled) {
+			return &DriverResponse{NotHandled: true}, nil
+		}
+		return nil, err
+	}
+	return mergeResponses(responses...), nil
+}
+
+func mergeResponses(responses ...*DriverResponse) *DriverResponse {
+	if len(responses) == 0 {
+		return nil
+	}
+	response := newDeduper()
+	response.dr.NotHandled = false
+	response.dr.Compiler = responses[0].Compiler
+	response.dr.Arch = responses[0].Arch
+	response.dr.GoVersion = responses[0].GoVersion
+	for _, v := range responses {
+		response.addAll(v)
+	}
+	return response.dr
 }
 
 // A Package describes a loaded Go package.

go/packages/packages_test.go

@@ -367,6 +367,32 @@ func TestLoadAbsolutePath(t *testing.T) {
 	}
 }
 
+func TestLoadArgumentListIsNotTooLong(t *testing.T) {
+	// NOTE: this test adds about 2s to the test suite running time
+
+	t.Parallel()
+
+	// using the real ARG_MAX for some platforms increases the running time of this test by a lot,
+	// 1_000_000 seems like enough to break Windows and macOS if Load doesn't split provided patterns
+	argMax := 1_000_000
+	exported := packagestest.Export(t, packagestest.GOPATH, []packagestest.Module{{
+		Name: "golang.org/mod",
+		Files: map[string]interface{}{
+			"main.go": `package main"`,
+		}}})
+	defer exported.Cleanup()
+	numOfPatterns := argMax/16 + 1 // the pattern below is approx. 16 chars
+	patterns := make([]string, numOfPatterns)
+	for i := 0; i < numOfPatterns; i++ {
+		patterns[i] = fmt.Sprintf("golang.org/mod/p%d", i)
+	} // patterns have more than argMax number of chars combined with whitespaces b/w patterns
+
+	_, err := packages.Load(exported.Config, patterns...)
+	if err != nil {
+		t.Fatalf("failed to load: %v", err)
+	}
+}
+
 func TestVendorImports(t *testing.T) {
 	t.Parallel()
 
@@ -1308,7 +1334,7 @@ func testNoPatterns(t *testing.T, exporter packagestest.Exporter) {
 
 func TestJSON(t *testing.T) { testAllOrModulesParallel(t, testJSON) }
 func testJSON(t *testing.T, exporter packagestest.Exporter) {
-	//TODO: add in some errors
+	// TODO: add in some errors
 	exported := packagestest.Export(t, exporter, []packagestest.Module{{
 		Name: "golang.org/fake",
 		Files: map[string]interface{}{