milo: Enable Swarming LogDog log loading.

milo/appengine/swarming/build.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package swarming
 
 import (
         "bytes"
         "fmt"
         "net/http"
         "net/url"
         "strings"
         "time"
 
         "golang.org/x/net/context"
 
         swarming "github.com/luci/luci-go/common/api/swarming/swarming/v1"
         "github.com/luci/luci-go/common/errors"
         "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/common/proto/google"
         miloProto "github.com/luci/luci-go/common/proto/milo"
         "github.com/luci/luci-go/common/sync/parallel"
         "github.com/luci/luci-go/logdog/client/annotee"
+        "github.com/luci/luci-go/logdog/client/coordinator"
         "github.com/luci/luci-go/logdog/common/types"
         "github.com/luci/luci-go/milo/api/resp"
         "github.com/luci/luci-go/milo/appengine/logdog"
         "github.com/luci/luci-go/server/auth"
 )
 
 // errNotMiloJob is returned if a Swarming task is fetched that does not self-
 // identify as a Milo job.
 var errNotMiloJob = errors.New("Not a Milo Job")
 
@@ skipping 74 matching lines @@
 }
 
 func (svc *prodSwarmingService) getSwarmingRequest(c context.Context, taskID string) (*swarming.SwarmingRpcsTaskRequest, error) {
         return svc.client.Task.Request(taskID).Context(c).Do()
 }
 
 type swarmingFetchParams struct {
         fetchReq bool
         fetchRes bool
         fetchLog bool
+
+        // taskTagCallback, if not nil, is a callback that will be invoked after
+        // fetching the result, if fetchRes is true. It will be passed a key/value map
+        // of the Swarming result's tags.
+        //
+        // If taskTagCallback returns true, any pending log fetch will be cancelled
+        // without error.
+        taskTagCallback func(map[string]string) bool
 }
 
 type swarmingFetchResult struct {
         req *swarming.SwarmingRpcsTaskRequest
         res *swarming.SwarmingRpcsTaskResult
+
+        // log is the log data content. If no log data was fetched, this will empty.
+        // If the log fetch was cancelled, this is undefined.
         log string
 }
 
 // swarmingFetch fetches (in parallel) the components that it is configured to
 // fetch.
 //
 // After fetching, an ACL check is performed to confirm that the user is
 // permitted to view the resulting data. If this check fails, get returns
 // errNotMiloJob.
 //
 // TODO(hinoka): Make this ACL check something more specific than the
 // resence of the "allow_milo" dimension.
 func swarmingFetch(c context.Context, svc swarmingService, taskID string, req swarmingFetchParams) (
         *swarmingFetchResult, error) {
 
         // logErr is managed separately from other fetch errors, since in some
         // situations it's acceptable to not have a log stream.
         var logErr error
         var fr swarmingFetchResult
+        var resTags map[string]string
+
+        // Special Context to enable the cancellation of log fetching.
+        logsCancelled := false
+        logCtx, cancelLogs := context.WithCancel(c)
+        defer cancelLogs()
 
         err := parallel.FanOutIn(func(workC chan<- func() error) {
                 if req.fetchReq {
                         workC <- func() (err error) {
                                 fr.req, err = svc.getSwarmingRequest(c, taskID)
                                 return
                         }
                 }
 
                 if req.fetchRes {
                         workC <- func() (err error) {
-                                fr.res, err = svc.getSwarmingResult(c, taskID)
+                                if fr.res, err = svc.getSwarmingResult(c, taskID); err == nil {
+                                        resTags = swarmingTags(fr.res.Tags)
+                                        if req.taskTagCallback != nil && req.taskTagCallback(resTags) {
+                                                logsCancelled = true
+                                                cancelLogs()
+                                        }
+                                }
                                 return
                         }
                 }
 
                 if req.fetchLog {
                         workC <- func() error {
-                                fr.log, logErr = svc.getTaskOutput(c, taskID)
+                                // Note: we're using the log Context here so we can cancel log fetch
+                                // explicitly.
+                                fr.log, logErr = svc.getTaskOutput(logCtx, taskID)
                                 return nil
                         }
                 }
         })
         if err != nil {
                 return nil, err
         }
 
         // Current ACL implementation: error if this is not a Milo job.
         switch {
         case req.fetchReq:
                 if !isMiloJob(fr.req.Tags) {
                         return nil, errNotMiloJob
                 }
+
         case req.fetchRes:
                 if !isMiloJob(fr.res.Tags) {
                         return nil, errNotMiloJob
                 }
+
         default:
                 // No metadata to decide if this is a Milo job, so assume that it is not.
                 return nil, errNotMiloJob
         }
 
         if req.fetchRes && logErr != nil {
                 switch fr.res.State {
                 case TaskCompleted, TaskRunning, TaskCanceled:
                 default:
                         //  Ignore log errors if the task might be pending, timed out, expired, etc.
                         if err != nil {
                                 fr.log = ""
                                 logErr = nil
                         }
                 }
         }
+
+        // If we explicitly cancelled logs, everything is OK.
+        if logErr == context.Canceled && logsCancelled {
+                logErr = nil
+        }
         return &fr, logErr
 }
 
 func taskProperties(sr *swarming.SwarmingRpcsTaskResult) *resp.PropertyGroup {
         props := &resp.PropertyGroup{GroupName: "Swarming"}
         if len(sr.CostsUsd) == 1 {
                 props.Property = append(props.Property, &resp.Property{
                         Key:   "Cost of job (USD)",
                         Value: fmt.Sprintf("$%.2f", sr.CostsUsd[0]),
                 })
@@ skipping 18 matching lines @@
                         Key: parts[0],
                 }
                 if len(parts) == 2 {
                         p.Value = parts[1]
                 }
                 props.Property = append(props.Property, p)
         }
         return props
 }
 
