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

     1  // Copyright 2022 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 compression provides different compression and decompression
    16  // functions for the usages in the buildbucket.
    17  package compression
    18  
    19  import (
    20  	"bytes"
    21  	"compress/zlib"
    22  	"io"
    23  
    24  	"github.com/klauspost/compress/zstd"
    25  
    26  	"go.chromium.org/luci/common/errors"
    27  )
    28  
    29  // Globally shared zstd encoder and decoder. We use only their EncodeAll and
    30  // DecodeAll methods which are allowed to be used concurrently. Internally, both
    31  // the encode and the decoder have worker pools (limited by GOMAXPROCS) and each
    32  // concurrent EncodeAll/DecodeAll call temporary consumes one worker (so overall
    33  // we do not run more jobs that we have cores for).
    34  var (
    35  	zstdEncoder *zstd.Encoder
    36  	zstdDecoder *zstd.Decoder
    37  )
    38  
    39  func init() {
    40  	var err error
    41  	if zstdEncoder, err = zstd.NewWriter(nil); err != nil {
    42  		panic(err) // this is impossible
    43  	}
    44  	if zstdDecoder, err = zstd.NewReader(nil); err != nil {
    45  		panic(err) // this is impossible
    46  	}
    47  }
    48  
    49  // ZstdCompress compresses src and append it to dst.
    50  func ZstdCompress(src, dst []byte) []byte {
    51  	return zstdEncoder.EncodeAll(src, dst)
    52  }
    53  
    54  // ZstdDecompress decompresses input and append it to dst.
    55  func ZstdDecompress(input, dst []byte) ([]byte, error) {
    56  	return zstdDecoder.DecodeAll(input, dst)
    57  }
    58  
    59  // ZlibCompress compresses data using zlib.
    60  func ZlibCompress(data []byte) ([]byte, error) {
    61  	buf := &bytes.Buffer{}
    62  	zw := zlib.NewWriter(buf)
    63  	if _, err := zw.Write(data); err != nil {
    64  		return nil, errors.Annotate(err, "failed to compress").Err()
    65  	}
    66  	if err := zw.Close(); err != nil {
    67  		return nil, errors.Annotate(err, "error closing zlib writer").Err()
    68  	}
    69  	return buf.Bytes(), nil
    70  }
    71  
    72  // ZlibDecompress decompresses data using zlib.
    73  func ZlibDecompress(compressed []byte) ([]byte, error) {
    74  	r, err := zlib.NewReader(bytes.NewReader(compressed))
    75  	if err != nil {
    76  		return nil, err
    77  	}
    78  	originalData, err := io.ReadAll(r)
    79  	if err != nil {
    80  		return nil, err
    81  	}
    82  	if err := r.Close(); err != nil {
    83  		return nil, err
    84  	}
    85  	return originalData, nil
    86  }