Skip to content

Commit 0bb331e

Browse files
tomhjptomhjp
authored
cachers,cmd/go-cacher: make HTTP PUTs async for BestEffortHTTP (#40)
Most of the time, PUTs don't significantly add to a build's latency, but occasionally we've seen an outsized amount of time spent on doing PUTs for a build. In addition to making BestEffortHTTP ignore errors, also make it run PUTs in a background goroutine with support for timeouts and concurrency control. This better represents best-effort PUTs as an optimisation, not a dependency for correctness. Disk writes still have to happen synchronously though, because cmd/go expects the file to exist and be readable/seekable as soon as the RPC returns. Updates tailscale/corp#45334 Signed-off-by: Tom Proctor <tomhjp@users.noreply.github.com>
1 parent a0c534b commit 0bb331e

3 files changed

Lines changed: 327 additions & 63 deletions

File tree

‎cachers/http.go‎

Lines changed: 153 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,15 @@ import (
44
"bytes"
55
"context"
66
"encoding/json"
7+
"errors"
78
"fmt"
89
"io"
910
"log"
1011
"net/http"
12+
"os"
1113
"strconv"
14+
"sync"
15+
"time"
1216

1317
"github.com/pierrec/lz4/v4"
1418
)
@@ -41,7 +45,75 @@ type HTTPClient struct {
4145
// BestEffortHTTP, when true, makes all HTTP errors non-fatal.
4246
// Get returns a cache miss and Put returns the local disk result,
4347
// silently ignoring any HTTP failures (connection errors, server errors, etc.).
48+
// It also makes the remote HTTP PUTs run asynchronously: Put returns as soon
49+
// as the blob is on local disk, and the upload happens in a background goroutine.
4450
BestEffortHTTP bool
51+
52+
// AsyncPutTimeout, if non-zero, bounds how long a background PUT may run
53+
// before its context is cancelled. It only applies when BestEffortHTTP is set.
54+
AsyncPutTimeout time.Duration
55+
56+
// AsyncPutMaxConcurrent, if positive, caps the number of background PUTs
57+
// running at once. New PUTs beyond the cap are queued until an existing PUT
58+
// finishes. If Close is called while some background PUTs are in-progress, it
59+
// will wait for up to AsyncPutTimeout before forcing shutdown. Disk writes
60+
// are not affected, and it only applies when BestEffortHTTP is set.
61+
AsyncPutMaxConcurrent int
62+
63+
asyncOnce sync.Once
64+
asyncPutSem chan struct{}
65+
inFlightWG sync.WaitGroup
66+
baseCtx context.Context
67+
baseCancel context.CancelFunc
68+
}
69+
70+
func (c *HTTPClient) ensureAsyncSetup() {
71+
c.asyncOnce.Do(func() {
72+
c.baseCtx, c.baseCancel = context.WithCancel(context.Background())
73+
if c.AsyncPutMaxConcurrent > 0 {
74+
c.asyncPutSem = make(chan struct{}, c.AsyncPutMaxConcurrent)
75+
}
76+
})
77+
}
78+
79+
// asyncPutSemaphore returns the buffered channel bounding concurrent background
80+
// PUTs, or nil if AsyncPutMaxConcurrent is unset (unbounded).
81+
func (c *HTTPClient) asyncPutSemaphore() chan struct{} {
82+
c.ensureAsyncSetup()
83+
return c.asyncPutSem
84+
}
85+
86+
// asyncPutContext returns the context shared by all background PUTs. Shutdown
87+
// cancels it when the drain deadline elapses.
88+
func (c *HTTPClient) asyncPutContext() context.Context {
89+
c.ensureAsyncSetup()
90+
return c.baseCtx
91+
}
92+
93+
// Shutdown blocks until all background PUTs have finished or AsyncPutTimeout
94+
// elapses, whichever comes first. If AsyncPutTimeout is zero (unbounded
95+
// PUTs), Shutdown waits indefinitely. It reports whether all PUTs drained in
96+
// time.
97+
func (c *HTTPClient) Shutdown() (drained bool) {
98+
done := make(chan struct{})
99+
go func() {
100+
c.inFlightWG.Wait()
101+
close(done)
102+
}()
103+
if c.AsyncPutTimeout == 0 {
104+
<-done
105+
return true
106+
}
107+
t := time.NewTimer(c.AsyncPutTimeout)
108+
defer t.Stop()
109+
select {
110+
case <-done:
111+
return true
112+
case <-t.C:
113+
c.ensureAsyncSetup()
114+
c.baseCancel()
115+
return false
116+
}
45117
}
46118

47119
func (c *HTTPClient) httpClient() *http.Client {
@@ -199,73 +271,98 @@ func (c *HTTPClient) Get(ctx context.Context, actionID string) (outputID, diskPa
199271
}
200272

201273
func (c *HTTPClient) Put(ctx context.Context, actionID, outputID string, size int64, body io.Reader) (diskPath string, _ error) {
202-
// Buffer the body so disk and HTTP can read from independent copies.
203-
// This avoids a race between our code and net/http's write loop when
204-
// the server responds (e.g. 403) before consuming the full request body.
205-
var buf []byte
206-
if size > 0 {
207-
var err error
208-
buf, err = io.ReadAll(body)
209-
if err != nil {
210-
return "", err
211-
}
274+
// Write to disk first.
275+
diskPath, err := c.Disk.Put(ctx, actionID, outputID, size, body)
276+
if err != nil {
277+
log.Printf("HTTPClient.Put local disk write error: %v", err)
278+
return "", err
212279
}
213280

214-
// Write to disk locally as we write it remotely, as we need to guarantee
215-
// it's on disk locally for the caller.
216-
diskPutCh := make(chan any, 1)
217-
go func() {
218-
diskPath, err := c.Disk.Put(ctx, actionID, outputID, size, bytes.NewReader(buf))
219-
if err != nil {
220-
diskPutCh <- err
221-
} else {
222-
diskPutCh <- diskPath
223-
}
224-
}()
281+
if c.BestEffortHTTP {
282+
sem := c.asyncPutSemaphore()
283+
// Background PUTs use a client-owned context, not the per-request ctx,
284+
// so they survive the RPC returning.
285+
putCtx := c.asyncPutContext()
286+
c.inFlightWG.Go(func() {
287+
// Acquire a concurrency slot before opening the file, so a queued
288+
// backlog doesn't hold an open file descriptor per pending upload.
289+
if sem != nil {
290+
select {
291+
case sem <- struct{}{}:
292+
case <-putCtx.Done():
293+
return
294+
}
295+
defer func() {
296+
<-sem
297+
}()
298+
}
299+
f, err := os.Open(diskPath)
300+
if err != nil {
301+
log.Printf("HTTPClient.Put local disk open after write error: %v", err)
302+
return
303+
}
304+
c.putRemote(putCtx, actionID, outputID, size, f)
305+
})
306+
return diskPath, nil
307+
}
308+
309+
f, err := os.Open(diskPath)
310+
if err != nil {
311+
log.Printf("HTTPClient.Put local disk open after write error: %v", err)
312+
return "", err
313+
}
314+
if err := c.putRemote(ctx, actionID, outputID, size, f); err != nil {
315+
return diskPath, err
316+
}
225317

226-
req, _ := http.NewRequestWithContext(ctx, "PUT", c.BaseURL+"/"+actionID+"/"+outputID, bytes.NewReader(buf))
318+
return diskPath, nil
319+
}
320+
321+
func (c *HTTPClient) putRemote(ctx context.Context, actionID, outputID string, size int64, body io.ReadCloser) error {
322+
defer body.Close()
323+
if c.BestEffortHTTP && c.AsyncPutTimeout != 0 {
324+
var cancel context.CancelFunc
325+
ctx, cancel = context.WithTimeout(ctx, c.AsyncPutTimeout)
326+
defer cancel()
327+
}
328+
// For a zero-length body, hand net/http NoBody so it sends an explicit
329+
// Content-Length: 0 rather than switching to chunked transfer encoding,
330+
// which the server rejects.
331+
var reqBody io.Reader = body
332+
if size == 0 {
333+
reqBody = http.NoBody
334+
}
335+
req, _ := http.NewRequestWithContext(ctx, "PUT", c.BaseURL+"/"+actionID+"/"+outputID, reqBody)
227336
req.ContentLength = size
228337
if c.AccessToken != "" {
229338
req.Header.Set("Authorization", "Bearer "+c.AccessToken)
230339
}
340+
231341
res, err := c.httpClient().Do(req)
232-
var httpErr error
233342
if err != nil {
234343
log.Printf("error PUT /%s/%s: %v", actionID, outputID, err)
235-
httpErr = err
236-
} else {
237-
defer res.Body.Close()
238-
if res.StatusCode != http.StatusNoContent {
239-
msg := tryReadErrorMessage(res)
240-
if c.BestEffortHTTP && (res.StatusCode == http.StatusUnauthorized || res.StatusCode == http.StatusForbidden) {
241-
// Known error codes that will repeatedly happen, so avoid
242-
// filling the logs with these errors.
243-
//
244-
// 401: can happen when gocached restarts during a session, as it
245-
// doesn't persist access tokens.
246-
// TODO(tomhjp): make the client retry auth in the background.
247-
//
248-
// 403: can happen when authed with a JWT that didn't grant global
249-
// write permissions.
250-
// TODO(tomhjp): support namespaces so all sessions can safely write.
251-
} else {
252-
log.Printf("error PUT /%s/%s: %v, %s", actionID, outputID, res.Status, msg)
253-
}
254-
httpErr = fmt.Errorf("unexpected PUT /%s/%s status %v", actionID, outputID, res.Status)
255-
}
344+
return err
256345
}
257-
// Wait for the disk write regardless of HTTP result.
258-
select {
259-
case v := <-diskPutCh:
260-
if diskErr, ok := v.(error); ok {
261-
log.Printf("HTTPClient.Put local disk error: %v", diskErr)
262-
return "", diskErr
263-
}
264-
if c.BestEffortHTTP {
265-
return v.(string), nil
346+
347+
defer res.Body.Close()
348+
if res.StatusCode != http.StatusNoContent {
349+
errMsg := fmt.Sprintf("error PUT /%s/%s: %s, %s", actionID, outputID, res.Status, tryReadErrorMessage(res))
350+
if c.BestEffortHTTP && (res.StatusCode == http.StatusUnauthorized || res.StatusCode == http.StatusForbidden) {
351+
// Known error codes that will repeatedly happen, so avoid
352+
// filling the logs with these errors.
353+
//
354+
// 401: can happen when gocached restarts during a session, as it
355+
// doesn't persist access tokens.
356+
// TODO(tomhjp): make the client retry auth in the background.
357+
//
358+
// 403: can happen when authed with a JWT that didn't grant global
359+
// write permissions.
360+
// TODO(tomhjp): support namespaces so all sessions can safely write.
361+
} else {
362+
log.Print(errMsg)
266363
}
267-
return v.(string), httpErr
268-
case <-ctx.Done():
269-
return "", ctx.Err()
364+
return errors.New(errMsg)
270365
}
366+
367+
return nil
271368
}

0 commit comments

Comments
 (0)