-// tagValue returns a value of the first tag matching the tag key. If not found
-// returns "".
-func tagValue(tags []string, key string) string {
-        prefix := key + ":"
-        for _, t := range tags {
-                if strings.HasPrefix(t, prefix) {
-                        return strings.TrimPrefix(t, prefix)
-                }
-        }
-        return ""
-}
-
 // addBuilderLink adds a link to the buildbucket builder view.
 func addBuilderLink(c context.Context, build *resp.MiloBuild, swarmingHostname string, sr *swarming.SwarmingRpcsTaskResult) {
-        bbHost := tagValue(sr.Tags, "buildbucket_hostname")
-        bucket := tagValue(sr.Tags, "buildbucket_bucket")
-        builder := tagValue(sr.Tags, "builder")
+        tags := swarmingTags(sr.Tags)
+
+        bbHost := tags["buildbucket_hostname"]
+        bucket := tags["buildbucket_bucket"]
+        builder := tags["builder"]
         if bucket == "" {
                 logging.Errorf(
                         c, "Could not extract buidlbucket bucket from task %s",
                         taskPageURL(swarmingHostname, sr.TaskId))
         }
         if builder == "" {
                 logging.Errorf(
                         c, "Could not extract builder name from task %s",
                         taskPageURL(swarmingHostname, sr.TaskId))
         }
@@ skipping 178 matching lines @@
         }
         // If this ever has more than one stream then memoryClient needs to become
         // goroutine safe
         if err := p.RunStreams([]*annotee.Stream{&is}); err != nil {
                 return nil, err
         }
         p.Finish()
         return c.ToLogDogStreams()
 }
 
