runtime/trace: implement annotation API

This implements the annotation API proposed in golang.org/cl/63274.

traceString is updated to protect the string map with trace.stringsLock
because the assumption that traceString is called by a single goroutine
(either at the beginning of tracing and at the end of tracing when
dumping all the symbols and function names) is no longer true.

traceString is used by the annotation apis (NewContext, StartSpan, Log)
to register frequently appearing strings (task and span names, and log
keys) after this change.

NewContext -> one or two records (EvString, EvUserTaskCreate)
end function -> one record (EvUserTaskEnd)
StartSpan -> one or two records (EvString, EvUserSpan)
span end function -> one or two records (EvString, EvUserSpan)
Log -> one or two records (EvString, EvUserLog)

EvUserLog record is of the typical record format written by traceEvent
except that it is followed by bytes that represents the value string.

In addition to runtime/trace change, this change includes
corresponding changes in internal/trace to parse the new record types.

Future work to improve efficiency:
  More efficient unique task id generation instead of atomic. (per-P
  counter).
  Instead of a centralized trace.stringsLock, consider using per-P
  string cache or something more efficient.

R=go1.11

Change-Id: Iec9276c6c51e5be441ccd52dec270f1e3b153970
Reviewed-on: https://go-review.googlesource.com/71690Reviewed-by: Austin Clements <austin@google.com>

src/runtime/trace.go

