trpc.group/trpc-go/trpc-go@v1.0.3/restful/compress_gzip.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package restful
    15  
    16  import (
    17  	"compress/gzip"
    18  	"io"
    19  	"sync"
    20  )
    21  
    22  func init() {
    23  	RegisterCompressor(&GZIPCompressor{})
    24  }
    25  
    26  var readerPool sync.Pool
    27  var writerPool sync.Pool
    28  
    29  // GZIPCompressor is the compressor for Content-Encoding: gzip.
    30  type GZIPCompressor struct{}
    31  
    32  // wrappedWriter wraps gzip.Writer for pooling.
    33  type wrappedWriter struct {
    34  	*gzip.Writer
    35  }
    36  
    37  // Close rewrites the underlying gzip.Writer's Close method.
    38  // The wrapped writer will be put back to the pool after its underlying gzip.Writer is closed.
    39  func (w *wrappedWriter) Close() error {
    40  	defer writerPool.Put(w)
    41  	return w.Writer.Close()
    42  }
    43  
    44  // wrappedReader wraps gzip.Reader for pooling.
    45  type wrappedReader struct {
    46  	*gzip.Reader
    47  }
    48  
    49  // Read rewrites the underlying gzip.Reader's Read method.
    50  // The wrapped reader will be put back to the pool after its underlying gzip.Reader is read.
    51  func (r *wrappedReader) Read(p []byte) (int, error) {
    52  	n, err := r.Reader.Read(p)
    53  	if err == io.EOF {
    54  		readerPool.Put(r)
    55  	}
    56  	return n, err
    57  }
    58  
    59  // Compress implements Compressor.
    60  func (*GZIPCompressor) Compress(w io.Writer) (io.WriteCloser, error) {
    61  	z, ok := writerPool.Get().(*wrappedWriter)
    62  	if !ok {
    63  		z = &wrappedWriter{
    64  			Writer: gzip.NewWriter(w),
    65  		}
    66  	}
    67  	z.Writer.Reset(w)
    68  	return z, nil
    69  }
    70  
    71  // Decompress implements Compressor.
    72  func (g *GZIPCompressor) Decompress(r io.Reader) (io.Reader, error) {
    73  	z, ok := readerPool.Get().(*wrappedReader)
    74  	if !ok {
    75  		gzipReader, err := gzip.NewReader(r)
    76  		if err != nil {
    77  			return nil, err
    78  		}
    79  		return &wrappedReader{
    80  			Reader: gzipReader,
    81  		}, nil
    82  	}
    83  	if err := z.Reader.Reset(r); err != nil {
    84  		readerPool.Put(z)
    85  		return nil, err
    86  	}
    87  	return z, nil
    88  }
    89  
    90  // Name implements Compressor.
    91  func (*GZIPCompressor) Name() string {
    92  	return "gzip"
    93  }
    94  
    95  // ContentEncoding implements Compressor.
    96  func (*GZIPCompressor) ContentEncoding() string {
    97  	return "gzip"
    98  }