github.com/Mrs4s/MiraiGo@v0.0.0-20240226124653-54bdd873e3fe/binary/pool.go (about)

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