github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/cmn/tests/bytepack_test.go (about)

     1  // Package test provides tests for common low-level types and utilities for all aistore projects
     2  /*
     3   * Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package tests_test
     6  
     7  import (
     8  	"bytes"
     9  	"testing"
    10  
    11  	"github.com/NVIDIA/aistore/cmn/cos"
    12  	"github.com/NVIDIA/aistore/tools/tassert"
    13  	"github.com/NVIDIA/aistore/tools/trand"
    14  )
    15  
    16  // A structure to test nested binary packing
    17  type pck struct {
    18  	id     int64
    19  	group  int16
    20  	name   string
    21  	data   []byte
    22  	parent *pck
    23  }
    24  
    25  func (p *pck) Pack(wr *cos.BytePack) {
    26  	// write POD and variable-length fields in turns
    27  	wr.WriteString(p.name)
    28  	wr.WriteInt64(p.id)
    29  	wr.WriteBytes(p.data)
    30  	wr.WriteInt16(p.group)
    31  	// A marker trick: before saving inner structure, put a boolean
    32  	// marker that indicates if unpacker should skip reading inner struct
    33  	if p.parent == nil {
    34  		wr.WriteByte(0)
    35  	} else {
    36  		wr.WriteByte(1)
    37  		wr.WriteAny(p.parent)
    38  	}
    39  }
    40  
    41  func (p *pck) Unpack(rd *cos.ByteUnpack) (err error) {
    42  	if p.name, err = rd.ReadString(); err != nil {
    43  		return
    44  	}
    45  	if p.id, err = rd.ReadInt64(); err != nil {
    46  		return
    47  	}
    48  	if p.data, err = rd.ReadBytes(); err != nil {
    49  		return
    50  	}
    51  	if p.group, err = rd.ReadInt16(); err != nil {
    52  		return
    53  	}
    54  
    55  	var exists byte
    56  	if exists, err = rd.ReadByte(); err != nil {
    57  		return
    58  	}
    59  	// Read inner struct only of the marker is `true`.
    60  	if exists != 0 {
    61  		// Do not forget to initialize inner field otherwise it may panic
    62  		// if it is `nil` at this point.
    63  		p.parent = &pck{}
    64  		rd.ReadAny(p.parent)
    65  	}
    66  	return
    67  }
    68  
    69  func (p *pck) PackedSize() int {
    70  	//    id              name&data len
    71  	sz := cos.SizeofI64 + cos.SizeofLen*2 +
    72  		// group        name len      data len      inner pointer marker
    73  		cos.SizeofI16 + len(p.name) + len(p.data) + 1
    74  	if p.parent != nil {
    75  		// If inner struct is not `nil`, add its size to the total.
    76  		sz += p.parent.PackedSize()
    77  	}
    78  	return sz
    79  }
    80  
    81  func TestBytePackStruct(t *testing.T) {
    82  	first := &pck{
    83  		id:     0x01020304,
    84  		group:  0x0507,
    85  		name:   "first",
    86  		data:   nil,
    87  		parent: nil,
    88  	}
    89  	second := &pck{
    90  		id:    0x11121314,
    91  		group: 0x1517,
    92  		name:  "second item",
    93  		data:  []byte("abcde"),
    94  		parent: &pck{
    95  			id:     0x21222324,
    96  			group:  0x2527,
    97  			name:   "inner item",
    98  			data:   []byte("hijkl"),
    99  			parent: nil,
   100  		},
   101  	}
   102  
   103  	packer := cos.NewPacker(nil, first.PackedSize()+second.PackedSize())
   104  	packer.WriteAny(first)
   105  	packer.WriteAny(second)
   106  
   107  	readFirst := &pck{}
   108  	readSecond := &pck{}
   109  	unpacker := cos.NewUnpacker(packer.Bytes())
   110  	err := unpacker.ReadAny(readFirst)
   111  	tassert.CheckFatal(t, err)
   112  	err = unpacker.ReadAny(readSecond)
   113  	tassert.CheckFatal(t, err)
   114  	if first.id != readFirst.id ||
   115  		first.group != readFirst.group ||
   116  		first.name != readFirst.name ||
   117  		len(readFirst.data) != 0 ||
   118  		readFirst.parent != nil {
   119  		t.Errorf("First: Read %+v mismatches original %+v", readFirst, first)
   120  	}
   121  	if second.id != readSecond.id ||
   122  		second.group != readSecond.group ||
   123  		second.name != readSecond.name ||
   124  		!bytes.Equal(second.data, readSecond.data) ||
   125  		readSecond.parent == nil {
   126  		t.Errorf("Second: Read %+v mismatches original %+v", readSecond, second)
   127  	}
   128  	if second.parent.id != readSecond.parent.id ||
   129  		second.parent.group != readSecond.parent.group ||
   130  		second.parent.name != readSecond.parent.name ||
   131  		!bytes.Equal(second.parent.data, readSecond.parent.data) ||
   132  		readSecond.parent.parent != nil {
   133  		t.Errorf("Second inner: Read %+v mismatches original %+v", readSecond.parent, second.parent)
   134  	}
   135  }
   136  
   137  func BenchmarkPackWriteString(b *testing.B) {
   138  	p := cos.NewPacker(nil, 90*b.N)
   139  
   140  	a := make([]string, 0, 1000)
   141  	for range 1000 {
   142  		a = append(a, trand.String(80))
   143  	}
   144  
   145  	b.ReportAllocs()
   146  	b.ResetTimer()
   147  
   148  	for i := range b.N {
   149  		p.WriteString(a[i%len(a)])
   150  	}
   151  }