Fix Annotee stream name generation.

logdog/client/annotee/processor.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 annotee
 
 import (
         "bufio"
         "bytes"
         "fmt"
@@ skipping 98 matching lines @@
 
 // Processor consumes data from a list of Stream entries and interacts with the
 // supplied Client instance.
 //
 // A Processor must be instantiated with New.
 type Processor struct {
         ctx context.Context
         o   *Options
 
         astate       *annotation.State
-        stepHandlers map[string]*stepHandler
+        stepHandlers map[*annotation.Step]*stepHandler
 
         annotationStream    streamclient.Stream
         annotationC         chan annotationSignal
         annotationFinishedC chan struct{}
 }
 
 type annotationSignal struct {
         data       []byte
         updateType annotation.UpdateType
 }
 
 // New instantiates a new Processor.
 func New(c context.Context, o Options) *Processor {
         p := Processor{
                 ctx: c,
                 o:   &o,
 
-                stepHandlers: make(map[string]*stepHandler),
+                stepHandlers: make(map[*annotation.Step]*stepHandler),
         }
         p.astate = &annotation.State{
                 LogNameBase: o.Base,
                 Callbacks:   &annotationCallbacks{&p},
                 Execution:   o.Execution,
                 Clock:       clock.Get(c),
                 Offline:     o.Offline,
         }
         return &p
 }
@@ skipping 78 matching lines @@
                 log.WithError(err).Errorf(p.ctx, "Failed to initialize.")
                 return err
         }
 
         a := extractAnnotation(line)
         if a != "" {
                 log.Debugf(p.ctx, "Annotation: %q", a)
         }
 
         var step *annotation.Step
-        var h *stepHandler
         if s.Annotate {
                 if a != "" {
                         // Append our annotation to the annotation state. This may cause our
                         // annotation callbacks to be invoked.
                         if err := p.astate.Append(a); err != nil {
                                 log.Fields{
                                         log.ErrorKey: err,
                                         "stream":     s.Name,
                                         "annotation": a,
                                 }.Errorf(p.ctx, "Failed to process annotation.")
                         }
                 }
 
                 // Use the step handler for the current step.
                 step = p.astate.CurrentStep()
         } else {
                 // Not handling annotations. Use our root step handler.
                 step = p.astate.RootStep()
         }
 
         h, err := p.getStepHandler(step, true)
         if err != nil {
                 return err
         }
 
         // Build our output, which will consist of the initial line and any extra
         // lines that have been registered.
         inject := h.flushInjectedLines()
-        output := make([]string, 1, 1+len(inject))
-        output[0] = line
-        output = append(output, inject...)
+        stepOutput := make([]string, 0, 1+len(inject))
 
