github.com/grafana/synthetic-monitoring-agent@v0.64.2/internal/pusher/v2/snappy_concat.go (about)

     1  package v2
     2  
     3  import (
     4  	encoding_binary "encoding/binary"
     5  	"io"
     6  )
     7  
     8  // SnappyConcatReader is an io.Reader that takes a list of snappy-compressed buffers (Streams)
     9  // and returns a single valid snappy-compressed stream equivalent to the concatenation of those buffers
    10  // without additional allocation.
    11  type SnappyConcatReader struct {
    12  	Streams [][]byte
    13  	init    bool
    14  }
    15  
    16  // Read fulfills the io.Reader interface.
    17  func (s *SnappyConcatReader) Read(out []byte) (n int, err error) {
    18  	if !s.init {
    19  		s.init = true
    20  		decodedLen := s.stripHeaders()
    21  		// Check that we have room to write the uvarint
    22  		if fit := len(out); fit < encoding_binary.MaxVarintLen64 {
    23  			if maxValue := uint64(1) << (7 * fit); decodedLen >= maxValue {
    24  				return 0, io.ErrShortBuffer
    25  			}
    26  		}
    27  
    28  		n = encoding_binary.PutUvarint(out, decodedLen)
    29  		out = out[n:]
    30  	}
    31  
    32  	if n == 0 && len(s.Streams) == 0 {
    33  		return 0, io.EOF
    34  	}
    35  
    36  	// Copy all the streams that fit
    37  	for len(s.Streams) > 0 && len(out) >= len(s.Streams[0]) {
    38  		chunkLen := len(s.Streams[0])
    39  		copy(out, s.Streams[0])
    40  		s.Streams = s.Streams[1:]
    41  		n += chunkLen
    42  		out = out[chunkLen:]
    43  	}
    44  
    45  	// Copy a partial stream if there is room
    46  	if fit := len(out); fit > 0 && len(s.Streams) > 0 {
    47  		// Here s.Streams[0] is always larger than what's left in out.
    48  		// Otherwise it would've been copied by the previous loop.
    49  		copy(out, s.Streams[0][:fit])
    50  		s.Streams[0] = s.Streams[0][fit:]
    51  		n += fit
    52  	}
    53  
    54  	return n, nil
    55  }
    56  
    57  func (s *SnappyConcatReader) stripHeaders() uint64 {
    58  	var totalLength uint64
    59  
    60  	for idx := range s.Streams {
    61  		streamLen, skip := encoding_binary.Uvarint(s.Streams[idx])
    62  		totalLength += streamLen
    63  		s.Streams[idx] = s.Streams[idx][skip:]
    64  	}
    65  
    66  	return totalLength
    67  }