serialize_test.go

     1  // Copyright 2024 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 pgo
     6  
     7  import (
     8  	"bytes"
     9  	"encoding/binary"
    10  	"fmt"
    11  	"reflect"
    12  	"strings"
    13  	"testing"
    14  )
    15  
    16  // equal returns an error if got and want are not equal.
    17  func equal(got, want *Profile) error {
    18  	if got.TotalWeight != want.TotalWeight {
    19  		return fmt.Errorf("got.TotalWeight %d != want.TotalWeight %d", got.TotalWeight, want.TotalWeight)
    20  	}
    21  	if !reflect.DeepEqual(got.NamedEdgeMap.ByWeight, want.NamedEdgeMap.ByWeight) {
    22  		return fmt.Errorf("got.NamedEdgeMap.ByWeight != want.NamedEdgeMap.ByWeight\ngot = %+v\nwant = %+v", got.NamedEdgeMap.ByWeight, want.NamedEdgeMap.ByWeight)
    23  	}
    24  	if !reflect.DeepEqual(got.NamedEdgeMap.Weight, want.NamedEdgeMap.Weight) {
    25  		return fmt.Errorf("got.NamedEdgeMap.Weight != want.NamedEdgeMap.Weight\ngot = %+v\nwant = %+v", got.NamedEdgeMap.Weight, want.NamedEdgeMap.Weight)
    26  	}
    27  
    28  	return nil
    29  }
    30  
    31  func testRoundTrip(t *testing.T, d *Profile) []byte {
    32  	var buf bytes.Buffer
    33  	n, err := d.WriteTo(&buf)
    34  	if err != nil {
    35  		t.Fatalf("WriteTo got err %v want nil", err)
    36  	}
    37  	if n != int64(buf.Len()) {
    38  		t.Errorf("WriteTo got n %d want %d", n, int64(buf.Len()))
    39  	}
    40  
    41  	b := buf.Bytes()
    42  
    43  	got, err := FromSerialized(&buf)
    44  	if err != nil {
    45  		t.Fatalf("processSerialized got err %v want nil", err)
    46  	}
    47  	if err := equal(got, d); err != nil {
    48  		t.Errorf("processSerialized output does not match input: %v", err)
    49  	}
    50  
    51  	return b
    52  }
    53  
    54  func TestEmpty(t *testing.T) {
    55  	d := emptyProfile()
    56  	b := testRoundTrip(t, d)
    57  
    58  	// Contents should consist of only a header.
    59  	if string(b) != serializationHeader {
    60  		t.Errorf("WriteTo got %q want %q", string(b), serializationHeader)
    61  	}
    62  }
    63  
    64  func TestRoundTrip(t *testing.T) {
    65  	d := &Profile{
    66  		TotalWeight: 3,
    67  		NamedEdgeMap: NamedEdgeMap{
    68  			ByWeight: []NamedCallEdge{
    69  				{
    70  					CallerName: "a",
    71  					CalleeName: "b",
    72  					CallSiteOffset: 14,
    73  				},
    74  				{
    75  					CallerName: "c",
    76  					CalleeName: "d",
    77  					CallSiteOffset: 15,
    78  				},
    79  			},
    80  			Weight: map[NamedCallEdge]int64{
    81  				{
    82  					CallerName: "a",
    83  					CalleeName: "b",
    84  					CallSiteOffset: 14,
    85  				}: 2,
    86  				{
    87  					CallerName: "c",
    88  					CalleeName: "d",
    89  					CallSiteOffset: 15,
    90  				}: 1,
    91  			},
    92  		},
    93  	}
    94  
    95  	testRoundTrip(t, d)
    96  }
    97  
    98  func constructFuzzProfile(t *testing.T, b []byte) *Profile {
    99  	// The fuzzer can't construct an arbitrary structure, so instead we
   100  	// consume bytes from b to act as our edge data.
   101  	r := bytes.NewReader(b)
   102  	consumeString := func() (string, bool) {
   103  		// First byte: how many bytes to read for this string? We only
   104  		// use a byte to avoid making humongous strings.
   105  		length, err := r.ReadByte()
   106  		if err != nil {
   107  			return "", false
   108  		}
   109  		if length == 0 {
   110  			return "", false
   111  		}
   112  
   113  		b := make([]byte, length)
   114  		_, err = r.Read(b)
   115  		if err != nil {
   116  			return "", false
   117  		}
   118  
   119  		return string(b), true
   120  	}
   121  	consumeInt64 := func() (int64, bool) {
   122  		b := make([]byte, 8)
   123  		_, err := r.Read(b)
   124  		if err != nil {
   125  			return 0, false
   126  		}
   127  
   128  		return int64(binary.LittleEndian.Uint64(b)), true
   129  	}
   130  
   131  	d := emptyProfile()
   132  
   133  	for {
   134  		caller, ok := consumeString()
   135  		if !ok {
   136  			break
   137  		}
   138  		if strings.ContainsAny(caller, " \r\n") {
   139  			t.Skip("caller contains space or newline")
   140  		}
   141  
   142  		callee, ok := consumeString()
   143  		if !ok {
   144  			break
   145  		}
   146  		if strings.ContainsAny(callee, " \r\n") {
   147  			t.Skip("callee contains space or newline")
   148  		}
   149  
   150  		line, ok := consumeInt64()
   151  		if !ok {
   152  			break
   153  		}
   154  		weight, ok := consumeInt64()
   155  		if !ok {
   156  			break
   157  		}
   158  
   159  		edge := NamedCallEdge{
   160  			CallerName: caller,
   161  			CalleeName: callee,
   162  			CallSiteOffset: int(line),
   163  		}
   164  
   165  		if _, ok := d.NamedEdgeMap.Weight[edge]; ok {
   166  			t.Skip("duplicate edge")
   167  		}
   168  
   169  		d.NamedEdgeMap.Weight[edge] = weight
   170  		d.TotalWeight += weight
   171  	}
   172  
   173  	byWeight := make([]NamedCallEdge, 0, len(d.NamedEdgeMap.Weight))
   174  	for namedEdge := range d.NamedEdgeMap.Weight {
   175  		byWeight = append(byWeight, namedEdge)
   176  	}
   177  	sortByWeight(byWeight, d.NamedEdgeMap.Weight)
   178  	d.NamedEdgeMap.ByWeight = byWeight
   179  
   180  	return d
   181  }
   182  
   183  func FuzzRoundTrip(f *testing.F) {
   184  	f.Add([]byte("")) // empty profile
   185  
   186  	f.Fuzz(func(t *testing.T, b []byte) {
   187  		d := constructFuzzProfile(t, b)
   188  		testRoundTrip(t, d)
   189  	})
   190  }
   191
View as plain text