LogDog/Kitchen: Misc fixes.

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"
@@ skipping 20 matching lines @@
         "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
+        host    string
+        project string
+        prefix  types.StreamName
+        annotee bool
+        tee     bool
+
+        filePath               string
+        serviceAccountJSONPath string
 }
 
 func (p *cookLogDogParams) addFlags(fs *flag.FlagSet) {
         fs.StringVar(
                 &p.host,
                 "logdog-host",
                 "",
                 "The name of the LogDog host.")
         fs.StringVar(
                 &p.project,
@@ skipping 14 matching lines @@
         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.")
+        fs.StringVar(
+                &p.serviceAccountJSONPath,
+                "logdog-service-account-json-path",
+                "",
+                "If specified, use the service account JSON file at this path. Otherwise, autodetect.")
 }
 
 func (p *cookLogDogParams) active() bool {
         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)")
         }
@@ skipping 55 matching lines @@
                 return 0, errors.Annotate(err).Reason("failed to get LogDog prefix").Err()
         }
         log.Fields{
                 "prefix": prefix,
         }.Infof(ctx, "Using LogDog prefix: %q", prefix)
 
         // 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.
+        switch {
+        case c.logdog.serviceAccountJSONPath != "":
+                authOpts.ServiceAccountJSONPath = c.logdog.serviceAccountJSONPath
+
+        case infraenv.OnGCE():
+                // Do nothing, auth will automatically use GCE metadata.
+                break
+
+        default:
+                // No service account specified, so load the LogDog credentials from the
+                // local bot deployment.
                 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
@@ skipping 16 matching lines @@
         } else {
                 // Debug: Use a file output.
                 ocfg := fileOut.Options{
                         Path: c.logdog.filePath,
                 }
                 o = ocfg.New(ctx)
         }
         defer o.Close()
 
         butlerCfg := butler.Config{
-                Output:     o,
-                Project:    config.ProjectName(c.logdog.project),
-                Prefix:     c.logdog.prefix,
-                BufferLogs: true,
+                Output:       o,
+                Project:      config.ProjectName(c.logdog.project),
+                Prefix:       prefix,
+                BufferLogs:   true,
+                MaxBufferAge: butler.DefaultMaxBufferAge,
         }
 
         // 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
@@ skipping 88 matching lines @@
                         Offline:                false,
                         CloseSteps:             true,
                 })
                 defer func() {
                         as := annoteeProcessor.Finish()
                         log.Infof(ctx, "Annotations finished:\n%s", proto.MarshalTextString(as.RootStep().Proto()))
                 }()
 
                 // Run STDOUT/STDERR streams through the processor. This will block until
                 // both streams are closed.
+                //
+                // If we're teeing, we will tee the full stream, including annotations.
                 streams := []*annotee.Stream{
                         {
                                 Reader:           stdout,
                                 Name:             annotee.STDOUT,
                                 Annotate:         true,
-                                StripAnnotations: true,
+                                StripAnnotations: !c.logdog.tee,
                         },
                         {
                                 Reader:           stderr,
                                 Name:             annotee.STDERR,
                                 Annotate:         true,
-                                StripAnnotations: true,
+                                StripAnnotations: !c.logdog.tee,
                         },
                 }
                 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 {
@@ skipping 58 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()
 }