-        for _, l := range output {
-                // If configured, tee to our tee stream.
-                if s.Tee != nil && (a == "" || !s.StripAnnotations) {
-                        // Tee this to the Stream's configured Tee output.
-                        if err := writeTextLine(s.Tee, l); err != nil {
-                                log.WithError(err).Errorf(p.ctx, "Failed to tee line.")
+        // If this is an annotation line, write it to our root handler. Otherwise,
+        // write it to our step's handler (by appending it to stepOutput).
+        if a != "" {
+                // If we're not stripping annotations, emit this to the root handler.
+                if !s.StripAnnotations {
+                        // Get our root log stream handler. As an optimization, if "step" is
+                        // the root step, then "h" is already the root handler, so we don't need
+                        // to duplicate the lookup.
+                        var rootHandler *stepHandler
+                        if rootStep := p.astate.RootStep(); rootStep != step {
+                                rootHandler, err = p.getStepHandler(rootStep, true)
+                                if err != nil {
+                                        return err
+                                }
+                        } else {
+                                rootHandler = h
+                        }
+
+                        if err := rootHandler.writeBaseStream(s, line); err != nil {
+                                log.WithError(err).Errorf(p.ctx, "Failed to send line to LogDog.")
                                 return err
                         }
                 }
+        } else {
+                stepOutput = append(stepOutput, line)
+        }
 
+        // Add any injected lines.
+        stepOutput = append(stepOutput, inject...)
+
+        for _, l := range stepOutput {
                 // Write to our LogDog stream.
                 if err := h.writeBaseStream(s, l); err != nil {
                         log.WithError(err).Errorf(p.ctx, "Failed to send line to LogDog.")
                         return err
                 }
         }
 
         return err
 }
 
@@ skipping 19 matching lines @@
                 if err := p.annotationStream.Close(); err != nil {
                         log.WithError(err).Errorf(p.ctx, "Failed to close annotation stream.")
                 }
                 p.annotationStream = nil
         }
 
         return p.astate
 }
 
 func (p *Processor) getStepHandler(step *annotation.Step, create bool) (*stepHandler, error) {
-        name := step.CanonicalName()
-        if h := p.stepHandlers[name]; h != nil {
+        if h := p.stepHandlers[step]; h != nil {
                 return h, nil
         }
         if !create {
                 return nil, nil
         }
 
         h, err := newStepHandler(p, step)
         if err != nil {
                 log.Fields{
                         log.ErrorKey: err,
-                        "step":       name,
+                        "step":       step,
                 }.Errorf(p.ctx, "Failed to create step handler.")
                 return nil, err
         }
-        p.stepHandlers[name] = h
+        p.stepHandlers[step] = h
         return h, nil
 }
 
 func (p *Processor) closeStep(step *annotation.Step) {
         if h, _ := p.getStepHandler(step, false); h != nil {
                 p.finishStepHandler(h, true)
         }
 }
 
 func (p *Processor) finishStepHandler(h *stepHandler, closeSteps bool) {
         // Remove this handler from our list. This will stop us from
         // double-finishing when finish() calls Close(), which calls the StepClosed
         // callback.
-        delete(p.stepHandlers, h.String())
+        delete(p.stepHandlers, h.step)
 
         // Finish the step.
         h.finish(closeSteps)
 }
 
 func (p *Processor) createStream(name types.StreamName, archetype *streamproto.Flags) (streamclient.Stream, error) {
         return p.o.Client.NewStream(streamFlagsFromArchetype(p.ctx, name, archetype))
 }
 
 func (p *Processor) annotationStateUpdated(ut annotation.UpdateType) {
@@ skipping 119 matching lines @@
         step      *annotation.Step
 
         client        streamclient.Client
         injectedLines []string
         streams       map[types.StreamName]streamclient.Stream
         finished      bool
 }
 
 func newStepHandler(p *Processor, step *annotation.Step) (*stepHandler, error) {
         h := stepHandler{
-                Context:   log.SetField(p.ctx, "step", step.CanonicalName()),
+                Context:   log.SetField(p.ctx, "step", step),
                 processor: p,
                 step:      step,
 
                 client:  p.o.Client,
                 streams: make(map[types.StreamName]streamclient.Stream),
         }
 
         // Send our initial annotation state.
         h.updated(annotation.UpdateStructural)
 
         return &h, nil
 }
 
 func (h *stepHandler) String() string {
-        return h.step.CanonicalName()
+        return h.step.String()
 }
 
 func (h *stepHandler) finish(closeSteps bool) {
         if h.finished {
                 return
         }
 
         if closeSteps {
                 h.step.Close(nil)
         }
 
         // Close all streams associated with this handler.
         if closeSteps {
                 h.closeAllStreams()
         }
 
         // Notify that the annotation state has updated (closed).
         h.processor.annotationStateUpdated(annotation.UpdateStructural)
         h.finished = true
 }
 
 func (h *stepHandler) writeBaseStream(s *Stream, line string) error {
+        // If we're teeing, also write this to our Processor's tee stream.
+        if s.Tee != nil {
+                // Tee this to the Stream's configured Tee output.
+                if err := writeTextLine(s.Tee, line); err != nil {
+                        log.WithError(err).Errorf(h, "Failed to tee line.")
+                        return err
+                }
+        }
+
         name := h.step.BaseStream(s.Name)
         stream, created, err := h.getStream(name, &textStreamArchetype)
         if err != nil {
                 return err
         }
         if created {
                 switch s.Name {
                 case STDOUT:
                         if h.step.SetSTDOUTStream(&milo.LogdogStream{Name: string(name)}) {
                                 h.updated(annotation.UpdateIterative)
@@ skipping 126 matching lines @@
                 return ""
         }
         return strings.TrimSpace(line[3 : len(line)-3])
 }
 
 func buildAnnotation(name string, params ...string) string {
         v := make([]string, 1, 1+len(params))
         v[0] = name
         return "@@@" + strings.Join(append(v, params...), "@") + "@@@"
 }