github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/grpc/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   */
    18  
    19  // Package gzip implements and registers the gzip compressor
    20  // during the initialization.
    21  //
    22  // Experimental
    23  //
    24  // Notice: This package is EXPERIMENTAL and may be changed or removed in a
    25  // later release.
    26  package gzip
    27  
    28  import (
    29  	"compress/gzip"
    30  	"encoding/binary"
    31  	"fmt"
    32  	"io"
    33  	"io/ioutil"
    34  	"sync"
    35  
    36  	"github.com/hxx258456/ccgo/grpc/encoding"
    37  )
    38  
    39  // Name is the name registered for the gzip compressor.
    40  const Name = "gzip"
    41  
    42  func init() {
    43  	c := &compressor{}
    44  	c.poolCompressor.New = func() interface{} {
    45  		return &writer{Writer: gzip.NewWriter(ioutil.Discard), pool: &c.poolCompressor}
    46  	}
    47  	encoding.RegisterCompressor(c)
    48  }
    49  
    50  type writer struct {
    51  	*gzip.Writer
    52  	pool *sync.Pool
    53  }
    54  
    55  // SetLevel updates the registered gzip compressor to use the compression level specified (gzip.HuffmanOnly is not supported).
    56  // NOTE: this function must only be called during initialization time (i.e. in an init() function),
    57  // and is not thread-safe.
    58  //
    59  // The error returned will be nil if the specified level is valid.
    60  func SetLevel(level int) error {
    61  	if level < gzip.DefaultCompression || level > gzip.BestCompression {
    62  		return fmt.Errorf("grpc: invalid gzip compression level: %d", level)
    63  	}
    64  	c := encoding.GetCompressor(Name).(*compressor)
    65  	c.poolCompressor.New = func() interface{} {
    66  		w, err := gzip.NewWriterLevel(ioutil.Discard, level)
    67  		if err != nil {
    68  			panic(err)
    69  		}
    70  		return &writer{Writer: w, pool: &c.poolCompressor}
    71  	}
    72  	return nil
    73  }
    74  
    75  func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
    76  	z := c.poolCompressor.Get().(*writer)
    77  	z.Writer.Reset(w)
    78  	return z, nil
    79  }
    80  
    81  func (z *writer) Close() error {
    82  	defer z.pool.Put(z)
    83  	return z.Writer.Close()
    84  }
    85  
    86  type reader struct {
    87  	*gzip.Reader
    88  	pool *sync.Pool
    89  }
    90  
    91  func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
    92  	z, inPool := c.poolDecompressor.Get().(*reader)
    93  	if !inPool {
    94  		newZ, err := gzip.NewReader(r)
    95  		if err != nil {
    96  			return nil, err
    97  		}
    98  		return &reader{Reader: newZ, pool: &c.poolDecompressor}, nil
    99  	}
   100  	if err := z.Reset(r); err != nil {
   101  		c.poolDecompressor.Put(z)
   102  		return nil, err
   103  	}
   104  	return z, nil
   105  }
   106  
   107  func (z *reader) Read(p []byte) (n int, err error) {
   108  	n, err = z.Reader.Read(p)
   109  	if err == io.EOF {
   110  		z.pool.Put(z)
   111  	}
   112  	return n, err
   113  }
   114  
   115  // RFC1952 specifies that the last four bytes "contains the size of
   116  // the original (uncompressed) input data modulo 2^32."
   117  // gRPC has a max message size of 2GB so we don't need to worry about wraparound.
   118  func (c *compressor) DecompressedSize(buf []byte) int {
   119  	last := len(buf)
   120  	if last < 4 {
   121  		return -1
   122  	}
   123  	return int(binary.LittleEndian.Uint32(buf[last-4 : last]))
   124  }
   125  
   126  func (c *compressor) Name() string {
   127  	return Name
   128  }
   129  
   130  type compressor struct {
   131  	poolCompressor   sync.Pool
   132  	poolDecompressor sync.Pool
   133  }