-func swarmingBuildImpl(c context.Context, svc swarmingService, linkBase, taskID string) (*resp.MiloBuild, error) {
+// buildLoader represents the ability to load a Milo build from a Swarming task.
+//
+// It exists so that the internal build loading functionality can be stubbed out
+// for testing.
+type buildLoader struct {
+        // logdogClientFunc returns a coordinator Client instance for the supplied
+        // parameters.
+        //
+        // If nil, a production client will be generated.
+        logDogClientFunc func(c context.Context, host string) (*coordinator.Client, error)
+}
+
+func (bl *buildLoader) newEmptyAnnotationStream(c context.Context, addr *types.StreamAddr) (
+        *logdog.AnnotationStream, error) {
+
+        fn := bl.logDogClientFunc
+        if fn == nil {
+                fn = logdog.NewClient
+        }
+        client, err := fn(c, addr.Host)
+        if err != nil {
+                return nil, errors.Annotate(err).Reason("failed to create LogDog client").Err()
+        }
+
+        as := logdog.AnnotationStream{
+                Client:  client,
+                Project: addr.Project,
+                Path:    addr.Path,
+        }
+        if err := as.Normalize(); err != nil {
+                return nil, errors.Annotate(err).Reason("failed to normalize annotation stream parameters").Err()
+        }
+
+        return &as, nil
+}
+
+func (bl *buildLoader) swarmingBuildImpl(c context.Context, svc swarmingService, linkBase, taskID string) (*resp.MiloBuild, error) {
         // Fetch the data from Swarming
+        var logDogStreamAddr *types.StreamAddr
+
         fetchParams := swarmingFetchParams{
                 fetchRes: true,
                 fetchLog: true,
+
+                // Cancel if LogDog annotation stream parameters are present in the tag set.
+                taskTagCallback: func(tags map[string]string) (cancelLogs bool) {
+                        var err error
+                        if logDogStreamAddr, err = resolveLogDogStreamAddrFromTags(tags); err != nil {
+                                logging.WithError(err).Debugf(c, "Not using LogDog annotation stream.")
+                                return false
+                        }
+                        return true
+                },
         }
         fr, err := swarmingFetch(c, svc, taskID, fetchParams)
         if err != nil {
                 return nil, err
         }
 
         var build resp.MiloBuild
         var s *miloProto.Step
         var lds *logdog.Streams
+        var ub logdog.URLBuilder
 
-        // Decode the data using annotee. The logdog stream returned here is assumed
-        // to be consistent, which is why the following block of code are not
-        // expected to ever err out.
-        if fr.log != "" {
+        // Load the build from the available data.
+        //
+        // If the Swarming task explicitly specifies its log location, we prefer that.
+        // As a fallback, we will try and parse the Swarming task's output for
+        // annotations.
+        switch {
+        case logDogStreamAddr != nil:
+                // If the LogDog stream is available, load the step from that.
+                as, err := bl.newEmptyAnnotationStream(c, logDogStreamAddr)
+                if err != nil {
+                        return nil, errors.Annotate(err).Reason("failed to create LogDog annotation stream").Err()
+                }
+
+                if s, err = as.Fetch(c); err != nil {
+                        return nil, errors.Annotate(err).Reason("failed to load LogDog annotation stream").Err()
+                }
+
+                prefix, _ := logDogStreamAddr.Path.Split()
+                ub = &logdog.ViewerURLBuilder{
+                        Host:    logDogStreamAddr.Host,
+                        Prefix:  prefix,
+                        Project: logDogStreamAddr.Project,
+                }
+
+        case fr.log != "":
+                // Decode the data using annotee. The logdog stream returned here is assumed
+                // to be consistent, which is why the following block of code are not
+                // expected to ever err out.
                 var err error
                 lds, err = streamsFromAnnotatedLog(c, fr.log)
                 if err != nil {
                         build.Components = []*resp.BuildComponent{{
                                 Type:   resp.Summary,
                                 Label:  "Milo annotation parser",
                                 Text:   []string{err.Error()},
                                 Status: resp.InfraFailure,
                                 SubLink: []*resp.Link{{
                                         Label: "swarming task",
                                         URL:   taskPageURL(svc.getHost(), taskID),
                                 }},
                         }}
                 }
-        }
 
-        if lds != nil && lds.MainStream != nil && lds.MainStream.Data != nil {
-                s = lds.MainStream.Data
-        } else {
+                if lds != nil && lds.MainStream != nil && lds.MainStream.Data != nil {
+                        s = lds.MainStream.Data
+                }
+                ub = swarmingURLBuilder(linkBase)
+
+        default:
                 s = &miloProto.Step{}
+                ub = swarmingURLBuilder(linkBase)
         }
 
         if err := addTaskToMiloStep(c, svc.getHost(), fr.res, s); err != nil {
                 return nil, err
         }
-        logdog.AddLogDogToBuild(c, swarmingURLBuilder(linkBase), s, &build)
+        logdog.AddLogDogToBuild(c, ub, s, &build)
 
         if err := addTaskToBuild(c, svc.getHost(), fr.res, &build); err != nil {
                 return nil, err
         }
 
         return &build, nil
 }
 
 func isMiloJob(tags []string) bool {
         for _, t := range tags {
@@ skipping 50 matching lines @@
         case *miloProto.Link_Url:
                 return &resp.Link{
                         Label: l.Label,
                         URL:   t.Url,
                 }
 
         default:
                 return nil
         }
 }
+
+func swarmingTags(v []string) map[string]string {
+        res := make(map[string]string, len(v))
+        for _, tag := range v {
+                var value string
+                parts := strings.SplitN(tag, ":", 2)
+                if len(parts) == 2 {
+                        value = parts[1]
+                }
+                res[parts[0]] = value
+        }
+        return res
+}