Add infraenv, setup creds and tee for Kitchen.

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"
+        "net/url"
+        "os"
         "os/exec"
         "sync"
         "time"
 
+        "infra/libs/infraenv"
+
         "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/client/logdog/butlerlib/streamproto"
         "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
+        tee      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.BoolVar(
+                &p.tee,
+                "logdog-tee",
+                true,
+                "Tee bootstrapped STDOUT and STDERR through Kitchen. If false, these will only be sent as LogDog streams")
         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 != ""
+        return p.host != "" || p.project != "" || p.prefix != "" || p.tee
 }
 
 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")
+        // Construct our LogDog prefix from the Swarming task parameters. The server
+        // will be exported as a URL. We want the "host" parameter from this URL.
+        server, _ := env.Get("SWARMING_SERVER")
+        serverURL, err := url.Parse(server)
+        if err != nil {
+                return "", errors.Annotate(err).Reason("failed to parse SWARMING_SERVER URL %(value)q").
+                        D("value", server).Err()
+        }
+
+        host := serverURL.Host
+        if serverURL.Scheme == "" {
+                // SWARMING_SERVER is not a full URL, so use its Path instead of its Host.
+                host = serverURL.Path
+        }
         if host == "" {
-                return "", errors.Reason("missing or empty SWARMING_SERVER").Err()
+                return "", errors.Reason("missing or empty SWARMING_SERVER host in %(value)q").
+                        D("value", 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:
@@ skipping 14 matching lines @@
         }
 
         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{
+        // Set up authentication.
+        authOpts := auth.Options{
                 Scopes: out.Scopes(),
-        })
+        }
+        if !infraenv.OnGCE() {
+                // If we're not on GCE, we will need to explicitly supply the LogDog
+                // credentials path.
+                credPath, err := infraenv.GetLogDogServiceAccountJSON()
+                if err != nil {
+                        return 0, errors.Annotate(err).Reason("failed to get LogDog service account JSON path").Err()
+                }
+                authOpts.ServiceAccountJSONPath = credPath
+        }
+        authenticator := auth.NewAuthenticator(ctx, auth.SilentLogin, authOpts)
 
         // 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:  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()
                 }
         } else {
                 // Debug: Use a file output.
                 ocfg := fileOut.Options{
                         Path: c.logdog.filePath,
                 }
                 o = ocfg.New(ctx)
         }
         defer o.Close()
 
-        ncCtx := withNonCancel(ctx)
-        b, err := butler.New(ncCtx, butler.Config{
+        butlerCfg := butler.Config{
                 Output:     o,
                 Project:    config.ProjectName(c.logdog.project),
                 Prefix:     c.logdog.prefix,
                 BufferLogs: true,
-        })
+        }
+
+        // If we're teeing and we're not using Annotee, tee our subprocess' STDOUT
+        // and STDERR through Kitchen's STDOUT/STDERR.
+        //
+        // If we're using Annotee, we will configure this in the Annotee setup
+        // directly.
+        if c.logdog.tee && !c.logdog.annotee {
+                butlerCfg.TeeStdout = os.Stdout
+                butlerCfg.TeeStderr = os.Stderr
+        }
+
+        ncCtx := withNonCancel(ctx)
+        b, err := butler.New(ncCtx, butlerCfg)
         if err != nil {
                 err = errors.Annotate(err).Reason("failed to create Butler instance").Err()
                 return
         }
         defer func() {
                 b.Activate()
                 if ierr := b.Wait(); ierr != nil {
                         ierr = errors.Annotate(ierr).Reason("failed to Wait() for Butler").Err()
                         logAnnotatedErr(ctx, ierr)
 
@@ skipping 88 matching lines @@
                                 Annotate:         true,
                                 StripAnnotations: true,
                         },
                         {
                                 Reader:           stderr,
                                 Name:             annotee.STDERR,
                                 Annotate:         true,
                                 StripAnnotations: true,
                         },
                 }
+                if c.logdog.tee {
+                        streams[0].Tee = os.Stdout
+                        streams[1].Tee = os.Stderr
+                }
 
                 // Run the process' output streams through Annotee. This will block until
                 // they are all consumed.
                 if err = annoteeProcessor.RunStreams(streams); err != nil {
                         err = errors.Annotate(err).Reason("failed to process streams through Annotee").Err()
                         return
                 }
         } else {
                 // Get our STDOUT / STDERR stream flags. Tailor them to match Annotee.
                 stdoutFlags := annotee.TextStreamFlags(ctx, annotee.STDOUT)
                 stderrFlags := annotee.TextStreamFlags(ctx, annotee.STDERR)
+                if c.logdog.tee {
+                        stdoutFlags.Tee = streamproto.TeeStdout
+                        stderrFlags.Tee = streamproto.TeeStderr
+                }
 
                 // Wait for our STDOUT / STDERR streams to complete.
                 var wg sync.WaitGroup
                 stdout = &callbackReadCloser{stdout, wg.Done}
                 stderr = &callbackReadCloser{stderr, wg.Done}
                 wg.Add(2)
 
                 // Explicitly add these streams to the Butler.
                 if err = b.AddStream(stdout, *stdoutFlags.Properties()); err != nil {
                         err = errors.Annotate(err).Reason("failed to add STDOUT stream to Butler").Err()
@@ skipping 37 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()
 }