go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/resultdb/internal/spanutil/compression.go (about)

     1  // Copyright 2020 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package spanutil
    16  
    17  import (
    18  	"bytes"
    19  
    20  	"github.com/klauspost/compress/zstd"
    21  
    22  	"go.chromium.org/luci/common/errors"
    23  )
    24  
    25  var zstdHeader = []byte("ztd\n")
    26  
    27  // Globally shared zstd encoder and decoder. We use only their EncodeAll and
    28  // DecodeAll methods which are allowed to be used concurrently. Internally, both
    29  // the encode and the decoder have worker pools (limited by GOMAXPROCS) and each
    30  // concurrent EncodeAll/DecodeAll call temporary consumes one worker (so overall
    31  // we do not run more jobs that we have cores for).
    32  var (
    33  	zstdEncoder *zstd.Encoder
    34  	zstdDecoder *zstd.Decoder
    35  )
    36  
    37  func init() {
    38  	var err error
    39  	if zstdEncoder, err = zstd.NewWriter(nil); err != nil {
    40  		panic(err) // this is impossible
    41  	}
    42  	if zstdDecoder, err = zstd.NewReader(nil); err != nil {
    43  		panic(err) // this is impossible
    44  	}
    45  }
    46  
    47  // Compressed instructs ToSpanner and FromSpanner functions to compress the
    48  // content with https://godoc.org/github.com/klauspost/compress/zstd encoding.
    49  type Compressed []byte
    50  
    51  // ToSpanner implements Value.
    52  func (c Compressed) ToSpanner() any {
    53  	if len(c) == 0 {
    54  		// Do not store empty bytes.
    55  		return []byte(nil)
    56  	}
    57  	return Compress(c)
    58  }
    59  
    60  // SpannerPtr implements Ptr.
    61  func (c *Compressed) SpannerPtr(b *Buffer) any {
    62  	return &b.ByteSlice
    63  }
    64  
    65  // FromSpanner implements Ptr.
    66  func (c *Compressed) FromSpanner(b *Buffer) error {
    67  	if len(b.ByteSlice) == 0 {
    68  		// do not set to nil; otherwise we lose the buffer.
    69  		*c = (*c)[:0]
    70  	} else {
    71  		// *c might be pointing to an existing memory buffer.
    72  		// Try to reuse it for decoding.
    73  		var err error
    74  		if *c, err = Decompress(b.ByteSlice, *c); err != nil {
    75  			return err
    76  		}
    77  	}
    78  
    79  	return nil
    80  }
    81  
    82  // Compress compresses data using zstd.
    83  func Compress(data []byte) []byte {
    84  	out := make([]byte, 0, len(data)/2+len(zstdHeader)) // hope for at least 2x compression
    85  	out = append(out, zstdHeader...)
    86  	return zstdEncoder.EncodeAll(data, out)
    87  }
    88  
    89  // Decompress decompresses the src compressed with Compress to dest.
    90  // dest is the buffer for decompressed content, it will be reset to 0 length
    91  // before taking the content.
    92  func Decompress(src, dest []byte) ([]byte, error) {
    93  	if !bytes.HasPrefix(src, zstdHeader) {
    94  		return nil, errors.Reason("expected ztd header").Err()
    95  	}
    96  
    97  	dest, err := zstdDecoder.DecodeAll(src[len(zstdHeader):], dest[:0])
    98  	if err != nil {
    99  		return nil, errors.Annotate(err, "failed to decode from zstd").Err()
   100  	}
   101  	return dest, nil
   102  }