src/internal/trace/v2/batch.go

   1 // Copyright 2023 The Go Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 package trace
   6
   7 import (
   8         "bufio"
   9         "bytes"
  10         "encoding/binary"
  11         "fmt"
  12         "io"
  13
  14         "internal/trace/v2/event"
  15         "internal/trace/v2/event/go122"
  16 )
  17
  18 // timestamp is an unprocessed timestamp.
  19 type timestamp uint64
  20
  21 // batch represents a batch of trace events.
  22 // It is unparsed except for its header.
  23 type batch struct {
  24         m    ThreadID
  25         time timestamp
  26         data []byte
  27 }
  28
  29 func (b *batch) isStringsBatch() bool {
  30         return len(b.data) > 0 && event.Type(b.data[0]) == go122.EvStrings
  31 }
  32
  33 func (b *batch) isStacksBatch() bool {
  34         return len(b.data) > 0 && event.Type(b.data[0]) == go122.EvStacks
  35 }
  36
  37 func (b *batch) isCPUSamplesBatch() bool {
  38         return len(b.data) > 0 && event.Type(b.data[0]) == go122.EvCPUSamples
  39 }
  40
  41 func (b *batch) isFreqBatch() bool {
  42         return len(b.data) > 0 && event.Type(b.data[0]) == go122.EvFrequency
  43 }
  44
  45 // readBatch reads the next full batch from r.
  46 func readBatch(r *bufio.Reader) (batch, uint64, error) {
  47         // Read batch header byte.
  48         b, err := r.ReadByte()
  49         if err != nil {
  50                 return batch{}, 0, err
  51         }
  52         if typ := event.Type(b); typ != go122.EvEventBatch {
  53                 return batch{}, 0, fmt.Errorf("expected batch event (%s), got %s", go122.EventString(go122.EvEventBatch), go122.EventString(typ))
  54         }
  55
  56         // Read the batch header: gen (generation), thread (M) ID, base timestamp
  57         // for the batch.
  58         gen, err := binary.ReadUvarint(r)
  59         if err != nil {
  60                 return batch{}, gen, fmt.Errorf("error reading batch gen: %w", err)
  61         }
  62         m, err := binary.ReadUvarint(r)
  63         if err != nil {
  64                 return batch{}, gen, fmt.Errorf("error reading batch M ID: %w", err)
  65         }
  66         ts, err := binary.ReadUvarint(r)
  67         if err != nil {
  68                 return batch{}, gen, fmt.Errorf("error reading batch timestamp: %w", err)
  69         }
  70
  71         // Read in the size of the batch to follow.
  72         size, err := binary.ReadUvarint(r)
  73         if err != nil {
  74                 return batch{}, gen, fmt.Errorf("error reading batch size: %w", err)
  75         }
  76         if size > go122.MaxBatchSize {
  77                 return batch{}, gen, fmt.Errorf("invalid batch size %d, maximum is %d", size, go122.MaxBatchSize)
  78         }
  79
  80         // Copy out the batch for later processing.
  81         var data bytes.Buffer
  82         data.Grow(int(size))
  83         n, err := io.CopyN(&data, r, int64(size))
  84         if n != int64(size) {
  85                 return batch{}, gen, fmt.Errorf("failed to read full batch: read %d but wanted %d", n, size)
  86         }
  87         if err != nil {
  88                 return batch{}, gen, fmt.Errorf("copying batch data: %w", err)
  89         }
  90
  91         // Return the batch.
  92         return batch{
  93                 m:    ThreadID(m),
  94                 time: timestamp(ts),
  95                 data: data.Bytes(),
  96         }, gen, nil
  97 }