nogo: enable bazel workers and other optimizations.

This is a suite of changes intended to dramatically speed up nogo speed.

First, there are minor changes that help efficiency significantly.

* Gob-based encoding is used internally, and JSON only used for the final
  set of findings. This is done to preserve the existing format (which is
  consumed by external tooling), and to facilitate manual debugging.

* Unnecessary regex compilation is elided in the configuration, and care is
  taken for merges to prevent redundant entries. I'm not sure quite sure how,
  but it turns out that this was consumed a significant amount of time,
  presumably compiling the same regexes over and over again.

Second, this change enables bazel workers for nogo analyzers.

Workers enable persistent processes instead of creating and tearing down a
sandbox every invocation. A library is introduced to abstraction these details,
and allow the tools to still be written using standard flags, etc.

The key here is that these binaries and the core of nogo become aware of
caches with worker.Cache. This allows us to save significant time loading the
same set of files and findings over and over again. These caches are keyed by
the digests that are provided by bazel, and are capped in overall size.

Note that the worker package attempts to capture output during each run, but
tools are no longer permitted to write to stdout. This necessitated dropping
some spurious output from checklocks.

PiperOrigin-RevId: 370505732

1 files changed, 325 insertions, 0 deletions

diff --git a/tools/worker/worker.go b/tools/worker/worker.go
new file mode 100644
index 000000000..669a5f203
--- /dev/null
+++ b/tools/worker/worker.go
@@ -0,0 +1,325 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package worker provides an implementation of the bazel worker protocol.
+//
+// Tools may be written as a normal command line utility, except the passed
+// run function may be invoked multiple times.
+package worker
+
+import (
+	"bufio"
+	"bytes"
+	"flag"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"log"
+	"net"
+	"net/http"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+	"time"
+
+	_ "net/http/pprof" // For profiling.
+
+	"golang.org/x/sys/unix"
+	"google.golang.org/protobuf/encoding/protowire"
+	"google.golang.org/protobuf/proto"
+	wpb "gvisor.dev/bazel/worker_protocol_go_proto"
+)
+
+var (
+	persistentWorker  = flag.Bool("persistent_worker", false, "enable persistent worker.")
+	workerDebug       = flag.Bool("worker_debug", false, "debug persistent workers.")
+	maximumCacheUsage = flag.Int64("maximum_cache_usage", 1024*1024*1024, "maximum cache size.")
+)
+
+var (
+	// inputFiles is the last set of input files.
+	//
+	// This is used for cache invalidation. The key is the *absolute* path
+	// name, and the value is the digest in the current run.
+	inputFiles = make(map[string]string)
+
+	// activeCaches is the set of active caches.
+	activeCaches = make(map[*Cache]struct{})
+
+	// totalCacheUsage is the total usage of all caches.
+	totalCacheUsage int64
+)
+
+// mustAbs returns the absolute path of a filename or dies.
+func mustAbs(filename string) string {
+	abs, err := filepath.Abs(filename)
+	if err != nil {
+		log.Fatalf("error getting absolute path: %v", err)
+	}
+	return abs
+}
+
+// updateInputFiles creates an entry in inputFiles.
+func updateInputFile(filename, digest string) {
+	inputFiles[mustAbs(filename)] = digest
+}
+
+// Sizer returns a size.
+type Sizer interface {
+	Size() int64
+}
+
+// CacheBytes is an example of a Sizer.
+type CacheBytes []byte
+
+// Size implements Sizer.Size.
+func (cb CacheBytes) Size() int64 {
+	return int64(len(cb))
+}
+
+// Cache is a worker cache.
+//
+// They can be created via NewCache.
+type Cache struct {
+	name    string
+	entries map[string]Sizer
+	size    int64
+	hits    int64
+	misses  int64
+}
+
+// NewCache returns a new cache.
+func NewCache(name string) *Cache {
+	return &Cache{
+		name: name,
+	}
+}
+
+// Lookup looks up an entry in the cache.
+//
+// It is a function of the given files.
+func (c *Cache) Lookup(filenames []string, generate func() Sizer) Sizer {
+	digests := make([]string, 0, len(filenames))
+	for _, filename := range filenames {
+		digest, ok := inputFiles[mustAbs(filename)]
+		if !ok {
+			// This is not a valid input. We may not be running as
+			// persistent worker in this cache. If that's the case,
+			// then the file's contents will not change across the
+			// run, and we just use the filename itself.
+			digest = filename
+		}
+		digests = append(digests, digest)
+	}
+
+	// Attempt the lookup.
+	sort.Slice(digests, func(i, j int) bool {
+		return digests[i] < digests[j]
+	})
+	cacheKey := strings.Join(digests, "+")
+	if c.entries == nil {
+		c.entries = make(map[string]Sizer)
+		activeCaches[c] = struct{}{}
+	}
+	entry, ok := c.entries[cacheKey]
+	if ok {
+		c.hits++
+		return entry
+	}
+
+	// Generate a new entry.
+	entry = generate()
+	c.misses++
+	c.entries[cacheKey] = entry
+	if entry != nil {
+		sz := entry.Size()
+		c.size += sz
+		totalCacheUsage += sz
+	}
+
+	// Check the capacity of all caches. If it greater than the maximum, we
+	// flush everything but still return this entry.
+	if totalCacheUsage > *maximumCacheUsage {
+		for entry, _ := range activeCaches {
+			// Drop all entries.
+			entry.size = 0
+			entry.entries = nil
+		}
+		totalCacheUsage = 0 // Reset.
+	}
+
+	return entry
+}
+
+// allCacheStats returns stats for all caches.
+func allCacheStats() string {
+	var sb strings.Builder
+	for entry, _ := range activeCaches {
+		ratio := float64(entry.hits) / float64(entry.hits+entry.misses)
+		fmt.Fprintf(&sb,
+			"% 10s: count: % 5d  size: % 10d  hits: % 7d  misses: % 7d  ratio: %2.2f\n",
+			entry.name, len(entry.entries), entry.size, entry.hits, entry.misses, ratio)
+	}
+	if len(activeCaches) > 0 {
+		fmt.Fprintf(&sb, "total: % 10d\n", totalCacheUsage)
+	}
+	return sb.String()
+}
+
+// LookupDigest returns a digest for the given file.
+func LookupDigest(filename string) (string, bool) {
+	digest, ok := inputFiles[filename]
+	return digest, ok
+}
+
+// Work invokes the main function.
+func Work(run func([]string) int) {
+	flag.CommandLine.Parse(os.Args[1:])
+	if !*persistentWorker {
+		// Handle the argument file.
+		args := flag.CommandLine.Args()
+		if len(args) == 1 && len(args[0]) > 1 && args[0][0] == '@' {
+			content, err := ioutil.ReadFile(args[0][1:])
+			if err != nil {
+				log.Fatalf("unable to parse args file: %v", err)
+			}
+			// Pull arguments from the file.
+			args = strings.Split(string(content), "\n")
+			flag.CommandLine.Parse(args)
+			args = flag.CommandLine.Args()
+		}
+		os.Exit(run(args))
+	}
+
+	var listenHeader string // Emitted always.
+	if *workerDebug {
+		// Bind a server for profiling.
+		listener, err := net.Listen("tcp", "localhost:0")
+		if err != nil {
+			log.Fatalf("unable to bind a server: %v", err)
+		}
+		// Construct the header for stats output, below.
+		listenHeader = fmt.Sprintf("Listening @ http://localhost:%d\n", listener.Addr().(*net.TCPAddr).Port)
+		go http.Serve(listener, nil)
+	}
+
+	// Move stdout. This is done to prevent anything else from accidentally
+	// printing to stdout, which must contain only the valid WorkerResponse
+	// serialized protos.
+	newOutput, err := unix.Dup(1)
+	if err != nil {
+		log.Fatalf("unable to move stdout: %v", err)
+	}
+	// Stderr may be closed or may be a copy of stdout. We make sure that
+	// we have an output that is in a completely separate range.
+	for newOutput <= 2 {
+		newOutput, err = unix.Dup(newOutput)
+		if err != nil {
+			log.Fatalf("unable to move stdout: %v", err)
+		}
+	}
+
+	// Best-effort: collect logs.
+	rPipe, wPipe, err := os.Pipe()
+	if err != nil {
+		log.Fatalf("unable to create pipe: %v", err)
+	}
+	if err := unix.Dup2(int(wPipe.Fd()), 1); err != nil {
+		log.Fatalf("error duping over stdout: %v", err)
+	}
+	if err := unix.Dup2(int(wPipe.Fd()), 2); err != nil {
+		log.Fatalf("error duping over stderr: %v", err)
+	}
+	wPipe.Close()
+	defer rPipe.Close()
+
+	// Read requests from stdin.
+	input := bufio.NewReader(os.NewFile(0, "input"))
+	output := bufio.NewWriter(os.NewFile(uintptr(newOutput), "output"))
+	for {
+		szBuf, err := input.Peek(4)
+		if err != nil {
+			log.Fatalf("unabel to read header: %v", err)
+		}
+
+		// Parse the size, and discard bits.
+		sz, szBytes := protowire.ConsumeVarint(szBuf)
+		if szBytes < 0 {
+			szBytes = 0
+		}
+		if _, err := input.Discard(szBytes); err != nil {
+			log.Fatalf("error discarding size: %v", err)
+		}
+
+		// Read a full message.
+		msg := make([]byte, int(sz))
+		if _, err := io.ReadFull(input, msg); err != nil {
+			log.Fatalf("error reading worker request: %v", err)
+		}
+		var wreq wpb.WorkRequest
+		if err := proto.Unmarshal(msg, &wreq); err != nil {
+			log.Fatalf("error unmarshaling worker request: %v", err)
+		}
+
+		// Flush relevant caches.
+		inputFiles = make(map[string]string)
+		for _, input := range wreq.GetInputs() {
+			updateInputFile(input.GetPath(), string(input.GetDigest()))
+		}
+
+		// Prepare logging.
+		outputBuffer := bytes.NewBuffer(nil)
+		outputBuffer.WriteString(listenHeader)
+		log.SetOutput(outputBuffer)
+
+		// Parse all arguments.
+		flag.CommandLine.Parse(wreq.GetArguments())
+		var exitCode int
+		exitChan := make(chan int)
+		go func() { exitChan <- run(flag.CommandLine.Args()) }()
+		for running := true; running; {
+			select {
+			case exitCode = <-exitChan:
+				running = false
+			default:
+			}
+			// N.B. rPipe is given a read deadline of 1ms. We expect
+			// this to turn a copy error after 1ms, and we just keep
+			// flushing this buffer while the task is running.
+			rPipe.SetReadDeadline(time.Now().Add(time.Millisecond))
+			outputBuffer.ReadFrom(rPipe)
+		}
+
+		if *workerDebug {
+			// Attach all cache stats.
+			outputBuffer.WriteString(allCacheStats())
+		}
+
+		// Send the response.
+		var wresp wpb.WorkResponse
+		wresp.ExitCode = int32(exitCode)
+		wresp.Output = string(outputBuffer.Bytes())
+		rmsg, err := proto.Marshal(&wresp)
+		if err != nil {
+			log.Fatalf("error marshaling response: %v", err)
+		}
+		if _, err := output.Write(append(protowire.AppendVarint(nil, uint64(len(rmsg))), rmsg...)); err != nil {
+			log.Fatalf("error sending worker response: %v", err)
+		}
+		if err := output.Flush(); err != nil {
+			log.Fatalf("error flushing output: %v", err)
+		}
+	}
+}