internal/trace: add end-of-generation signal to trace

This change takes the EvEndOfGeneration event and promotes it to a real
event that appears in the trace.

This allows the trace parser to unambiguously identify truncated traces
vs. broken traces. It also makes a lot of the logic around parsing
simpler, because there's no more batch spilling necessary.

Fixes golang#73904.

Change-Id: I37c359b32b6b5f894825aafc02921adeaacf2595
Reviewed-on: https://go-review.googlesource.com/c/go/+/693398
Reviewed-by: Carlos Amedee <[email protected]>
Reviewed-by: Michael Pratt <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: mknyszek

src/internal/trace/batch.go

@@ -44,6 +44,10 @@ func (b *batch) isSyncBatch(ver version.Version) bool {
 			(tracev2.EventType(b.data[0]) == tracev2.EvSync && ver >= version.Go125))
 }
 
+func (b *batch) isEndOfGeneration() bool {
+	return b.exp == tracev2.NoExperiment && len(b.data) > 0 && tracev2.EventType(b.data[0]) == tracev2.EvEndOfGeneration
+}
+
 // readBatch reads the next full batch from r.
 func readBatch(r interface {
 	io.Reader
@@ -54,6 +58,9 @@ func readBatch(r interface {
 	if err != nil {
 		return batch{}, 0, err
 	}
+	if typ := tracev2.EventType(b); typ == tracev2.EvEndOfGeneration {
+		return batch{m: NoThread, exp: tracev2.NoExperiment, data: []byte{b}}, 0, nil
+	}
 	if typ := tracev2.EventType(b); typ != tracev2.EvEventBatch && typ != tracev2.EvExperimentalBatch {
 		return batch{}, 0, fmt.Errorf("expected batch event, got event %d", typ)
 	}

src/internal/trace/generation.go

@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"cmp"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"slices"
@@ -32,22 +33,102 @@ type generation struct {
 	*evTable
 }
 
+// readGeneration buffers and decodes the structural elements of a trace generation
+// out of r.
+func readGeneration(r *bufio.Reader, ver version.Version) (*generation, error) {
+	if ver < version.Go126 {
+		return nil, errors.New("internal error: readGeneration called for <1.26 trace")
+	}
+	g := &generation{
+		evTable: &evTable{
+			pcs: make(map[uint64]frame),
+		},
+		batches: make(map[ThreadID][]batch),
+	}
+
+	// Read batches one at a time until we either hit the next generation.
+	for {
+		b, gen, err := readBatch(r)
+		if err == io.EOF {
+			if len(g.batches) != 0 {
+				return nil, errors.New("incomplete generation found; trace likely truncated")
+			}
+			return nil, nil // All done.
+		}
+		if err != nil {
+			return nil, err
+		}
+		if g.gen == 0 {
+			// Initialize gen.
+			g.gen = gen
+		}
+		if b.isEndOfGeneration() {
+			break
+		}
+		if gen == 0 {
+			// 0 is a sentinel used by the runtime, so we'll never see it.
+			return nil, fmt.Errorf("invalid generation number %d", gen)
+		}
+		if gen != g.gen {
+			return nil, fmt.Errorf("broken trace: missing end-of-generation event, or generations are interleaved")
+		}
+		if g.minTs == 0 || b.time < g.minTs {
+			g.minTs = b.time
+		}
+		if err := processBatch(g, b, ver); err != nil {
+			return nil, err
+		}
+	}
+
+	// Check some invariants.
+	if g.freq == 0 {
+		return nil, fmt.Errorf("no frequency event found")
+	}
+	if !g.hasClockSnapshot {
+		return nil, fmt.Errorf("no clock snapshot event found")
+	}
+
+	// N.B. Trust that the batch order is correct. We can't validate the batch order
+	// by timestamp because the timestamps could just be plain wrong. The source of
+	// truth is the order things appear in the trace and the partial order sequence
+	// numbers on certain events. If it turns out the batch order is actually incorrect
+	// we'll very likely fail to advance a partial order from the frontier.
+
+	// Compactify stacks and strings for better lookup performance later.
+	g.stacks.compactify()
+	g.strings.compactify()
+
+	// Validate stacks.
+	if err := validateStackStrings(&g.stacks, &g.strings, g.pcs); err != nil {
+		return nil, err
+	}
+
+	// Now that we have the frequency, fix up CPU samples.
+	fixUpCPUSamples(g.cpuSamples, g.freq)
+	return g, nil
+}
+
 // spilledBatch represents a batch that was read out for the next generation,
 // while reading the previous one. It's passed on when parsing the next
 // generation.
+//
+// Used only for trace versions < Go126.
 type spilledBatch struct {
 	gen uint64
 	*batch
 }
 
-// readGeneration buffers and decodes the structural elements of a trace generation
+// readGenerationWithSpill buffers and decodes the structural elements of a trace generation
 // out of r. spill is the first batch of the new generation (already buffered and
 // parsed from reading the last generation). Returns the generation and the first
 // batch read of the next generation, if any.
 //
 // If gen is non-nil, it is valid and must be processed before handling the returned
 // error.
-func readGeneration(r *bufio.Reader, spill *spilledBatch, ver version.Version) (*generation, *spilledBatch, error) {
+func readGenerationWithSpill(r *bufio.Reader, spill *spilledBatch, ver version.Version) (*generation, *spilledBatch, error) {
+	if ver >= version.Go126 {
+		return nil, nil, errors.New("internal error: readGenerationWithSpill called for Go 1.26+ trace")
+	}
 	g := &generation{
 		evTable: &evTable{
 			pcs: make(map[uint64]frame),
@@ -56,15 +137,15 @@ func readGeneration(r *bufio.Reader, spill *spilledBatch, ver version.Version) (
 	}
 	// Process the spilled batch.
 	if spill != nil {
+		// Process the spilled batch, which contains real data.
 		g.gen = spill.gen
 		g.minTs = spill.batch.time
 		if err := processBatch(g, *spill.batch, ver); err != nil {
 			return nil, nil, err
 		}
 		spill = nil
 	}
-	// Read batches one at a time until we either hit EOF or
-	// the next generation.
+	// Read batches one at a time until we either hit the next generation.
 	var spillErr error
 	for {
 		b, gen, err := readBatch(r)
@@ -73,7 +154,7 @@ func readGeneration(r *bufio.Reader, spill *spilledBatch, ver version.Version) (
 		}
 		if err != nil {
 			if g.gen != 0 {
-				// This is an error reading the first batch of the next generation.
+				// This may be an error reading the first batch of the next generation.
 				// This is fine. Let's forge ahead assuming that what we've got so
 				// far is fine.
 				spillErr = err
@@ -89,7 +170,8 @@ func readGeneration(r *bufio.Reader, spill *spilledBatch, ver version.Version) (
 			// Initialize gen.
 			g.gen = gen
 		}
-		if gen == g.gen+1 { // TODO: advance this the same way the runtime does.
+		if gen == g.gen+1 {
+			// TODO: Increment the generation with wraparound the same way the runtime does.
 			spill = &spilledBatch{gen: gen, batch: &b}
 			break
 		}
@@ -134,15 +216,8 @@ func readGeneration(r *bufio.Reader, spill *spilledBatch, ver version.Version) (
 		return nil, nil, err
 	}
 
-	// Fix up the CPU sample timestamps, now that we have freq.
-	for i := range g.cpuSamples {
-		s := &g.cpuSamples[i]
-		s.time = g.freq.mul(timestamp(s.time))
-	}
-	// Sort the CPU samples.
-	slices.SortFunc(g.cpuSamples, func(a, b cpuSample) int {
-		return cmp.Compare(a.time, b.time)
-	})
+	// Now that we have the frequency, fix up CPU samples.
+	fixUpCPUSamples(g.cpuSamples, g.freq)
 	return g, spill, spillErr
 }
 
@@ -174,6 +249,8 @@ func processBatch(g *generation, b batch, ver version.Version) error {
 		if err := addExperimentalBatch(g.expBatches, b); err != nil {
 			return err
 		}
+	case b.isEndOfGeneration():
+		return errors.New("internal error: unexpectedly processing EndOfGeneration; broken trace?")
 	default:
 		if _, ok := g.batches[b.m]; !ok {
 			g.batchMs = append(g.batchMs, b.m)
@@ -512,3 +589,15 @@ func addExperimentalBatch(expBatches map[tracev2.Experiment][]ExperimentalBatch,
 	})
 	return nil
 }
+
+func fixUpCPUSamples(samples []cpuSample, freq frequency) {
+	// Fix up the CPU sample timestamps.
+	for i := range samples {
+		s := &samples[i]
+		s.time = freq.mul(timestamp(s.time))
+	}
+	// Sort the CPU samples.
+	slices.SortFunc(samples, func(a, b cpuSample) int {
+		return cmp.Compare(a.time, b.time)
+	})
+}

src/internal/trace/internal/testgen/trace.go

@@ -322,6 +322,14 @@ func (g *Generation) writeEventsTo(tw *raw.TextWriter) {
 		}
 	}
 	b.writeEventsTo(tw)
+
+	// Write end-of-generation event if necessary.
+	if g.trace.ver >= version.Go126 {
+		tw.WriteEvent(raw.Event{
+			Version: g.trace.ver,
+			Ev:      tracev2.EvEndOfGeneration,
+		})
+	}
 }
 
 func (g *Generation) newStructuralBatch() *Batch {