trpc.group/trpc-go/trpc-go@v1.0.3/codec/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 codec
    15  
    16  import (
    17  	"bytes"
    18  	"compress/gzip"
    19  	"io"
    20  	"sync"
    21  )
    22  
    23  func init() {
    24  	RegisterCompressor(CompressTypeGzip, &GzipCompress{})
    25  }
    26  
    27  // GzipCompress is gzip compressor.
    28  type GzipCompress struct {
    29  	readerPool sync.Pool
    30  	writerPool sync.Pool
    31  }
    32  
    33  // Compress returns binary data compressed by gzip.
    34  func (c *GzipCompress) Compress(in []byte) ([]byte, error) {
    35  	if len(in) == 0 {
    36  		return in, nil
    37  	}
    38  
    39  	buffer := &bytes.Buffer{}
    40  	z, ok := c.writerPool.Get().(*gzip.Writer)
    41  	if !ok {
    42  		z = gzip.NewWriter(buffer)
    43  	} else {
    44  		z.Reset(buffer)
    45  	}
    46  	defer c.writerPool.Put(z)
    47  
    48  	if _, err := z.Write(in); err != nil {
    49  		return nil, err
    50  	}
    51  	if err := z.Close(); err != nil {
    52  		return nil, err
    53  	}
    54  
    55  	return buffer.Bytes(), nil
    56  }
    57  
    58  // Decompress returns binary data decompressed by gzip.
    59  func (c *GzipCompress) Decompress(in []byte) ([]byte, error) {
    60  	if len(in) == 0 {
    61  		return in, nil
    62  	}
    63  	br := bytes.NewReader(in)
    64  	z, ok := c.readerPool.Get().(*gzip.Reader)
    65  	defer func() {
    66  		if z != nil {
    67  			c.readerPool.Put(z)
    68  		}
    69  	}()
    70  	if !ok {
    71  		gr, err := gzip.NewReader(br)
    72  		if err != nil {
    73  			return nil, err
    74  		}
    75  		z = gr
    76  	} else {
    77  		if err := z.Reset(br); err != nil {
    78  			return nil, err
    79  		}
    80  	}
    81  	out, err := io.ReadAll(z)
    82  	if err != nil {
    83  		return nil, err
    84  	}
    85  	return out, nil
    86  }