milo: fix running steps

client/logdog/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 90 matching lines @@
         //      - If this is < 0, metadata will only be pushed at the beginning and end of
         //        a step.
         //      - If this equals 0, metadata will be pushed every time it's updated.
         //      - If this is 0, DefaultMetadataUpdateInterval will be used.
         MetadataUpdateInterval time.Duration
 
         // Offline specifies whether parsing happens not at the same time as
         // emitting. If true and CURRENT_TIMESTAMP annotations are not provided
         // then step start/end times are left empty.
         Offline bool
+
+        // CloseSteps specified whether outstanding open steps must be closed.
+        CloseSteps bool
 }
 
 // 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
 
@@ skipping 123 matching lines @@
                 // 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
 }
 
-// Finish instructs the Processor to close any outstanding state. This should be
-// called when all automatic state updates have completed in case any steps
-// didn't properly close their state.
+// Finish instructs the Processor to finish any outstanding state.
+// It is mandatory to call Finish.
 func (p *Processor) Finish() *annotation.State {
-        // Close our step handlers.
+        // Finish our step handlers.
         for _, h := range p.stepHandlers {
-                p.closeStepHandler(h)
+                p.finishStepHandler(h, p.o.CloseSteps)
         }
         return p.astate
 }
 
 func (p *Processor) getStepHandler(step *annotation.Step, create bool) (*stepHandler, error) {
         name := step.CanonicalName()
         if h := p.stepHandlers[name]; 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,
                 }.Errorf(p.ctx, "Failed to create step handler.")
                 return nil, err
         }
         p.stepHandlers[name] = h
         return h, nil
 }
 
 func (p *Processor) closeStep(step *annotation.Step) {
         if h, _ := p.getStepHandler(step, false); h != nil {
-                p.closeStepHandler(h)
+                p.finishStepHandler(h, true)
         }
 }
 
-func (p *Processor) closeStepHandler(h *stepHandler) {
+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())
 
         // Finish the step.
-        h.finish()
+        h.finish(closeSteps)

[hinoka, 2016/06/16 21:13:58]

I don' think you have to pass it in here either

[nodir, 2016/06/16 21:23:14]

it does not always equal p.o.CloseSteps, e.g. 2 hunks above

[nodir, 2016/06/16 21:23:14]

why?

 }
 
 type annotationCallbacks struct {
         *Processor
 }
 
 func (c *annotationCallbacks) StepClosed(step *annotation.Step) {
         c.closeStep(step)
 }
 
@@ skipping 41 matching lines @@
         context.Context
 
         processor *Processor
         step      *annotation.Step
 
         client              streamclient.Client
         injectedLines       []string
         streams             map[types.StreamName]streamclient.Stream
         annotationC         chan []byte
         annotationFinishedC chan struct{}
-        closed              bool
+        finished            bool
 }
 
 func newStepHandler(p *Processor, step *annotation.Step) (*stepHandler, error) {
         h := stepHandler{
                 Context:   log.SetField(p.ctx, "step", step.CanonicalName()),
                 processor: p,
                 step:      step,
 
                 client:              p.o.Client,
                 streams:             make(map[types.StreamName]streamclient.Stream),
@@ skipping 76 matching lines @@
                                 timerRunning = true
                         }
 
                 case <-t.GetC():
                         timerRunning = false
                         sendLatest()
                 }
         }
 }
 
-func (h *stepHandler) finish() {
-        if h.closed {
+func (h *stepHandler) finish(closeSteps bool) {
+        if h.finished {
                 return
         }
 
-        // This may send one last annotation to summarize the
-        // state if closing changed it.
-        if h.step.Close(nil) {
-                // Manually mark it updated, since Close callbacks will have unregistered
-                // us from the standard Updated() reporting loop.
-                h.updated()
+        if closeSteps {
+                h.step.Close(nil)
         }
-
+        // Send last annotation unconditionally.
+        // It may send the same annotation it sent last time; could be optimized.
+        h.sendAnnotation()
         // Close and reap our meter goroutine.
         close(h.annotationC)
         <-h.annotationFinishedC
 
         // Close all streams associated with this handler.
-        h.closeAllStreams()
+        if closeSteps {
+                h.closeAllStreams()
+        }
 
-        h.closed = true
+        h.finished = true
 }
 
 func (h *stepHandler) writeBaseStream(s *Stream, line string) error {
         name := h.step.BaseStream(s.Name)
         stream, created, err := h.getStream(name, &textStreamArchetype)
         if err != nil {
                 return err
         }
         if created {
                 segs := s.Name.Segments()
                 h.maybeInjectLink("stdio", segs[len(segs)-1], name)
         }
         return writeTextLine(stream, line)
 }
 
 func (h *stepHandler) updated() {
-        // Ignore updates after the step has closed.
-        if h.closed {
-                return
+        if !h.finished {
+                h.sendAnnotation()
         }
+}
 
+func (h *stepHandler) sendAnnotation() {
         // Serialize immediately, as the Step's internal state may change in future
         // annotation runs.
         p := h.step.Proto()
 
         data, err := proto.Marshal(p)
         if err != nil {
                 log.WithError(err).Errorf(h, "Failed to marshal state.")
                 return
         }
         h.annotationC <- data
 }
 
 func (h *stepHandler) getStream(name types.StreamName, flags *streamproto.Flags) (s streamclient.Stream, created bool, err error) {
-        if h.closed {
-                err = fmt.Errorf("refusing to get stream %q for closed handler", name)
+        if h.finished {
+                err = fmt.Errorf("refusing to get stream %q for finished handler", name)
                 return
         }
         if s = h.streams[name]; s != nil {
                 return
         }
 
         // New stream, will be creating.
         if flags == nil {
                 return
         }
@@ skipping 100 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...), "@") + "@@@"
 }