cmd/trace: distinguish task endTimestamp and lastTimestamp

A task may have other user annotation events after the task ends.
So far, task.lastTimestamp returned the task end event if the
event available. This change introduces task.endTimestamp for that
and makes task.lastTimestamp returns the "last" seen event's timestamp
if the task is ended.

If the task is not ended, both returns the last timestamp of the entire
trace assuming the task is still active.

This fixes the task-oriented trace view mode not to drop user
annotation instances when they appear outside a task's lifespan.
Adds a test.

Change-Id: Iba1062914f224edd521b9ee55c6cd5e180e55359
Reviewed-on: https://go-review.googlesource.com/109175
Reviewed-by: Heschi Kreinick <heschi@google.com>

3 files changed, 89 insertions, 17 deletions

diff --git a/src/cmd/trace/annotations.go b/src/cmd/trace/annotations.go
index 5f672de3b4..dcec42eb5b 100644
--- a/src/cmd/trace/annotations.go
+++ b/src/cmd/trace/annotations.go
@@ -230,7 +230,7 @@ func httpUserTask(w http.ResponseWriter, r *http.Request) {
 			Complete:   task.complete(),
 			Events:     events,
 			Start:      time.Duration(task.firstTimestamp()) * time.Nanosecond,
-			End:        time.Duration(task.lastTimestamp()) * time.Nanosecond,
+			End:        time.Duration(task.endTimestamp()) * time.Nanosecond,
 			GCTime:     task.overlappingGCDuration(res.gcEvents),
 		})
 	}
@@ -365,7 +365,7 @@ func (task *taskDesc) String() string {
 	}
 	wb := new(bytes.Buffer)
 	fmt.Fprintf(wb, "task %d:\t%s\n", task.id, task.name)
