github.com/LagrangeDev/LagrangeGo@v0.0.0-20240512064304-ad4a85e10cb4/utils/binary/pool.go (about)

     1  package binary
     2  
     3  // from https://github.com/Mrs4s/MiraiGo/blob/master/binary/pool.go
     4  
     5  import (
     6  	"bytes"
     7  	"compress/gzip"
     8  	"compress/zlib"
     9  	"sync"
    10  )
    11  
    12  var bufferPool = sync.Pool{
    13  	New: func() any {
    14  		return new(Builder)
    15  	},
    16  }
    17  
    18  // SelectBuilder 从池中取出一个 Builder
    19  func SelectBuilder(key []byte) *Builder {
    20  	// 因为 bufferPool 定义有 New 函数
    21  	// 所以 bufferPool.Get() 永不为 nil
    22  	// 不用判空
    23  	return bufferPool.Get().(*Builder).init(key)
    24  }
    25  
    26  // PutBuilder 将 Builder 放回池中
    27  func PutBuilder(w *Builder) {
    28  	// See https://golang.org/issue/23199
    29  	const maxSize = 32 * 1024
    30  	if w.hasput {
    31  		return
    32  	}
    33  	w.hasput = true
    34  	if w.buffer.Cap() < maxSize { // 对于大Buffer直接丢弃
    35  		w.buffer.Reset()
    36  		bufferPool.Put(w)
    37  	}
    38  }
    39  
    40  var gzipPool = sync.Pool{
    41  	New: func() any {
    42  		buf := new(bytes.Buffer)
    43  		w := gzip.NewWriter(buf)
    44  		return &GzipWriter{
    45  			w:   w,
    46  			buf: buf,
    47  		}
    48  	},
    49  }
    50  
    51  func AcquireGzipWriter() *GzipWriter {
    52  	ret := gzipPool.Get().(*GzipWriter)
    53  	ret.buf.Reset()
    54  	ret.w.Reset(ret.buf)
    55  	return ret
    56  }
    57  
    58  func ReleaseGzipWriter(w *GzipWriter) {
    59  	// See https://golang.org/issue/23199
    60  	const maxSize = 1 << 16
    61  	if w.buf.Cap() < maxSize {
    62  		w.buf.Reset()
    63  		gzipPool.Put(w)
    64  	}
    65  }
    66  
    67  type zlibWriter struct {
    68  	w   *zlib.Writer
    69  	buf *bytes.Buffer
    70  }
    71  
    72  var zlibPool = sync.Pool{
    73  	New: func() any {
    74  		buf := new(bytes.Buffer)
    75  		w := zlib.NewWriter(buf)
    76  		return &zlibWriter{
    77  			w:   w,
    78  			buf: buf,
    79  		}
    80  	},
    81  }
    82  
    83  func acquireZlibWriter() *zlibWriter {
    84  	ret := zlibPool.Get().(*zlibWriter)
    85  	ret.buf.Reset()
    86  	ret.w.Reset(ret.buf)
    87  	return ret
    88  }
    89  
    90  func releaseZlibWriter(w *zlibWriter) {
    91  	// See https://golang.org/issue/23199
    92  	const maxSize = 1 << 16
    93  	if w.buf.Cap() < maxSize {
    94  		w.buf.Reset()
    95  		zlibPool.Put(w)
    96  	}
    97  }