github.com/TeaOSLab/EdgeNode@v1.3.8/internal/compressions/utils.go (about)

     1  // Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package compressions
     4  
     5  import (
     6  	"errors"
     7  	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
     8  	memutils "github.com/TeaOSLab/EdgeNode/internal/utils/mem"
     9  	"io"
    10  	"net/http"
    11  	"runtime"
    12  )
    13  
    14  type ContentEncoding = string
    15  
    16  const (
    17  	ContentEncodingBr      ContentEncoding = "br"
    18  	ContentEncodingGzip    ContentEncoding = "gzip"
    19  	ContentEncodingDeflate ContentEncoding = "deflate"
    20  	ContentEncodingZSTD    ContentEncoding = "zstd"
    21  )
    22  
    23  var ErrNotSupportedContentEncoding = errors.New("not supported content encoding")
    24  
    25  // AllEncodings 当前支持的所有编码
    26  func AllEncodings() []ContentEncoding {
    27  	return []ContentEncoding{
    28  		ContentEncodingBr,
    29  		ContentEncodingGzip,
    30  		ContentEncodingZSTD,
    31  		ContentEncodingDeflate,
    32  	}
    33  }
    34  
    35  // NewReader 获取Reader
    36  func NewReader(reader io.Reader, contentEncoding ContentEncoding) (Reader, error) {
    37  	switch contentEncoding {
    38  	case ContentEncodingBr:
    39  		return NewBrotliReader(reader)
    40  	case ContentEncodingGzip:
    41  		return NewGzipReader(reader)
    42  	case ContentEncodingDeflate:
    43  		return NewDeflateReader(reader)
    44  	case ContentEncodingZSTD:
    45  		return NewZSTDReader(reader)
    46  	}
    47  	return nil, ErrNotSupportedContentEncoding
    48  }
    49  
    50  // NewWriter 获取Writer
    51  func NewWriter(writer io.Writer, compressType serverconfigs.HTTPCompressionType, level int) (Writer, error) {
    52  	switch compressType {
    53  	case serverconfigs.HTTPCompressionTypeGzip:
    54  		return NewGzipWriter(writer, level)
    55  	case serverconfigs.HTTPCompressionTypeDeflate:
    56  		return NewDeflateWriter(writer, level)
    57  	case serverconfigs.HTTPCompressionTypeBrotli:
    58  		return NewBrotliWriter(writer, level)
    59  	case serverconfigs.HTTPCompressionTypeZSTD:
    60  		return NewZSTDWriter(writer, level)
    61  	}
    62  	return nil, errors.New("invalid compression type '" + compressType + "'")
    63  }
    64  
    65  // SupportEncoding 检查是否支持某个编码
    66  func SupportEncoding(encoding string) bool {
    67  	return encoding == ContentEncodingBr ||
    68  		encoding == ContentEncodingGzip ||
    69  		encoding == ContentEncodingDeflate ||
    70  		encoding == ContentEncodingZSTD
    71  }
    72  
    73  // WrapHTTPResponse 包装http.Response对象
    74  func WrapHTTPResponse(resp *http.Response) {
    75  	if resp == nil {
    76  		return
    77  	}
    78  
    79  	var contentEncoding = resp.Header.Get("Content-Encoding")
    80  	if len(contentEncoding) == 0 || !SupportEncoding(contentEncoding) {
    81  		return
    82  	}
    83  
    84  	reader, err := NewReader(resp.Body, contentEncoding)
    85  	if err != nil {
    86  		// unable to decode, we ignore the error
    87  		return
    88  	}
    89  	resp.Header.Del("Content-Encoding")
    90  	resp.Header.Del("Content-Length")
    91  	resp.Body = reader
    92  }
    93  
    94  // 系统CPU线程数
    95  var countCPU = runtime.NumCPU()
    96  
    97  // GenerateCompressLevel 根据系统资源自动生成压缩级别
    98  func GenerateCompressLevel(minLevel int, maxLevel int) (level int) {
    99  	if countCPU < 16 {
   100  		return minLevel
   101  	}
   102  
   103  	if countCPU < 32 {
   104  		return min(3, maxLevel)
   105  	}
   106  
   107  	return min(5, maxLevel)
   108  }
   109  
   110  // CalculatePoolSize 计算Pool尺寸
   111  func CalculatePoolSize() int {
   112  	var maxSize = memutils.SystemMemoryGB() * 32
   113  	if maxSize == 0 {
   114  		maxSize = 128
   115  	}
   116  	if maxSize > 2048 {
   117  		maxSize = 2048
   118  	}
   119  	return maxSize
   120  }