github.com/grafana/pyroscope@v1.18.0/pkg/metastore/index/store/shard_strings.go (about)

     1  package store
     2  
     3  import (
     4  	"encoding/binary"
     5  
     6  	"github.com/grafana/pyroscope/pkg/iter"
     7  	"github.com/grafana/pyroscope/pkg/metastore/store"
     8  )
     9  
    10  type stringIterator struct {
    11  	iter.Iterator[store.KV]
    12  	batch []string
    13  	cur   int
    14  	err   error
    15  }
    16  
    17  func newStringIter(i iter.Iterator[store.KV]) *stringIterator {
    18  	return &stringIterator{Iterator: i}
    19  }
    20  
    21  func (i *stringIterator) Err() error {
    22  	if err := i.Iterator.Err(); err != nil {
    23  		return err
    24  	}
    25  	return i.err
    26  }
    27  
    28  func (i *stringIterator) At() string { return i.batch[i.cur] }
    29  
    30  func (i *stringIterator) Next() bool {
    31  	if n := i.cur + 1; n < len(i.batch) {
    32  		i.cur = n
    33  		return true
    34  	}
    35  	i.cur = 0
    36  	i.batch = i.batch[:0]
    37  	if !i.Iterator.Next() {
    38  		return false
    39  	}
    40  	var err error
    41  	if i.batch, err = decodeStrings(i.batch, i.Iterator.At().Value); err != nil {
    42  		i.err = err
    43  		return false
    44  	}
    45  	return len(i.batch) > 0
    46  }
    47  
    48  func encodeStrings(strings []string) []byte {
    49  	size := 4
    50  	for _, s := range strings {
    51  		size += 4 + len(s)
    52  	}
    53  	data := make([]byte, 0, size)
    54  	data = binary.BigEndian.AppendUint32(data, uint32(len(strings)))
    55  	for _, s := range strings {
    56  		data = binary.BigEndian.AppendUint32(data, uint32(len(s)))
    57  		data = append(data, s...)
    58  	}
    59  	return data
    60  }
    61  
    62  func decodeStrings(dst []string, data []byte) ([]string, error) {
    63  	offset := 0
    64  	if len(data) < offset+4 {
    65  		return dst, ErrInvalidStringTable
    66  	}
    67  	n := binary.BigEndian.Uint32(data[offset:])
    68  	offset += 4
    69  	for i := uint32(0); i < n; i++ {
    70  		if len(data) < offset+4 {
    71  			return dst, ErrInvalidStringTable
    72  		}
    73  		size := binary.BigEndian.Uint32(data[offset:])
    74  		offset += 4
    75  		if len(data) < offset+int(size) {
    76  			return dst, ErrInvalidStringTable
    77  		}
    78  		dst = append(dst, string(data[offset:offset+int(size)]))
    79  		offset += int(size)
    80  	}
    81  	return dst, nil
    82  }