@@ -128,11 +128,13 @@ var trace struct {
 
 	// Dictionary for traceEvString.
 	//
-	// Currently this is used only at trace setup and for
-	// func/file:line info after tracing session, so we assume
-	// single-threaded access.
-	strings   map[string]uint64
-	stringSeq uint64
+	// TODO: central lock to access the map is not ideal.
+	//   option: pre-assign ids to all user annotation span names and tags
+	//   option: per-P cache
+	//   option: sync.Map like data structure
+	stringsLock mutex
+	strings     map[string]uint64
+	stringSeq   uint64
 
 	// markWorkerLabels maps gcMarkWorkerMode to string ID.
 	markWorkerLabels [len(gcMarkWorkerModeStrings)]uint64
@@ -513,12 +515,26 @@ func traceEvent(ev byte, skip int, args ...uint64) {
 	// so if we see trace.enabled == true now, we know it's true for the rest of the function.
 	// Exitsyscall can run even during stopTheWorld. The race with StartTrace/StopTrace
 	// during tracing in exitsyscall is resolved by locking trace.bufLock in traceLockBuffer.
+	//
+	// Note trace_userTaskCreate runs the same check.
 	if !trace.enabled && !mp.startingtrace {
 		traceReleaseBuffer(pid)
 		return
 	}
+
+	if skip > 0 {
+		if getg() == mp.curg {
+			skip++ // +1 because stack is captured in traceEventLocked.
+		}
+	}
+	traceEventLocked(0, mp, pid, bufp, ev, skip, args...)
+	traceReleaseBuffer(pid)
+}
+
+func traceEventLocked(extraBytes int, mp *m, pid int32, bufp *traceBufPtr, ev byte, skip int, args ...uint64) {
 	buf := (*bufp).ptr()
-	const maxSize = 2 + 5*traceBytesPerNumber // event type, length, sequence, timestamp, stack id and two add params
+	// TODO: test on non-zero extraBytes param.
+	maxSize := 2 + 5*traceBytesPerNumber + extraBytes // event type, length, sequence, timestamp, stack id and two add params
 	if buf == nil || len(buf.arr)-buf.pos < maxSize {
 		buf = traceFlush(traceBufPtrOf(buf), pid).ptr()
 		(*bufp).set(buf)
@@ -561,7 +577,6 @@ func traceEvent(ev byte, skip int, args ...uint64) {
 		// Fill in actual length.
 		*lenp = byte(evSize - 2)
 	}
-	traceReleaseBuffer(pid)
 }
 
 func traceStackID(mp *m, buf []uintptr, skip int) uint64 {
@@ -643,7 +658,20 @@ func traceString(bufp *traceBufPtr, pid int32, s string) (uint64, *traceBufPtr)
 	if s == "" {
 		return 0, bufp
 	}
+
+	lock(&trace.stringsLock)
+	if raceenabled {
+		// raceacquire is necessary because the map access
+		// below is race annotated.
+		raceacquire(unsafe.Pointer(&trace.stringsLock))
+	}
+
 	if id, ok := trace.strings[s]; ok {
+		if raceenabled {
+			racerelease(unsafe.Pointer(&trace.stringsLock))
+		}
+		unlock(&trace.stringsLock)
+
 		return id, bufp
 	}
 
@@ -651,6 +679,11 @@ func traceString(bufp *traceBufPtr, pid int32, s string) (uint64, *traceBufPtr)
 	id := trace.stringSeq
 	trace.strings[s] = id
 
+	if raceenabled {
+		racerelease(unsafe.Pointer(&trace.stringsLock))
+	}
+	unlock(&trace.stringsLock)
+
 	// memory allocation in above may trigger tracing and
 	// cause *bufp changes. Following code now works with *bufp,
 	// so there must be no memory allocation or any activities
@@ -664,8 +697,16 @@ func traceString(bufp *traceBufPtr, pid int32, s string) (uint64, *traceBufPtr)
 	}
 	buf.byte(traceEvString)
 	buf.varint(id)
-	buf.varint(uint64(len(s)))
-	buf.pos += copy(buf.arr[buf.pos:], s)
+
+	// double-check the string and the length can fit.
+	// Otherwise, truncate the string.
+	slen := len(s)
+	if room := len(buf.arr) - buf.pos; room < slen+traceBytesPerNumber {
+		slen = room
+	}
+
+	buf.varint(uint64(slen))
+	buf.pos += copy(buf.arr[buf.pos:], s[:slen])
 
 	(*bufp).set(buf)
 	return id, bufp
@@ -1109,24 +1150,72 @@ func traceNextGC() {
 
 //go:linkname trace_userTaskCreate runtime/trace.userTaskCreate
 func trace_userTaskCreate(id, parentID uint64, taskType string) {
-	// TODO: traceEvUserTaskCreate
-	// TODO: truncate the name if too long.
+	if !trace.enabled {
+		return
+	}
+
+	// Same as in traceEvent.
+	mp, pid, bufp := traceAcquireBuffer()
+	if !trace.enabled && !mp.startingtrace {
+		traceReleaseBuffer(pid)
+		return
+	}
+
+	typeStringID, bufp := traceString(bufp, pid, taskType)
+	traceEventLocked(0, mp, pid, bufp, traceEvUserTaskCreate, 3, id, parentID, typeStringID)
+	traceReleaseBuffer(pid)
 }
 
 //go:linkname trace_userTaskEnd runtime/trace.userTaskEnd
 func trace_userTaskEnd(id uint64) {
-	// TODO: traceEvUserSpan
+	traceEvent(traceEvUserTaskEnd, 2, id)
 }
 
 //go:linkname trace_userSpan runtime/trace.userSpan
-func trace_userSpan(id, mode uint64, spanType string) {
-	// TODO: traceEvString for name.
-	// TODO: truncate the name if too long.
-	// TODO: traceEvSpan.
+func trace_userSpan(id, mode uint64, name string) {
+	if !trace.enabled {
+		return
+	}
+
+	mp, pid, bufp := traceAcquireBuffer()
+	if !trace.enabled && !mp.startingtrace {
+		traceReleaseBuffer(pid)
+		return
+	}
+
+	nameStringID, bufp := traceString(bufp, pid, name)
+	traceEventLocked(0, mp, pid, bufp, traceEvUserSpan, 3, id, mode, nameStringID)
+	traceReleaseBuffer(pid)
 }
 
 //go:linkname trace_userLog runtime/trace.userLog
 func trace_userLog(id uint64, category, message string) {
-	// TODO: traceEvString for key.
-	// TODO: traceEvUserLog.
+	if !trace.enabled {
+		return
+	}
+
+	mp, pid, bufp := traceAcquireBuffer()
+	if !trace.enabled && !mp.startingtrace {
+		traceReleaseBuffer(pid)
+		return
+	}
+
+	categoryID, bufp := traceString(bufp, pid, category)
+
+	extraSpace := traceBytesPerNumber + len(message) // extraSpace for the value string
+	traceEventLocked(extraSpace, mp, pid, bufp, traceEvUserLog, 3, id, categoryID)
+	// traceEventLocked reserved extra space for val and len(val)
+	// in buf, so buf now has room for the following.
+	buf := (*bufp).ptr()
+
+	// double-check the message and its length can fit.
+	// Otherwise, truncate the message.
+	slen := len(message)
+	if room := len(buf.arr) - buf.pos; room < slen+traceBytesPerNumber {
+		slen = room
+	}
+	buf.varint(uint64(slen))
+	buf.pos += copy(buf.arr[buf.pos:], message[:slen])
+
+	traceReleaseBuffer(pid)
 }

src/runtime/trace/annotation.go

@@ -72,9 +72,11 @@ type task struct {
 	// TODO(hyangah): record parent id?
 }
 
+var lastTaskID uint64 = 0 // task id issued last time
+
 func newID() uint64 {
-	// TODO(hyangah): implement
-	return 0
+	// TODO(hyangah): use per-P cache
+	return atomic.AddUint64(&lastTaskID, 1)
 }
 
 var bgTask = task{id: uint64(0)}
@@ -90,7 +92,10 @@ func Log(ctx context.Context, category, message string) {
 // Logf is like Log, but the value is formatted using the specified format spec.
 func Logf(ctx context.Context, category, format string, args ...interface{}) {
 	if IsEnabled() {
-		Log(ctx, category, fmt.Sprintf(format, args...))
+		// Ideally this should be just Log, but that will
+		// add one more frame in the stack trace.
+		id := fromContext(ctx).id
+		userLog(id, category, fmt.Sprintf(format, args...))
 	}
 }
 

src/runtime/trace/annotation_test.go

+package trace_test
+
+import (
+	"bytes"
+	"context"
+	"internal/trace"
+	"reflect"
+	. "runtime/trace"
+	"sync"
+	"testing"
+)
+
+func TestUserTaskSpan(t *testing.T) {
+	bgctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// TODO(hyangah): test pre-existing spans don't cause troubles
+
+	buf := new(bytes.Buffer)
+	if err := Start(buf); err != nil {
+		t.Fatalf("failed to start tracing: %v", err)
+	}
+
+	// Beginning of traced execution
+	var wg sync.WaitGroup
+	ctx, end := NewContext(bgctx, "task0") // EvUserTaskCreate("task0")
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		defer end() // EvUserTaskEnd("span0")
+
+		WithSpan(ctx, "span0", func(ctx context.Context) {
+			// EvUserSpanCreate("span0", start)
+			Log(ctx, "key0", "0123456789abcdef") // EvUserLog("task0", "key0", "0....f")
+			// EvUserSpan("span0", end)
+		})
+	}()
+	wg.Wait()
+	// End of traced execution
+	Stop()
+	saveTrace(t, buf, "TestUserTaskSpan")
+	res, err := trace.Parse(buf, "")
+	if err != nil {
+		t.Fatalf("Parse failed: %v", err)
+	}
+
+	// Check whether we see all user annotation related records in order
+	type testData struct {
+		typ  byte
+		strs []string
+		args []uint64
+	}
+
+	var got []testData
+	tasks := map[uint64]string{}
+	for _, e := range res.Events {
+		t.Logf("%s", e)
+		switch e.Type {
+		case trace.EvUserTaskCreate:
+			taskName := e.SArgs[0]
+			got = append(got, testData{trace.EvUserTaskCreate, []string{taskName}, nil})
+			tasks[e.Args[0]] = taskName
+		case trace.EvUserLog:
+			key, val := e.SArgs[0], e.SArgs[1]
+			taskName := tasks[e.Args[0]]
+			got = append(got, testData{trace.EvUserLog, []string{taskName, key, val}, nil})
+		case trace.EvUserTaskEnd:
+			taskName := tasks[e.Args[0]]
+			got = append(got, testData{trace.EvUserTaskEnd, []string{taskName}, nil})
+		case trace.EvUserSpan:
+			taskName := tasks[e.Args[0]]
+			spanName := e.SArgs[0]
+			got = append(got, testData{trace.EvUserSpan, []string{taskName, spanName}, []uint64{e.Args[1]}})
+		}
+	}
+	want := []testData{
+		{trace.EvUserTaskCreate, []string{"task0"}, nil},
+		{trace.EvUserSpan, []string{"task0", "span0"}, []uint64{0}},
+		{trace.EvUserLog, []string{"task0", "key0", "0123456789abcdef"}, nil},
+		{trace.EvUserSpan, []string{"task0", "span0"}, []uint64{1}},
+		{trace.EvUserTaskEnd, []string{"task0"}, nil},
+	}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("Got user span related events %+v\nwant: %+v", got, want)
+	}
+}