-	fmt.Fprintf(wb, "\tstart: %v end: %v complete: %t\n", task.firstTimestamp(), task.lastTimestamp(), task.complete())
+	fmt.Fprintf(wb, "\tstart: %v end: %v complete: %t\n", task.firstTimestamp(), task.endTimestamp(), task.complete())
 	fmt.Fprintf(wb, "\t%d goroutines\n", len(task.goroutines))
 	fmt.Fprintf(wb, "\t%d regions:\n", len(task.regions))
 	for _, s := range task.regions {
@@ -463,6 +463,17 @@ func (task *taskDesc) firstTimestamp() int64 {
 // trace. If the trace does not contain the task end event, the last
 // timestamp of the trace will be returned.
 func (task *taskDesc) lastTimestamp() int64 {
+	endTs := task.endTimestamp()
+	if last := task.lastEvent(); last != nil && last.Ts > endTs {
+		return last.Ts
+	}
+	return endTs
+}
+
+// endTimestamp returns the timestamp of this task's end event.
+// If the trace does not contain the task end event, the last
+// timestamp of the trace will be returned.
+func (task *taskDesc) endTimestamp() int64 {
 	if task != nil && task.end != nil {
 		return task.end.Ts
 	}
@@ -470,7 +481,7 @@ func (task *taskDesc) lastTimestamp() int64 {
 }
 
 func (task *taskDesc) duration() time.Duration {
-	return time.Duration(task.lastTimestamp()-task.firstTimestamp()) * time.Nanosecond
+	return time.Duration(task.endTimestamp()-task.firstTimestamp()) * time.Nanosecond
 }
 
 func (region *regionDesc) duration() time.Duration {
@@ -496,13 +507,13 @@ func (task *taskDesc) overlappingGCDuration(evs []*trace.Event) (overlapping tim
 // any of the task's region if ev is a goroutine-local event, or overlaps with the
 // task's lifetime if ev is a global event.
 func (task *taskDesc) overlappingInstant(ev *trace.Event) bool {
-	if isUserAnnotationEvent(ev) && task.id != ev.Args[0] {
+	if _, ok := isUserAnnotationEvent(ev); ok && task.id != ev.Args[0] {
 		return false // not this task's user event.
 	}
 
 	ts := ev.Ts
 	taskStart := task.firstTimestamp()
-	taskEnd := task.lastTimestamp()
+	taskEnd := task.endTimestamp()
 	if ts < taskStart || taskEnd < ts {
 		return false
 	}
@@ -547,7 +558,7 @@ func (task *taskDesc) overlappingDuration(ev *trace.Event) (time.Duration, bool)
 	// This event is a global GC event
 	if ev.P == trace.GCP {
 		taskStart := task.firstTimestamp()
-		taskEnd := task.lastTimestamp()
+		taskEnd := task.endTimestamp()
 		o := overlappingDuration(taskStart, taskEnd, start, end)
 		return o, o > 0
 	}
@@ -627,7 +638,7 @@ func (region *regionDesc) lastTimestamp() int64 {
 // If non-zero depth is provided, this searches all events with BFS and includes
 // goroutines unblocked any of related goroutines to the result.
 func (task *taskDesc) RelatedGoroutines(events []*trace.Event, depth int) map[uint64]bool {
-	start, end := task.firstTimestamp(), task.lastTimestamp()
+	start, end := task.firstTimestamp(), task.endTimestamp()
 
 	gmap := map[uint64]bool{}
 	for k := range task.goroutines {
@@ -1127,12 +1138,12 @@ func describeEvent(ev *trace.Event) string {
 	return ""
 }
 
-func isUserAnnotationEvent(ev *trace.Event) bool {
+func isUserAnnotationEvent(ev *trace.Event) (taskID uint64, ok bool) {
 	switch ev.Type {
 	case trace.EvUserLog, trace.EvUserRegion, trace.EvUserTaskCreate, trace.EvUserTaskEnd:
-		return true
+		return ev.Args[0], true
 	}
-	return false
+	return 0, false
 }
 
 var templUserRegionType = template.Must(template.New("").Funcs(template.FuncMap{
diff --git a/src/cmd/trace/trace.go b/src/cmd/trace/trace.go
index 19d6fa36a7..96e912159b 100644
--- a/src/cmd/trace/trace.go
+++ b/src/cmd/trace/trace.go
@@ -751,7 +751,7 @@ func generateTrace(params *traceParams, consumer traceConsumer) error {
 			}
 			ctx.emit(tBegin)
 
-			tEnd := &ViewerEvent{Category: "task", Name: taskName, Phase: "e", Time: float64(task.lastTimestamp()) / 1e3, Tid: taskRow, ID: task.id, Cname: colorBlue}
+			tEnd := &ViewerEvent{Category: "task", Name: taskName, Phase: "e", Time: float64(task.endTimestamp()) / 1e3, Tid: taskRow, ID: task.id, Cname: colorBlue}
 			if task.end != nil {
 				tEnd.Stack = ctx.stack(task.end.Stk)
 			}
@@ -955,17 +955,19 @@ func (ctx *traceContext) emitThreadCounters(ev *trace.Event) {
 func (ctx *traceContext) emitInstant(ev *trace.Event, name, category string) {
 	cname := ""
 	if ctx.mode&modeTaskOriented != 0 {
-		overlapping := false
+		taskID, isUserAnnotation := isUserAnnotationEvent(ev)
+
+		show := false
 		for _, task := range ctx.tasks {
-			if task.overlappingInstant(ev) {
-				overlapping = true
+			if isUserAnnotation && task.id == taskID || task.overlappingInstant(ev) {
+				show = true
 				break
 			}
 		}
 		// grey out or skip if non-overlapping instant.
-		if !overlapping {
-			if isUserAnnotationEvent(ev) {
-				return // don't display unrelated task events.
+		if !show {
+			if isUserAnnotation {
+				return // don't display unrelated user annotation events.
 			}
 			cname = colorLightGrey
 		}
diff --git a/src/cmd/trace/trace_test.go b/src/cmd/trace/trace_test.go
index aff3863802..852d745b29 100644
--- a/src/cmd/trace/trace_test.go
+++ b/src/cmd/trace/trace_test.go
@@ -5,8 +5,10 @@
 package main
 
 import (
+	"context"
 	"internal/trace"
 	"io/ioutil"
+	rtrace "runtime/trace"
 	"strings"
 	"testing"
 )
@@ -172,3 +174,60 @@ func TestPreemptedMarkAssist(t *testing.T) {
 		t.Errorf("Got %v MARK ASSIST events, want %v", marks, 2)
 	}
 }
+
+func TestFoo(t *testing.T) {
+	prog0 := func() {
+		ctx, task := rtrace.NewTask(context.Background(), "ohHappyDay")
+		rtrace.Log(ctx, "", "log before task ends")
+		task.End()
+		rtrace.Log(ctx, "", "log after task ends") // log after task ends
+	}
+	if err := traceProgram(t, prog0, "TestFoo"); err != nil {
+		t.Fatalf("failed to trace the program: %v", err)
+	}
+	res, err := parseTrace()
+	if err != nil {
+		t.Fatalf("failed to parse the trace: %v", err)
+	}
+	annotRes, _ := analyzeAnnotations()
+	var task *taskDesc
+	for _, t := range annotRes.tasks {
+		if t.name == "ohHappyDay" {
+			task = t
+			break
+		}
+	}
+	if task == nil {
+		t.Fatal("failed to locate expected task event")
+	}
+
+	params := &traceParams{
+		parsed:    res,
+		mode:      modeTaskOriented,
+		startTime: task.firstTimestamp() - 1,
+		endTime:   task.lastTimestamp() + 1,
+		tasks:     []*taskDesc{task},
+	}
+
+	c := viewerDataTraceConsumer(ioutil.Discard, 0, 1<<63-1)
+
+	var logBeforeTaskEnd, logAfterTaskEnd bool
+	c.consumeViewerEvent = func(ev *ViewerEvent, _ bool) {
+		if ev.Name == "log before task ends" {
+			logBeforeTaskEnd = true
+		}
+		if ev.Name == "log after task ends" {
+			logAfterTaskEnd = true
+		}
+	}
+	if err := generateTrace(params, c); err != nil {
+		t.Fatalf("generateTrace failed: %v", err)
+	}
+	if !logBeforeTaskEnd {
+		t.Error("failed to find 'log before task ends'")
+	}
+	if !logAfterTaskEnd {
+		t.Error("failed to find 'log after task ends'")
+	}
+
+}