prog.go

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package cache
     6  
     7  import (
     8  	"bufio"
     9  	"cmd/go/internal/base"
    10  	"cmd/go/internal/cacheprog"
    11  	"cmd/internal/quoted"
    12  	"context"
    13  	"crypto/sha256"
    14  	"encoding/base64"
    15  	"encoding/json"
    16  	"errors"
    17  	"fmt"
    18  	"io"
    19  	"log"
    20  	"os"
    21  	"os/exec"
    22  	"sync"
    23  	"sync/atomic"
    24  	"time"
    25  )
    26  
    27  // ProgCache implements Cache via JSON messages over stdin/stdout to a child
    28  // helper process which can then implement whatever caching policy/mechanism it
    29  // wants.
    30  //
    31  // See https://github.com/golang/go/issues/59719
    32  type ProgCache struct {
    33  	cmd    *exec.Cmd
    34  	stdout io.ReadCloser  // from the child process
    35  	stdin  io.WriteCloser // to the child process
    36  	bw     *bufio.Writer  // to stdin
    37  	jenc   *json.Encoder  // to bw
    38  
    39  	// can are the commands that the child process declared that it supports.
    40  	// This is effectively the versioning mechanism.
    41  	can map[cacheprog.Cmd]bool
    42  
    43  	// fuzzDirCache is another Cache implementation to use for the FuzzDir
    44  	// method. In practice this is the default GOCACHE disk-based
    45  	// implementation.
    46  	//
    47  	// TODO(bradfitz): maybe this isn't ideal. But we'd need to extend the Cache
    48  	// interface and the fuzzing callers to be less disk-y to do more here.
    49  	fuzzDirCache Cache
    50  
    51  	closing      atomic.Bool
    52  	ctx          context.Context    // valid until Close via ctxClose
    53  	ctxCancel    context.CancelFunc // called on Close
    54  	readLoopDone chan struct{}      // closed when readLoop returns
    55  
    56  	mu         sync.Mutex // guards following fields
    57  	nextID     int64
    58  	inFlight   map[int64]chan<- *cacheprog.Response
    59  	outputFile map[OutputID]string // object => abs path on disk
    60  
    61  	// writeMu serializes writing to the child process.
    62  	// It must never be held at the same time as mu.
    63  	writeMu sync.Mutex
    64  }
    65  
    66  // startCacheProg starts the prog binary (with optional space-separated flags)
    67  // and returns a Cache implementation that talks to it.
    68  //
    69  // It blocks a few seconds to wait for the child process to successfully start
    70  // and advertise its capabilities.
    71  func startCacheProg(progAndArgs string, fuzzDirCache Cache) Cache {
    72  	if fuzzDirCache == nil {
    73  		panic("missing fuzzDirCache")
    74  	}
    75  	args, err := quoted.Split(progAndArgs)
    76  	if err != nil {
    77  		base.Fatalf("GOCACHEPROG args: %v", err)
    78  	}
    79  	var prog string
    80  	if len(args) > 0 {
    81  		prog = args[0]
    82  		args = args[1:]
    83  	}
    84  
    85  	ctx, ctxCancel := context.WithCancel(context.Background())
    86  
    87  	cmd := exec.CommandContext(ctx, prog, args...)
    88  	out, err := cmd.StdoutPipe()
    89  	if err != nil {
    90  		base.Fatalf("StdoutPipe to GOCACHEPROG: %v", err)
    91  	}
    92  	in, err := cmd.StdinPipe()
    93  	if err != nil {
    94  		base.Fatalf("StdinPipe to GOCACHEPROG: %v", err)
    95  	}
    96  	cmd.Stderr = os.Stderr
    97  	// On close, we cancel the context. Rather than killing the helper,
    98  	// close its stdin.
    99  	cmd.Cancel = in.Close
   100  
   101  	if err := cmd.Start(); err != nil {
   102  		base.Fatalf("error starting GOCACHEPROG program %q: %v", prog, err)
   103  	}
   104  
   105  	pc := &ProgCache{
   106  		ctx:          ctx,
   107  		ctxCancel:    ctxCancel,
   108  		fuzzDirCache: fuzzDirCache,
   109  		cmd:          cmd,
   110  		stdout:       out,
   111  		stdin:        in,
   112  		bw:           bufio.NewWriter(in),
   113  		inFlight:     make(map[int64]chan<- *cacheprog.Response),
   114  		outputFile:   make(map[OutputID]string),
   115  		readLoopDone: make(chan struct{}),
   116  	}
   117  
   118  	// Register our interest in the initial protocol message from the child to
   119  	// us, saying what it can do.
   120  	capResc := make(chan *cacheprog.Response, 1)
   121  	pc.inFlight[0] = capResc
   122  
   123  	pc.jenc = json.NewEncoder(pc.bw)
   124  	go pc.readLoop(pc.readLoopDone)
   125  
   126  	// Give the child process a few seconds to report its capabilities. This
   127  	// should be instant and not require any slow work by the program.
   128  	timer := time.NewTicker(5 * time.Second)
   129  	defer timer.Stop()
   130  	for {
   131  		select {
   132  		case <-timer.C:
   133  			log.Printf("# still waiting for GOCACHEPROG %v ...", prog)
   134  		case capRes := <-capResc:
   135  			can := map[cacheprog.Cmd]bool{}
   136  			for _, cmd := range capRes.KnownCommands {
   137  				can[cmd] = true
   138  			}
   139  			if len(can) == 0 {
   140  				base.Fatalf("GOCACHEPROG %v declared no supported commands", prog)
   141  			}
   142  			pc.can = can
   143  			return pc
   144  		}
   145  	}
   146  }
   147  
   148  func (c *ProgCache) readLoop(readLoopDone chan<- struct{}) {
   149  	defer close(readLoopDone)
   150  	jd := json.NewDecoder(c.stdout)
   151  	for {
   152  		res := new(cacheprog.Response)
   153  		if err := jd.Decode(res); err != nil {
   154  			if c.closing.Load() {
   155  				c.mu.Lock()
   156  				for _, ch := range c.inFlight {
   157  					close(ch)
   158  				}
   159  				c.inFlight = nil
   160  				c.mu.Unlock()
   161  				return // quietly
   162  			}
   163  			if err == io.EOF {
   164  				c.mu.Lock()
   165  				inFlight := len(c.inFlight)
   166  				c.mu.Unlock()
   167  				base.Fatalf("GOCACHEPROG exited pre-Close with %v pending requests", inFlight)
   168  			}
   169  			base.Fatalf("error reading JSON from GOCACHEPROG: %v", err)
   170  		}
   171  		c.mu.Lock()
   172  		ch, ok := c.inFlight[res.ID]
   173  		delete(c.inFlight, res.ID)
   174  		c.mu.Unlock()
   175  		if ok {
   176  			ch <- res
   177  		} else {
   178  			base.Fatalf("GOCACHEPROG sent response for unknown request ID %v", res.ID)
   179  		}
   180  	}
   181  }
   182  
   183  var errCacheprogClosed = errors.New("GOCACHEPROG program closed unexpectedly")
   184  
   185  func (c *ProgCache) send(ctx context.Context, req *cacheprog.Request) (*cacheprog.Response, error) {
   186  	resc := make(chan *cacheprog.Response, 1)
   187  	if err := c.writeToChild(req, resc); err != nil {
   188  		return nil, err
   189  	}
   190  	select {
   191  	case res := <-resc:
   192  		if res == nil {
   193  			return nil, errCacheprogClosed
   194  		}
   195  		if res.Err != "" {
   196  			return nil, errors.New(res.Err)
   197  		}
   198  		return res, nil
   199  	case <-ctx.Done():
   200  		return nil, ctx.Err()
   201  	}
   202  }
   203  
   204  func (c *ProgCache) writeToChild(req *cacheprog.Request, resc chan<- *cacheprog.Response) (err error) {
   205  	c.mu.Lock()
   206  	if c.inFlight == nil {
   207  		c.mu.Unlock()
   208  		return errCacheprogClosed
   209  	}
   210  	c.nextID++
   211  	req.ID = c.nextID
   212  	c.inFlight[req.ID] = resc
   213  	c.mu.Unlock()
   214  
   215  	defer func() {
   216  		if err != nil {
   217  			c.mu.Lock()
   218  			if c.inFlight != nil {
   219  				delete(c.inFlight, req.ID)
   220  			}
   221  			c.mu.Unlock()
   222  		}
   223  	}()
   224  
   225  	c.writeMu.Lock()
   226  	defer c.writeMu.Unlock()
   227  
   228  	if err := c.jenc.Encode(req); err != nil {
   229  		return err
   230  	}
   231  	if err := c.bw.WriteByte('\n'); err != nil {
   232  		return err
   233  	}
   234  	if req.Body != nil && req.BodySize > 0 {
   235  		if err := c.bw.WriteByte('"'); err != nil {
   236  			return err
   237  		}
   238  		e := base64.NewEncoder(base64.StdEncoding, c.bw)
   239  		wrote, err := io.Copy(e, req.Body)
   240  		if err != nil {
   241  			return err
   242  		}
   243  		if err := e.Close(); err != nil {
   244  			return err
   245  		}
   246  		if wrote != req.BodySize {
   247  			return fmt.Errorf("short write writing body to GOCACHEPROG for action %x, output %x: wrote %v; expected %v",
   248  				req.ActionID, req.OutputID, wrote, req.BodySize)
   249  		}
   250  		if _, err := c.bw.WriteString("\"\n"); err != nil {
   251  			return err
   252  		}
   253  	}
   254  	if err := c.bw.Flush(); err != nil {
   255  		return err
   256  	}
   257  	return nil
   258  }
   259  
   260  func (c *ProgCache) Get(a ActionID) (Entry, error) {
   261  	if !c.can[cacheprog.CmdGet] {
   262  		// They can't do a "get". Maybe they're a write-only cache.
   263  		//
   264  		// TODO(bradfitz,bcmills): figure out the proper error type here. Maybe
   265  		// errors.ErrUnsupported? Is entryNotFoundError even appropriate? There
   266  		// might be places where we rely on the fact that a recent Put can be
   267  		// read through a corresponding Get. Audit callers and check, and document
   268  		// error types on the Cache interface.
   269  		return Entry{}, &entryNotFoundError{}
   270  	}
   271  	res, err := c.send(c.ctx, &cacheprog.Request{
   272  		Command:  cacheprog.CmdGet,
   273  		ActionID: a[:],
   274  	})
   275  	if err != nil {
   276  		return Entry{}, err // TODO(bradfitz): or entryNotFoundError? Audit callers.
   277  	}
   278  	if res.Miss {
   279  		return Entry{}, &entryNotFoundError{}
   280  	}
   281  	e := Entry{
   282  		Size: res.Size,
   283  	}
   284  	if res.Time != nil {
   285  		e.Time = *res.Time
   286  	} else {
   287  		e.Time = time.Now()
   288  	}
   289  	if res.DiskPath == "" {
   290  		return Entry{}, &entryNotFoundError{errors.New("GOCACHEPROG didn't populate DiskPath on get hit")}
   291  	}
   292  	if copy(e.OutputID[:], res.OutputID) != len(res.OutputID) {
   293  		return Entry{}, &entryNotFoundError{errors.New("incomplete ProgResponse OutputID")}
   294  	}
   295  	c.noteOutputFile(e.OutputID, res.DiskPath)
   296  	return e, nil
   297  }
   298  
   299  func (c *ProgCache) noteOutputFile(o OutputID, diskPath string) {
   300  	c.mu.Lock()
   301  	defer c.mu.Unlock()
   302  	c.outputFile[o] = diskPath
   303  }
   304  
   305  func (c *ProgCache) OutputFile(o OutputID) string {
   306  	c.mu.Lock()
   307  	defer c.mu.Unlock()
   308  	return c.outputFile[o]
   309  }
   310  
   311  func (c *ProgCache) Put(a ActionID, file io.ReadSeeker) (_ OutputID, size int64, _ error) {
   312  	// Compute output ID.
   313  	h := sha256.New()
   314  	if _, err := file.Seek(0, 0); err != nil {
   315  		return OutputID{}, 0, err
   316  	}
   317  	size, err := io.Copy(h, file)
   318  	if err != nil {
   319  		return OutputID{}, 0, err
   320  	}
   321  	var out OutputID
   322  	h.Sum(out[:0])
   323  
   324  	if _, err := file.Seek(0, 0); err != nil {
   325  		return OutputID{}, 0, err
   326  	}
   327  
   328  	if !c.can[cacheprog.CmdPut] {
   329  		// Child is a read-only cache. Do nothing.
   330  		return out, size, nil
   331  	}
   332  
   333  	res, err := c.send(c.ctx, &cacheprog.Request{
   334  		Command:  cacheprog.CmdPut,
   335  		ActionID: a[:],
   336  		OutputID: out[:],
   337  		Body:     file,
   338  		BodySize: size,
   339  	})
   340  	if err != nil {
   341  		return OutputID{}, 0, err
   342  	}
   343  	if res.DiskPath == "" {
   344  		return OutputID{}, 0, errors.New("GOCACHEPROG didn't return DiskPath in put response")
   345  	}
   346  	c.noteOutputFile(out, res.DiskPath)
   347  	return out, size, err
   348  }
   349  
   350  func (c *ProgCache) Close() error {
   351  	c.closing.Store(true)
   352  	var err error
   353  
   354  	// First write a "close" message to the child so it can exit nicely
   355  	// and clean up if it wants. Only after that exchange do we cancel
   356  	// the context that kills the process.
   357  	if c.can[cacheprog.CmdClose] {
   358  		_, err = c.send(c.ctx, &cacheprog.Request{Command: cacheprog.CmdClose})
   359  		if errors.Is(err, errCacheprogClosed) {
   360  			// Allow the child to quit without responding to close.
   361  			err = nil
   362  		}
   363  	}
   364  	// Cancel the context, which will close the helper's stdin.
   365  	c.ctxCancel()
   366  	// Wait until the helper closes its stdout.
   367  	<-c.readLoopDone
   368  	return err
   369  }
   370  
   371  func (c *ProgCache) FuzzDir() string {
   372  	// TODO(bradfitz): figure out what to do here. For now just use the
   373  	// disk-based default.
   374  	return c.fuzzDirCache.FuzzDir()
   375  }
   376
View as plain text