github.com/cloudwego/kitex@v0.9.0/pkg/remote/codec/protobuf/encoding/gzip/gzip.go (about)

     1  /*
     2   *
     3   * Copyright 2017 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   * This file may have been modified by CloudWeGo authors. All CloudWeGo
    18   * Modifications are Copyright 2023 CloudWeGo Authors.
    19   */
    20  
    21  // Package gzip implements and registers the gzip compressor
    22  // during the initialization.
    23  //
    24  // # Experimental
    25  //
    26  // Notice: This package is EXPERIMENTAL and may be changed or removed in a
    27  // later release.
    28  package gzip
    29  
    30  import (
    31  	"compress/gzip"
    32  	"encoding/binary"
    33  	"fmt"
    34  	"io"
    35  	"io/ioutil"
    36  	"sync"
    37  
    38  	"github.com/cloudwego/kitex/pkg/remote/codec/protobuf/encoding"
    39  )
    40  
    41  // Name is the name registered for the gzip compressor.
    42  const Name = "gzip"
    43  
    44  func init() {
    45  	c := &compressor{}
    46  	c.poolCompressor.New = func() interface{} {
    47  		return &writer{Writer: gzip.NewWriter(ioutil.Discard), pool: &c.poolCompressor}
    48  	}
    49  	encoding.RegisterCompressor(c)
    50  }
    51  
    52  type writer struct {
    53  	*gzip.Writer
    54  	pool *sync.Pool
    55  }
    56  
    57  // SetLevel updates the registered gzip compressor to use the compression level specified (gzip.HuffmanOnly is not supported).
    58  // NOTE: this function must only be called during initialization time (i.e. in an init() function),
    59  // and is not thread-safe.
    60  //
    61  // The error returned will be nil if the specified level is valid.
    62  func SetLevel(level int) error {
    63  	if level < gzip.DefaultCompression || level > gzip.BestCompression {
    64  		return fmt.Errorf("grpc: invalid gzip compression level: %d", level)
    65  	}
    66  	c := encoding.GetCompressor(Name).(*compressor)
    67  	c.poolCompressor.New = func() interface{} {
    68  		w, err := gzip.NewWriterLevel(ioutil.Discard, level)
    69  		if err != nil {
    70  			panic(err)
    71  		}
    72  		return &writer{Writer: w, pool: &c.poolCompressor}
    73  	}
    74  	return nil
    75  }
    76  
    77  func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
    78  	z := c.poolCompressor.Get().(*writer)
    79  	z.Writer.Reset(w)
    80  	return z, nil
    81  }
    82  
    83  func (z *writer) Close() error {
    84  	defer z.pool.Put(z)
    85  	return z.Writer.Close()
    86  }
    87  
    88  type reader struct {
    89  	*gzip.Reader
    90  	pool *sync.Pool
    91  }
    92  
    93  func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
    94  	z, inPool := c.poolDecompressor.Get().(*reader)
    95  	if !inPool {
    96  		newZ, err := gzip.NewReader(r)
    97  		if err != nil {
    98  			return nil, err
    99  		}
   100  		return &reader{Reader: newZ, pool: &c.poolDecompressor}, nil
   101  	}
   102  	if err := z.Reset(r); err != nil {
   103  		c.poolDecompressor.Put(z)
   104  		return nil, err
   105  	}
   106  	return z, nil
   107  }
   108  
   109  func (z *reader) Read(p []byte) (n int, err error) {
   110  	n, err = z.Reader.Read(p)
   111  	if err == io.EOF {
   112  		z.pool.Put(z)
   113  	}
   114  	return n, err
   115  }
   116  
   117  // RFC1952 specifies that the last four bytes "contains the size of
   118  // the original (uncompressed) input data modulo 2^32."
   119  // gRPC has a max message size of 2GB so we don't need to worry about wraparound.
   120  func (c *compressor) DecompressedSize(buf []byte) int {
   121  	last := len(buf)
   122  	if last < 4 {
   123  		return -1
   124  	}
   125  	return int(binary.LittleEndian.Uint32(buf[last-4 : last]))
   126  }
   127  
   128  func (c *compressor) Name() string {
   129  	return Name
   130  }
   131  
   132  type compressor struct {
   133  	poolCompressor   sync.Pool
   134  	poolDecompressor sync.Pool
   135  }