Kitchen: Generate LogDog prefix from Swarming.

go/src/infra/tools/kitchen/cook_logdog.go

 // Copyright 2016 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 main
 
 import (
+        "flag"
+        "fmt"
         "io"
         "os/exec"
         "sync"
         "time"
 
         "github.com/luci/luci-go/client/logdog/annotee"
         "github.com/luci/luci-go/client/logdog/annotee/annotation"
         "github.com/luci/luci-go/client/logdog/butler"
         "github.com/luci/luci-go/client/logdog/butler/bootstrap"
         "github.com/luci/luci-go/client/logdog/butler/output"
         fileOut "github.com/luci/luci-go/client/logdog/butler/output/file"
         out "github.com/luci/luci-go/client/logdog/butler/output/logdog"
         "github.com/luci/luci-go/client/logdog/butlerlib/streamclient"
         "github.com/luci/luci-go/common/auth"
         "github.com/luci/luci-go/common/config"
         "github.com/luci/luci-go/common/ctxcmd"
         "github.com/luci/luci-go/common/environ"
         "github.com/luci/luci-go/common/errors"
+        "github.com/luci/luci-go/common/logdog/types"
         log "github.com/luci/luci-go/common/logging"
 
         "github.com/golang/protobuf/proto"
         "golang.org/x/net/context"
 )
 
+type cookLogDogParams struct {
+        host     string
+        project  string
+        prefix   types.StreamName
+        annotee  bool
+        filePath string
+}
+
+func (p *cookLogDogParams) addFlags(fs *flag.FlagSet) {
+        fs.StringVar(
+                &p.host,
+                "logdog-host",
+                "",
+                "The name of the LogDog host.")
+        fs.StringVar(
+                &p.project,
+                "logdog-project",
+                "",
+                "The name of the LogDog project to log into. Projects have different ACL sets, "+
+                        "so choose this appropriately.")
+        fs.Var(
+                &p.prefix,
+                "logdog-prefix",
+                "The LogDog stream Prefix to use. If empty, one will be constructed from the Swarming "+
+                        "task parameters (found in enviornment).")
+        fs.BoolVar(
+                &p.annotee,
+                "logdog-enable-annotee",
+                true,
+                "Process bootstrap STDOUT/STDERR annotations through Annotee.")
+        fs.StringVar(
+                &p.filePath,
+                "logdog-debug-out-file",
+                "",
+                "If specified, write all generated logs to this path instead of sending them.")
+}
+
+func (p *cookLogDogParams) active() bool {
+        return p.host != "" || p.project != "" || p.prefix != ""
+}
+
+func (p *cookLogDogParams) validate() error {
+        if p.project == "" {
+                return fmt.Errorf("a LogDog project must be supplied (-logdog-project)")
+        }
+        return nil
+}
+
+func (p *cookLogDogParams) getPrefix(env environ.Env) (types.StreamName, error) {
+        if p.prefix != "" {
+                return p.prefix, nil
+        }
+
+        // Construct our LogDog prefix from the Swarming task parameters.
+        host, _ := env.Get("SWARMING_SERVER")
+        if host == "" {
+                return "", errors.Reason("missing or empty SWARMING_SERVER").Err()
+        }
+        taskID, _ := env.Get("SWARMING_TASK_ID")
+        if taskID == "" {
+                return "", errors.Reason("missing or empty SWARMING_TASK_ID").Err()
+        }
+
+        return types.MakeStreamName("", "swarm", host, taskID)
+}
+
 // runWithLogdogButler rus the supplied command through the a LogDog Butler
 // engine instance. This involves:
 //      - Determine a LogDog Prefix.
 //      - Configuring / setting up the Butler.
 //      - Initiating a LogDog Pub/Sub Output, registering with remote server.
 //      - Running the recipe process.
 //        - Optionally, hook its output streams up through an Annotee processor.
 //        - Otherwise, wait for the process to finish.
 //      - Shut down the Butler instance.
 func (c *cookRun) runWithLogdogButler(ctx context.Context, cmd *exec.Cmd) (rc int, err error) {
-        _ = auth.Authenticator{}
+        // Get our task's environment.
+        var env environ.Env
+        if cmd.Env != nil {
+                env = environ.New(cmd.Env)
+        } else {
+                env = environ.System()
+        }
+
+        prefix, err := c.logdog.getPrefix(env)
+        if err != nil {
+                return 0, errors.Annotate(err).Reason("failed to get LogDog prefix").Err()
+        }
+        log.Fields{
+                "prefix": prefix,
+        }.Infof(ctx, "Using LogDog prefix: %q", prefix)
 
         authenticator := auth.NewAuthenticator(ctx, auth.SilentLogin, auth.Options{
                 Scopes: out.Scopes(),
         })
 
         // Register and instantiate our LogDog Output.
         var o output.Output
         if c.logdog.filePath == "" {
                 ocfg := out.Config{
                         Auth:    authenticator,
                         Host:    c.logdog.host,
                         Project: config.ProjectName(c.logdog.project),
-                        Prefix:  c.logdog.prefix,
+                        Prefix:  prefix,
                         SourceInfo: []string{
                                 "Kitchen",
                         },
                         PublishContext: withNonCancel(ctx),
                 }
 
                 var err error
                 if o, err = ocfg.Register(ctx); err != nil {
                         return 0, errors.Annotate(err).Reason("failed to create LogDog Output instance").Err()
                 }
@@ skipping 28 matching lines @@
                                 err = ierr
                         }
                 }
         }()
 
         // Wrap our incoming command in a CtxCmd.
         proc := ctxcmd.CtxCmd{
                 Cmd: cmd,
         }
 
-        var env environ.Env
-        if proc.Env != nil {
-                env = environ.New(proc.Env)
-        } else {
-                env = environ.System()
-        }
-
         // Augment our environment with Butler parameters.
         bsEnv := bootstrap.Environment{
                 Project: config.ProjectName(c.logdog.project),
                 Prefix:  c.logdog.prefix,
         }
         bsEnv.Augment(env)
         proc.Env = env.Sorted()
 
         // Build pipes for our STDOUT and STDERR streams.
         stdout, err := proc.StdoutPipe()
@@ skipping 134 matching lines @@
 // callbackReadCloser invokes a callback method when closed.
 type callbackReadCloser struct {
         io.ReadCloser
         callback func()
 }
 
 func (c *callbackReadCloser) Close() error {
         defer c.callback()
         return c.ReadCloser.Close()
 }