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 during
+        // the swarmingFetch of the fetch result (if enabled). It will contain the

[hinoka, 2017/02/16 02:08:21]

invoked after fetching the swarm fetchRes call?

[dnj, 2017/02/16 02:20:51]

Done.

+        // mapped result tags.
+        //
+        // If taskTagCallback returns true, any pending log fetch will be cancelled
+        // without error.
+        taskTagCallback func(map[string]string) bool

[hinoka, 2017/02/16 02:08:21]

This is funky, but I don't have any better ideas :)

[dnj, 2017/02/16 02:20:51]

Acknowledged.

 }
 
 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) {
+type buildLoader struct {

[hinoka, 2017/02/16 02:08:21]

Comment on why this exists (testing)

[dnj, 2017/02/16 02:20:51]

Done.

+        // 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) newAnnotationStream(c context.Context, stream *types.Stream) (*logdog.AnnotationStream, error) {

[hinoka, 2017/02/16 02:08:21]

newEmptyAnnotationStream might be more appropriate?  To make it clear that it's
not a populated one.

At some point we probably want to go through and do a refactor and rename things
to make things easier to read.

[dnj, 2017/02/16 02:20:51]

Done.

+        fn := bl.logDogClientFunc
+        if fn == nil {
+                fn = logdog.NewClient
+        }
+        client, err := fn(c, stream.Host)
+        if err != nil {
+                return nil, errors.Annotate(err).Reason("failed to create LogDog client").Err()
+        }
+
+        as := logdog.AnnotationStream{
+                Client:  client,
+                Project: stream.Project,
+                Path:    stream.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 logDogStream *types.Stream
+
         fetchParams := swarmingFetchParams{
                 fetchRes: true,
                 fetchLog: true,
+
+                // Cancel if LogDog annotation strem parameters are present in the tag set.

[hinoka, 2017/02/16 02:08:21]

stream

[dnj, 2017/02/16 02:20:51]

Done.

+                taskTagCallback: func(tags map[string]string) (cancelLogs bool) {
+                        var err error
+                        if logDogStream, err = resolveLogDogStreamFromTags(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
 
-        // 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 != "" {
+        // If the LogDog stream is available, load the step from that.

[hinoka, 2017/02/16 02:08:21]

Add to comment: Explanation on why we're preferring one over the other.  ie
tasks with log_location will probably not contain annotated swarming logs.

[dnj, 2017/02/16 02:20:51]

Done.

+        switch {
+        case logDogStream != nil:
+                as, err := bl.newAnnotationStream(c, logDogStream)
+                if err != nil {
+                        return nil, errors.Annotate(err).Reason("failed to create LogDog annotation stream").Err()
+                }
+
+                if s, err = as.Load(c); err != nil {

[hinoka, 2017/02/16 02:08:21]

Fetch seems like a more appropriate name

[dnj, 2017/02/16 02:20:51]

Done.

+                        return nil, errors.Annotate(err).Reason("failed to load LogDog annotation stream").Err()
+                }
+
+        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
+                }
+
+        default:
                 s = &miloProto.Step{}
         }
 
         if err := addTaskToMiloStep(c, svc.getHost(), fr.res, s); err != nil {
                 return nil, err
         }
         logdog.AddLogDogToBuild(c, swarmingURLBuilder(linkBase), s, &build)
 
         if err := addTaskToBuild(c, svc.getHost(), fr.res, &build); err != nil {
                 return nil, err
@@ skipping 57 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
+}