storage: add Context arguments for queries

Change-Id: I0ba5b534e0d5262e2626351d10564219d18f93a8
Reviewed-on: https://go-review.googlesource.com/41372Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

analysis/app/compare.go

@@ -17,6 +17,7 @@ import (
 	"strings"
 	"unicode"
 
+	"golang.org/x/net/context"
 	"golang.org/x/perf/benchstat"
 	"golang.org/x/perf/storage/benchfmt"
 	"golang.org/x/perf/storage/query"
@@ -134,6 +135,8 @@ func addToQuery(query, add string) string {
 
 // compare handles queries that require comparison of the groups in the query.
 func (a *App) compare(w http.ResponseWriter, r *http.Request) {
+	ctx := requestContext(r)
+
 	if err := r.ParseForm(); err != nil {
 		http.Error(w, err.Error(), 500)
 		return
@@ -155,7 +158,7 @@ func (a *App) compare(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	data := a.compareQuery(q)
+	data := a.compareQuery(ctx, q)
 
 	w.Header().Set("Content-Type", "text/html; charset=utf-8")
 	if err := t.Execute(w, data); err != nil {
@@ -225,7 +228,7 @@ func elideKeyValues(content string, keys map[string]bool) string {
 
 // fetchCompareResults fetches the matching results for a given query string.
 // The results will be grouped into one or more groups based on either the query string or heuristics.
-func (a *App) fetchCompareResults(q string) ([]*resultGroup, error) {
+func (a *App) fetchCompareResults(ctx context.Context, q string) ([]*resultGroup, error) {
 	// Parse query
 	prefix, queries := parseQueryString(q)
 
@@ -239,7 +242,7 @@ func (a *App) fetchCompareResults(q string) ([]*resultGroup, error) {
 		if prefix != "" {
 			qPart = prefix + " " + qPart
 		}
-		res := a.StorageClient.Query(qPart)
+		res := a.StorageClient.Query(ctx, qPart)
 		for res.Next() {
 			result := res.Result()
 			result.Content = elideKeyValues(result.Content, keys)
@@ -275,12 +278,12 @@ func (a *App) fetchCompareResults(q string) ([]*resultGroup, error) {
 	return groups, nil
 }
 
-func (a *App) compareQuery(q string) *compareData {
+func (a *App) compareQuery(ctx context.Context, q string) *compareData {
 	if len(q) == 0 {
 		return &compareData{}
 	}
 
-	groups, err := a.fetchCompareResults(q)
+	groups, err := a.fetchCompareResults(ctx, q)
 	if err != nil {
 		return &compareData{
 			Q:     q,
@@ -349,6 +352,8 @@ func (a *App) compareQuery(q string) *compareData {
 
 // textCompare is called if benchsave is requesting a text-only analysis.
 func (a *App) textCompare(w http.ResponseWriter, r *http.Request) {
+	ctx := requestContext(r)
+
 	if err := r.ParseForm(); err != nil {
 		http.Error(w, err.Error(), 500)
 		return
@@ -358,7 +363,7 @@ func (a *App) textCompare(w http.ResponseWriter, r *http.Request) {
 
 	q := r.Form.Get("q")
 
-	groups, err := a.fetchCompareResults(q)
+	groups, err := a.fetchCompareResults(ctx, q)
 	if err != nil {
 		// TODO(quentin): Should we serve this with a 500 or 404? This means the query was invalid or had no results.
 		fmt.Fprintf(w, "unable to analyze results: %v", err)

analysis/app/compare_test.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"testing"
 
+	"golang.org/x/net/context"
 	"golang.org/x/perf/storage"
 	"golang.org/x/perf/storage/benchfmt"
 )
@@ -89,7 +90,7 @@ func TestCompareQuery(t *testing.T) {
 
 	for _, q := range []string{"one vs two", "onetwo"} {
 		t.Run(q, func(t *testing.T) {
-			data := a.compareQuery(q)
+			data := a.compareQuery(context.Background(), q)
 			if data.Error != "" {
 				t.Fatalf("compareQuery failed: %s", data.Error)
 			}

analysis/app/index.go

@@ -30,7 +30,7 @@ func (a *App) index(w http.ResponseWriter, r *http.Request) {
 	}
 
 	var uploads []storage.UploadInfo
-	ul := a.StorageClient.ListUploads("", []string{"by", "upload-time"}, 16)
+	ul := a.StorageClient.ListUploads(ctx, "", []string{"by", "upload-time"}, 16)
 	defer ul.Close()
 	for ul.Next() {
 		uploads = append(uploads, ul.Info())

analysis/app/trend.go

@@ -78,7 +78,7 @@ type trendData struct {
 func (a *App) trendQuery(ctx context.Context, q string, opt plotOptions) *trendData {
 	d := &trendData{Q: q}
 	if q == "" {
-		ul := a.StorageClient.ListUploads(`trend>`, []string{"by", "upload-time", "trend"}, 16)
+		ul := a.StorageClient.ListUploads(ctx, `trend>`, []string{"by", "upload-time", "trend"}, 16)
 		defer ul.Close()
 		for ul.Next() {
 			d.TrendUploads = append(d.TrendUploads, ul.Info())
@@ -90,7 +90,7 @@ func (a *App) trendQuery(ctx context.Context, q string, opt plotOptions) *trendD
 	}
 
 	// TODO(quentin): Chunk query based on matching upload IDs.
-	res := a.StorageClient.Query(q)
+	res := a.StorageClient.Query(ctx, q)
 	defer res.Close()
 	t, resultCols := queryToTable(res)
 	if err := res.Err(); err != nil {

cmd/benchsave/benchsave.go

@@ -102,7 +102,7 @@ func main() {
 
 	start := time.Now()
 
-	u := client.NewUpload()
+	u := client.NewUpload(context.Background())
 
 	for _, name := range files {
 		if err := writeOneFile(u, name, headerData); err != nil {

storage/client.go

@@ -14,6 +14,8 @@ import (
 	"net/http"
 	"net/url"
 
+	"golang.org/x/net/context"
+	"golang.org/x/net/context/ctxhttp"
 	"golang.org/x/perf/storage/benchfmt"
 )
 
@@ -41,10 +43,10 @@ func (c *Client) httpClient() *http.Client {
 // key:value - exact match on label "key" = "value"
 // key>value - value greater than (useful for dates)
 // key<value - value less than (also useful for dates)
-func (c *Client) Query(q string) *Query {
+func (c *Client) Query(ctx context.Context, q string) *Query {
 	hc := c.httpClient()
 
-	resp, err := hc.Get(c.BaseURL + "/search?" + url.Values{"q": []string{q}}.Encode())
+	resp, err := ctxhttp.Get(ctx, hc, c.BaseURL+"/search?"+url.Values{"q": []string{q}}.Encode())
 	if err != nil {
 		return &Query{err: err}
 	}
@@ -124,7 +126,7 @@ type UploadInfo struct {
 // extraLabels specifies other labels to be retrieved.
 // If limit is 0, no limit will be provided to the server.
 // The uploads are returned starting with the most recent upload.
-func (c *Client) ListUploads(q string, extraLabels []string, limit int) *UploadList {
+func (c *Client) ListUploads(ctx context.Context, q string, extraLabels []string, limit int) *UploadList {
 	hc := c.httpClient()
 
 	v := url.Values{"extra_label": extraLabels}
@@ -139,7 +141,7 @@ func (c *Client) ListUploads(q string, extraLabels []string, limit int) *UploadL
 	if len(v) > 0 {
 		u += "?" + v.Encode()
 	}
-	resp, err := hc.Get(u)
+	resp, err := ctxhttp.Get(ctx, hc, u)
 	if err != nil {
 		return &UploadList{err: err}
 	}
@@ -213,7 +215,7 @@ func (ul *UploadList) Close() error {
 // NewUpload starts a new upload to the storage server.
 // The upload must have Abort or Commit called on it.
 // If the server requires authentication for uploads, c.HTTPClient should be set to the result of oauth2.NewClient.
-func (c *Client) NewUpload() *Upload {
+func (c *Client) NewUpload(ctx context.Context) *Upload {
 	hc := c.httpClient()
 
 	pr, pw := io.Pipe()
@@ -228,7 +230,7 @@ func (c *Client) NewUpload() *Upload {
 	errCh := make(chan error)
 	u := &Upload{pw: pw, mpw: mpw, errCh: errCh}
 	go func() {
-		resp, err := hc.Do(req)
+		resp, err := ctxhttp.Do(ctx, hc, req)
 		if err != nil {
 			errCh <- err
 			return

storage/client_test.go

@@ -14,6 +14,7 @@ import (
 	"reflect"
 	"testing"
 
+	"golang.org/x/net/context"
 	"golang.org/x/perf/internal/diff"
 	"golang.org/x/perf/storage/benchfmt"
 )
@@ -26,7 +27,7 @@ func TestQueryError(t *testing.T) {
 
 	c := &Client{BaseURL: ts.URL}
 
-	q := c.Query("invalid query")
+	q := c.Query(context.Background(), "invalid query")
 	defer q.Close()
 
 	if q.Next() {
@@ -48,7 +49,7 @@ func TestQuery(t *testing.T) {
 
 	c := &Client{BaseURL: ts.URL}
 
-	q := c.Query("key1:value key2:value")
+	q := c.Query(context.Background(), "key1:value key2:value")
 	defer q.Close()
 
 	var buf bytes.Buffer
@@ -79,7 +80,7 @@ func TestListUploads(t *testing.T) {
 
 	c := &Client{BaseURL: ts.URL}
 
-	r := c.ListUploads("key1:value key2:value", []string{"key1", "key2"}, 10)
+	r := c.ListUploads(context.Background(), "key1:value key2:value", []string{"key1", "key2"}, 10)
 	defer r.Close()
 
 	if !r.Next() {
@@ -132,7 +133,7 @@ func TestNewUpload(t *testing.T) {
 
 	c := &Client{BaseURL: ts.URL}
 
-	u := c.NewUpload()
+	u := c.NewUpload(context.Background())
 	for i := 0; i < 2; i++ {
 		w, err := u.CreateFile(fmt.Sprintf("want%d.txt", i))
 		if err != nil {
@@ -190,7 +191,7 @@ func TestNewUploadAbort(t *testing.T) {
 
 	c := &Client{BaseURL: ts.URL}
 
-	u := c.NewUpload()
+	u := c.NewUpload(context.Background())
 	for i := 0; i < 2; i++ {
 		w, err := u.CreateFile(fmt.Sprintf("want%d.txt", i))
 		if err != nil {