milo: Add a build info Swarming implementation.

milo/appengine/swarming/buildinfo.go

+// Copyright 2017 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 (
+        "strings"
+
+        swarming "github.com/luci/luci-go/common/api/swarming/swarming/v1"
+        "github.com/luci/luci-go/common/logging"
+        miloProto "github.com/luci/luci-go/common/proto/milo"
+        "github.com/luci/luci-go/grpc/grpcutil"
+        "github.com/luci/luci-go/logdog/client/coordinator"
+        "github.com/luci/luci-go/logdog/common/types"
+        "github.com/luci/luci-go/luci_config/common/cfgtypes"
+        milo "github.com/luci/luci-go/milo/api/proto"
+        "github.com/luci/luci-go/milo/appengine/logdog"
+
+        "golang.org/x/net/context"
+)
+
+// BuildInfoProvider is a configuration that provides build information.
+//
+// In a production system, this will be completely defaults. For testing, the
+// various services and data sources may be substituted for testing stubs.
+type BuildInfoProvider struct {
+        // LogdogClientFunc returns a coordinator Client instance for the supplied
+        // parameters.
+        //
+        // If nil, a production client will be generated.
+        LogdogClientFunc func(c context.Context) (*coordinator.Client, error)
+
+        // swarmingServiceFunc returns a swarmingService instance for the supplied
+        // parameters.
+        //
+        // If nil, a production fetcher will be generated.
+        swarmingServiceFunc func(c context.Context) (swarmingService, error)
+}
+
+func (p *BuildInfoProvider) newLogdogClient(c context.Context) (*coordinator.Client, error) {
+        if p.LogdogClientFunc != nil {
+                return p.LogdogClientFunc(c)
+        }
+        return logdog.NewClient(c, "")
+}
+
+func (p *BuildInfoProvider) newSwarmingService(c context.Context) (swarmingService, error) {
+        if p.swarmingServiceFunc != nil {
+                return p.swarmingServiceFunc(c)
+        }
+        return newProdService(c, defaultSwarmingServer)
+}
+
+// GetBuildInfo resolves a Milo protobuf Step for a given Swarming task.
+func (p *BuildInfoProvider) GetBuildInfo(c context.Context, req *milo.BuildInfoRequest_Swarming,
+        projectHint cfgtypes.ProjectName) (*milo.BuildInfoResponse, error) {
+
+        // Load the Swarming task (no log content).
+        sf, err := p.newSwarmingService(c)
+        if err != nil {
+                logging.WithError(err).Errorf(c, "Failed to create Swarming fetcher.")
+                return nil, grpcutil.Internal
+        }
+
+        // Use default Swarming host.
+        host := sf.getHost()
+        logging.Infof(c, "Loading build info for Swarming host %s, task %s.", host, req.Task)
+
+        fetch := swarmingFetch{
+                fetchReq: true,
+                fetchRes: true,
+        }
+        if err := fetch.get(c, sf, req.Task); err != nil {
+                if err == errNotMiloJob {
+                        logging.Warningf(c, "User requested non-Milo task.")
+                        return nil, grpcutil.NotFound
+                }
+
+                logging.WithError(err).Errorf(c, "Failed to load Swarming task.")
+                return nil, grpcutil.Internal
+        }
+
+        // Determine the LogDog annotation stream path for this Swarming task.
+        //
+        // On failure, will return a gRPC error.
+        as, err := getLogDogAnnotations(c, fetch.req, projectHint, host, req.Task)
+        if err != nil {
+                logging.WithError(err).Warningf(c, "Failed to get annotation stream parameters.")
+                return nil, err
+        }
+        if err := as.Normalize(); err != nil {
+                logging.WithError(err).Warningf(c, "Failed to normalize annotation stream parameters.")
+                return nil, grpcutil.Internal
+        }
+
+        client, err := p.newLogdogClient(c)
+        if err != nil {
+                logging.WithError(err).Errorf(c, "Failed to create LogDog client.")
+                return nil, grpcutil.Internal
+        }
+
+        as.Client = client
+        if err := as.Load(c); err != nil {
+                logging.WithError(err).Warningf(c, "Failed to load annotation stream.")
+                return nil, grpcutil.Internal
+        }
+
+        // Add Swarming task parameters to the Milo step.
+        step := as.Step()
+        if err := addTaskToMiloStep(c, sf.getHost(), fetch.res, step); err != nil {
+                return nil, err
+        }
+
+        prefix, name := as.Path.Split()
+        return &milo.BuildInfoResponse{
+                Project: string(as.Project),
+                Step:    step,
+                AnnotationStream: &miloProto.LogdogStream{
+                        Server: client.Host,
+                        Prefix: string(prefix),
+                        Name:   string(name),
+                },
+        }, nil
+}
+
+// getLogDogAnnotations returns a configured AnnotationStream given the input
+// parameters.
+//
+// This will return a gRPC error on failure.
+//
+// This function is messy and implementation-specific. That's the point of this
+// endpoint, though. All of the nastiness here should be replaced with something
+// more elegant once that becomes available. In the meantime...
+func getLogDogAnnotations(c context.Context, sr *swarming.SwarmingRpcsTaskRequest, projectHint cfgtypes.ProjectName,
+        host, taskID string) (*logdog.AnnotationStream, error) {
+
+        // If this is a Kitchen command, maybe we can infer our LogDog project from
+        // the command-line.
+        var as logdog.AnnotationStream
+        if sr.Properties == nil {
+                logging.Warningf(c, "No request properties, can't infer annotation stream path.")
+                return nil, grpcutil.NotFound
+        }
+
+        var isKitchen bool
+        if as.Project, isKitchen = getLogDogProjectFromKitchen(sr.Properties.Command); !isKitchen {
+                logging.Warningf(c, "Not a Kitchen CLI, can't infer annotation stream path.")
+                return nil, grpcutil.NotFound
+        }
+
+        if as.Project == "" {
+                as.Project = projectHint
+        }
+        if as.Project == "" {
+                logging.Warningf(c, "Don't know how to get annotation stream path.")
+                return nil, grpcutil.NotFound
+        }
+
+        // This is a Kitchen run, and it has a project! Construct the annotation
+        // stream path.
+        //
+        // This is an implementation of:
+        // https://chromium.googlesource.com/infra/infra/+/a7032e3e240d4b81a1912bfaf29a20d02f665cc1/go/src/infra/tools/kitchen/cook_logdog.go#129
+        prefix, err := types.MakeStreamName("", "swarm", host, taskID)
+        if err != nil {
+                logging.WithError(err).Errorf(c, "Failed to generate Swarming prefix.")
+                return nil, grpcutil.Internal
+        }
+        as.Path = prefix.Join("annotations")
+        return &as, nil
+}
+
+func getLogDogProjectFromKitchen(cmd []string) (proj cfgtypes.ProjectName, isKitchen bool) {
+        // Is this a Kitchen command?
+        switch {
+        case len(cmd) == 0:
+                return
+        case !strings.HasPrefix(cmd[0], "kitchen"):
+                return
+        }
+        isKitchen = true
+        cmd = cmd[1:]
+
+        // Scan through for the "-logdog-project" argument.
+        for i, arg := range cmd {
+                if arg == "-logdog-project" {
+                        if i < len(cmd)-2 {
+                                proj = cfgtypes.ProjectName(cmd[i+1])
+                                return
+                        }
+                        break
+                }
+        }
+        return
+}