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

     1  // Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package bfs
     4  
     5  import (
     6  	"github.com/TeaOSLab/EdgeNode/internal/utils/percpu"
     7  	"github.com/klauspost/compress/gzip"
     8  	"io"
     9  	"runtime"
    10  )
    11  
    12  var SharedCompressPool = NewGzipWriterPool()
    13  
    14  type GzipWriterPool struct {
    15  	c     chan *gzip.Writer
    16  	cList []chan *gzip.Writer
    17  }
    18  
    19  func NewGzipWriterPool() *GzipWriterPool {
    20  	const poolSize = 16
    21  
    22  	var countProcs = runtime.GOMAXPROCS(0)
    23  	if countProcs <= 0 {
    24  		countProcs = runtime.NumCPU()
    25  	}
    26  	countProcs *= 4
    27  
    28  	var cList []chan *gzip.Writer
    29  	for i := 0; i < countProcs; i++ {
    30  		cList = append(cList, make(chan *gzip.Writer, poolSize))
    31  	}
    32  
    33  	return &GzipWriterPool{
    34  		c:     make(chan *gzip.Writer, poolSize),
    35  		cList: cList,
    36  	}
    37  }
    38  
    39  func (this *GzipWriterPool) Get(rawWriter io.Writer) (*gzip.Writer, error) {
    40  	select {
    41  	case w := <-this.getC():
    42  		w.Reset(rawWriter)
    43  		return w, nil
    44  	default:
    45  		return gzip.NewWriterLevel(rawWriter, gzip.BestSpeed)
    46  	}
    47  }
    48  
    49  func (this *GzipWriterPool) Put(writer *gzip.Writer) {
    50  	select {
    51  	case this.getC() <- writer:
    52  	default:
    53  		// 不需要close，因为已经在使用的时候调用了
    54  	}
    55  }
    56  
    57  func (this *GzipWriterPool) getC() chan *gzip.Writer {
    58  	var procId = percpu.GetProcId()
    59  	if procId < len(this.cList) {
    60  		return this.cList[procId]
    61  	}
    62  	return this.c
    63  }