encode.go

     1  // Copyright 2021 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 encodecounter
     6  
     7  import (
     8  	"bufio"
     9  	"encoding/binary"
    10  	"fmt"
    11  	"internal/coverage"
    12  	"internal/coverage/slicewriter"
    13  	"internal/coverage/stringtab"
    14  	"internal/coverage/uleb128"
    15  	"io"
    16  	"maps"
    17  	"os"
    18  	"slices"
    19  )
    20  
    21  // This package contains APIs and helpers for encoding initial portions
    22  // of the counter data files emitted at runtime when coverage instrumentation
    23  // is enabled.  Counter data files may contain multiple segments; the file
    24  // header and first segment are written via the "Write" method below, and
    25  // additional segments can then be added using "AddSegment".
    26  
    27  type CoverageDataWriter struct {
    28  	stab    *stringtab.Writer
    29  	w       *bufio.Writer
    30  	csh     coverage.CounterSegmentHeader
    31  	tmp     []byte
    32  	cflavor coverage.CounterFlavor
    33  	segs    uint32
    34  	debug   bool
    35  }
    36  
    37  func NewCoverageDataWriter(w io.Writer, flav coverage.CounterFlavor) *CoverageDataWriter {
    38  	r := &CoverageDataWriter{
    39  		stab: &stringtab.Writer{},
    40  		w:    bufio.NewWriter(w),
    41  
    42  		tmp:     make([]byte, 64),
    43  		cflavor: flav,
    44  	}
    45  	r.stab.InitWriter()
    46  	r.stab.Lookup("")
    47  	return r
    48  }
    49  
    50  // CounterVisitor describes a helper object used during counter file
    51  // writing; when writing counter data files, clients pass a
    52  // CounterVisitor to the write/emit routines, then the expectation is
    53  // that the VisitFuncs method will then invoke the callback "f" with
    54  // data for each function to emit to the file.
    55  type CounterVisitor interface {
    56  	VisitFuncs(f CounterVisitorFn) error
    57  }
    58  
    59  // CounterVisitorFn describes a callback function invoked when writing
    60  // coverage counter data.
    61  type CounterVisitorFn func(pkid uint32, funcid uint32, counters []uint32) error
    62  
    63  // Write writes the contents of the count-data file to the writer
    64  // previously supplied to NewCoverageDataWriter. Returns an error
    65  // if something went wrong somewhere with the write.
    66  func (cfw *CoverageDataWriter) Write(metaFileHash [16]byte, args map[string]string, visitor CounterVisitor) error {
    67  	if err := cfw.writeHeader(metaFileHash); err != nil {
    68  		return err
    69  	}
    70  	return cfw.AppendSegment(args, visitor)
    71  }
    72  
    73  func padToFourByteBoundary(ws *slicewriter.WriteSeeker) error {
    74  	sz := len(ws.BytesWritten())
    75  	zeros := []byte{0, 0, 0, 0}
    76  	rem := uint32(sz) % 4
    77  	if rem != 0 {
    78  		pad := zeros[:(4 - rem)]
    79  		if nw, err := ws.Write(pad); err != nil {
    80  			return err
    81  		} else if nw != len(pad) {
    82  			return fmt.Errorf("error: short write")
    83  		}
    84  	}
    85  	return nil
    86  }
    87  
    88  func (cfw *CoverageDataWriter) patchSegmentHeader(ws *slicewriter.WriteSeeker) error {
    89  	// record position
    90  	off, err := ws.Seek(0, io.SeekCurrent)
    91  	if err != nil {
    92  		return fmt.Errorf("error seeking in patchSegmentHeader: %v", err)
    93  	}
    94  	// seek back to start so that we can update the segment header
    95  	if _, err := ws.Seek(0, io.SeekStart); err != nil {
    96  		return fmt.Errorf("error seeking in patchSegmentHeader: %v", err)
    97  	}
    98  	if cfw.debug {
    99  		fmt.Fprintf(os.Stderr, "=-= writing counter segment header: %+v", cfw.csh)
   100  	}
   101  	if err := binary.Write(ws, binary.LittleEndian, cfw.csh); err != nil {
   102  		return err
   103  	}
   104  	// ... and finally return to the original offset.
   105  	if _, err := ws.Seek(off, io.SeekStart); err != nil {
   106  		return fmt.Errorf("error seeking in patchSegmentHeader: %v", err)
   107  	}
   108  	return nil
   109  }
   110  
   111  func (cfw *CoverageDataWriter) writeSegmentPreamble(args map[string]string, ws *slicewriter.WriteSeeker) error {
   112  	if err := binary.Write(ws, binary.LittleEndian, cfw.csh); err != nil {
   113  		return err
   114  	}
   115  	hdrsz := uint32(len(ws.BytesWritten()))
   116  
   117  	// Write string table and args to a byte slice (since we need
   118  	// to capture offsets at various points), then emit the slice
   119  	// once we are done.
   120  	cfw.stab.Freeze()
   121  	if err := cfw.stab.Write(ws); err != nil {
   122  		return err
   123  	}
   124  	cfw.csh.StrTabLen = uint32(len(ws.BytesWritten())) - hdrsz
   125  
   126  	akeys := slices.Sorted(maps.Keys(args))
   127  
   128  	wrULEB128 := func(v uint) error {
   129  		cfw.tmp = cfw.tmp[:0]
   130  		cfw.tmp = uleb128.AppendUleb128(cfw.tmp, v)
   131  		if _, err := ws.Write(cfw.tmp); err != nil {
   132  			return err
   133  		}
   134  		return nil
   135  	}
   136  
   137  	// Count of arg pairs.
   138  	if err := wrULEB128(uint(len(args))); err != nil {
   139  		return err
   140  	}
   141  	// Arg pairs themselves.
   142  	for _, k := range akeys {
   143  		ki := uint(cfw.stab.Lookup(k))
   144  		if err := wrULEB128(ki); err != nil {
   145  			return err
   146  		}
   147  		v := args[k]
   148  		vi := uint(cfw.stab.Lookup(v))
   149  		if err := wrULEB128(vi); err != nil {
   150  			return err
   151  		}
   152  	}
   153  	if err := padToFourByteBoundary(ws); err != nil {
   154  		return err
   155  	}
   156  	cfw.csh.ArgsLen = uint32(len(ws.BytesWritten())) - (cfw.csh.StrTabLen + hdrsz)
   157  
   158  	return nil
   159  }
   160  
   161  // AppendSegment appends a new segment to a counter data, with a new
   162  // args section followed by a payload of counter data clauses.
   163  func (cfw *CoverageDataWriter) AppendSegment(args map[string]string, visitor CounterVisitor) error {
   164  	cfw.stab = &stringtab.Writer{}
   165  	cfw.stab.InitWriter()
   166  	cfw.stab.Lookup("")
   167  
   168  	var err error
   169  	for k, v := range args {
   170  		cfw.stab.Lookup(k)
   171  		cfw.stab.Lookup(v)
   172  	}
   173  
   174  	ws := &slicewriter.WriteSeeker{}
   175  	if err = cfw.writeSegmentPreamble(args, ws); err != nil {
   176  		return err
   177  	}
   178  	if err = cfw.writeCounters(visitor, ws); err != nil {
   179  		return err
   180  	}
   181  	if err = cfw.patchSegmentHeader(ws); err != nil {
   182  		return err
   183  	}
   184  	if err := cfw.writeBytes(ws.BytesWritten()); err != nil {
   185  		return err
   186  	}
   187  	if err = cfw.writeFooter(); err != nil {
   188  		return err
   189  	}
   190  	if err := cfw.w.Flush(); err != nil {
   191  		return fmt.Errorf("write error: %v", err)
   192  	}
   193  	cfw.stab = nil
   194  	return nil
   195  }
   196  
   197  func (cfw *CoverageDataWriter) writeHeader(metaFileHash [16]byte) error {
   198  	// Emit file header.
   199  	ch := coverage.CounterFileHeader{
   200  		Magic:     coverage.CovCounterMagic,
   201  		Version:   coverage.CounterFileVersion,
   202  		MetaHash:  metaFileHash,
   203  		CFlavor:   cfw.cflavor,
   204  		BigEndian: false,
   205  	}
   206  	if err := binary.Write(cfw.w, binary.LittleEndian, ch); err != nil {
   207  		return err
   208  	}
   209  	return nil
   210  }
   211  
   212  func (cfw *CoverageDataWriter) writeBytes(b []byte) error {
   213  	if len(b) == 0 {
   214  		return nil
   215  	}
   216  	nw, err := cfw.w.Write(b)
   217  	if err != nil {
   218  		return fmt.Errorf("error writing counter data: %v", err)
   219  	}
   220  	if len(b) != nw {
   221  		return fmt.Errorf("error writing counter data: short write")
   222  	}
   223  	return nil
   224  }
   225  
   226  func (cfw *CoverageDataWriter) writeCounters(visitor CounterVisitor, ws *slicewriter.WriteSeeker) error {
   227  	// Notes:
   228  	// - this version writes everything little-endian, which means
   229  	//   a call is needed to encode every value (expensive)
   230  	// - we may want to move to a model in which we just blast out
   231  	//   all counters, or possibly mmap the file and do the write
   232  	//   implicitly.
   233  	ctrb := make([]byte, 4)
   234  	wrval := func(val uint32) error {
   235  		var buf []byte
   236  		var towr int
   237  		if cfw.cflavor == coverage.CtrRaw {
   238  			binary.LittleEndian.PutUint32(ctrb, val)
   239  			buf = ctrb
   240  			towr = 4
   241  		} else if cfw.cflavor == coverage.CtrULeb128 {
   242  			cfw.tmp = cfw.tmp[:0]
   243  			cfw.tmp = uleb128.AppendUleb128(cfw.tmp, uint(val))
   244  			buf = cfw.tmp
   245  			towr = len(buf)
   246  		} else {
   247  			panic("internal error: bad counter flavor")
   248  		}
   249  		if sz, err := ws.Write(buf); err != nil {
   250  			return err
   251  		} else if sz != towr {
   252  			return fmt.Errorf("writing counters: short write")
   253  		}
   254  		return nil
   255  	}
   256  
   257  	// Write out entries for each live function.
   258  	emitter := func(pkid uint32, funcid uint32, counters []uint32) error {
   259  		cfw.csh.FcnEntries++
   260  		if err := wrval(uint32(len(counters))); err != nil {
   261  			return err
   262  		}
   263  
   264  		if err := wrval(pkid); err != nil {
   265  			return err
   266  		}
   267  
   268  		if err := wrval(funcid); err != nil {
   269  			return err
   270  		}
   271  		for _, val := range counters {
   272  			if err := wrval(val); err != nil {
   273  				return err
   274  			}
   275  		}
   276  		return nil
   277  	}
   278  	if err := visitor.VisitFuncs(emitter); err != nil {
   279  		return err
   280  	}
   281  	return nil
   282  }
   283  
   284  func (cfw *CoverageDataWriter) writeFooter() error {
   285  	cfw.segs++
   286  	cf := coverage.CounterFileFooter{
   287  		Magic:       coverage.CovCounterMagic,
   288  		NumSegments: cfw.segs,
   289  	}
   290  	if err := binary.Write(cfw.w, binary.LittleEndian, cf); err != nil {
   291  		return err
   292  	}
   293  	return nil
   294  }
   295
View as plain text