diff --git a/dfget/core/downloader/p2p_downloader/client_stream_writer_test.go b/dfget/core/downloader/p2p_downloader/client_stream_writer_test.go index 7e5e38179..ae3216a21 100644 --- a/dfget/core/downloader/p2p_downloader/client_stream_writer_test.go +++ b/dfget/core/downloader/p2p_downloader/client_stream_writer_test.go @@ -17,11 +17,11 @@ package downloader import ( - "bytes" "io" "sort" "github.com/dragonflyoss/Dragonfly/dfget/config" + "github.com/dragonflyoss/Dragonfly/pkg/pool" "github.com/go-check/check" ) @@ -49,7 +49,7 @@ func (s *ClientStreamWriterTestSuite) TestWrite(c *check.C) { piece: &Piece{ PieceNum: 0, PieceSize: 6, - Content: bytes.NewBufferString("000010"), + Content: pool.NewBufferString("000010"), }, noWrapper: false, expected: "1", @@ -58,7 +58,7 @@ func (s *ClientStreamWriterTestSuite) TestWrite(c *check.C) { piece: &Piece{ PieceNum: 1, PieceSize: 6, - Content: bytes.NewBufferString("000020"), + Content: pool.NewBufferString("000020"), }, noWrapper: false, expected: "2", @@ -67,7 +67,7 @@ func (s *ClientStreamWriterTestSuite) TestWrite(c *check.C) { piece: &Piece{ PieceNum: 3, PieceSize: 6, - Content: bytes.NewBufferString("000040"), + Content: pool.NewBufferString("000040"), }, noWrapper: false, expected: "4", @@ -76,7 +76,7 @@ func (s *ClientStreamWriterTestSuite) TestWrite(c *check.C) { piece: &Piece{ PieceNum: 4, PieceSize: 6, - Content: bytes.NewBufferString("000050"), + Content: pool.NewBufferString("000050"), }, noWrapper: false, expected: "5", @@ -85,7 +85,7 @@ func (s *ClientStreamWriterTestSuite) TestWrite(c *check.C) { piece: &Piece{ PieceNum: 2, PieceSize: 6, - Content: bytes.NewBufferString("000030"), + Content: pool.NewBufferString("000030"), }, noWrapper: false, expected: "3", diff --git a/dfget/core/downloader/p2p_downloader/client_writer.go b/dfget/core/downloader/p2p_downloader/client_writer.go index a19bbef57..1833aba10 100644 --- a/dfget/core/downloader/p2p_downloader/client_writer.go +++ b/dfget/core/downloader/p2p_downloader/client_writer.go @@ -17,7 +17,6 @@ package downloader import ( - "bufio" "context" "io" "math/rand" @@ -31,6 +30,7 @@ import ( "github.com/dragonflyoss/Dragonfly/dfget/core/helper" "github.com/dragonflyoss/Dragonfly/dfget/types" "github.com/dragonflyoss/Dragonfly/pkg/fileutils" + "github.com/dragonflyoss/Dragonfly/pkg/pool" "github.com/dragonflyoss/Dragonfly/pkg/queue" "github.com/sirupsen/logrus" @@ -244,9 +244,10 @@ func writePieceToFile(piece *Piece, file *os.File, cdnSource apiTypes.CdnSource) return err } - buf := bufio.NewWriterSize(file, 4*1024*1024) - _, err := io.Copy(buf, piece.RawContent(noWrapper)) - buf.Flush() + writer := pool.AcquireWriter(file) + _, err := io.Copy(writer, piece.RawContent(noWrapper)) + pool.ReleaseWriter(writer) + writer = nil return err } diff --git a/dfget/core/downloader/p2p_downloader/client_writer_test.go b/dfget/core/downloader/p2p_downloader/client_writer_test.go index 22326e085..cf3497a9f 100644 --- a/dfget/core/downloader/p2p_downloader/client_writer_test.go +++ b/dfget/core/downloader/p2p_downloader/client_writer_test.go @@ -17,7 +17,6 @@ package downloader import ( - "bytes" "fmt" "io/ioutil" "os" @@ -25,6 +24,7 @@ import ( apiTypes "github.com/dragonflyoss/Dragonfly/apis/types" "github.com/dragonflyoss/Dragonfly/pkg/fileutils" + "github.com/dragonflyoss/Dragonfly/pkg/pool" "github.com/go-check/check" ) @@ -63,7 +63,7 @@ func (s *ClientWriterTestSuite) TestWrite(c *check.C) { piece: &Piece{ PieceNum: 0, PieceSize: 6, - Content: bytes.NewBufferString("000010"), + Content: pool.NewBufferString("000010"), }, cdnSource: apiTypes.CdnSourceSupernode, expected: "1", @@ -72,7 +72,7 @@ func (s *ClientWriterTestSuite) TestWrite(c *check.C) { piece: &Piece{ PieceNum: 1, PieceSize: 6, - Content: bytes.NewBufferString("000020"), + Content: pool.NewBufferString("000020"), }, cdnSource: apiTypes.CdnSourceSupernode, expected: "2", @@ -81,7 +81,7 @@ func (s *ClientWriterTestSuite) TestWrite(c *check.C) { piece: &Piece{ PieceNum: 1, PieceSize: 6, - Content: bytes.NewBufferString("000030"), + Content: pool.NewBufferString("000030"), }, cdnSource: apiTypes.CdnSourceSource, expected: "000030", diff --git a/dfget/core/downloader/p2p_downloader/p2p_downloader.go b/dfget/core/downloader/p2p_downloader/p2p_downloader.go index dcfbadcf4..8566c3c8d 100644 --- a/dfget/core/downloader/p2p_downloader/p2p_downloader.go +++ b/dfget/core/downloader/p2p_downloader/p2p_downloader.go @@ -17,7 +17,6 @@ package downloader import ( - "bytes" "context" "fmt" "io" @@ -212,7 +211,7 @@ func (p2p *P2PDownloader) run(ctx context.Context, pieceWriter PieceWriter) erro logrus.Infof("downloading piece:%v", lastItem) curItem := *lastItem - curItem.Content = &bytes.Buffer{} + curItem.Content = nil lastItem = nil response, err := p2p.pullPieceTask(&curItem) @@ -412,7 +411,7 @@ func (p2p *P2PDownloader) getItem(latestItem *Piece) (bool, *Piece) { } if !v && (item.Result == constants.ResultSemiSuc || item.Result == constants.ResultSuc) { - p2p.total += int64(item.Content.Len()) + p2p.total += item.ContentLength() p2p.pieceSet[item.Range] = true } else if !v { delete(p2p.pieceSet, item.Range) diff --git a/dfget/core/downloader/p2p_downloader/piece.go b/dfget/core/downloader/p2p_downloader/piece.go index 8b048293f..67a932d19 100644 --- a/dfget/core/downloader/p2p_downloader/piece.go +++ b/dfget/core/downloader/p2p_downloader/piece.go @@ -22,6 +22,7 @@ import ( apiTypes "github.com/dragonflyoss/Dragonfly/apis/types" "github.com/dragonflyoss/Dragonfly/pkg/constants" + "github.com/dragonflyoss/Dragonfly/pkg/pool" ) // Piece contains all information of a piece. @@ -51,7 +52,13 @@ type Piece struct { PieceNum int `json:"pieceNum"` // Content uses a buffer to temporarily store the piece content. - Content *bytes.Buffer `json:"-"` + Content *pool.Buffer `json:"-"` + + // length the length of the content. + length int64 + + // autoReset automatically reset content after reading. + autoReset bool } // RawContent returns raw contents, @@ -59,6 +66,11 @@ type Piece struct { func (p *Piece) RawContent(noWrapper bool) *bytes.Buffer { contents := p.Content.Bytes() length := len(contents) + defer func() { + if p.autoReset { + p.ResetContent() + } + }() if noWrapper { return bytes.NewBuffer(contents[:]) @@ -69,6 +81,14 @@ func (p *Piece) RawContent(noWrapper bool) *bytes.Buffer { return nil } +// ContentLength returns the content length. +func (p *Piece) ContentLength() int64 { + if p.length <= 0 && p.Content != nil { + p.length = int64(p.Content.Len()) + } + return p.length +} + func (p *Piece) String() string { if b, e := json.Marshal(p); e == nil { return string(b) @@ -76,6 +96,18 @@ func (p *Piece) String() string { return "" } +// ResetContent reset contents and returns it back to buffer pool. +func (p *Piece) ResetContent() { + if p.Content == nil { + return + } + if p.length == 0 { + p.length = int64(p.Content.Len()) + } + pool.ReleaseBuffer(p.Content) + p.Content = nil +} + // NewPiece creates a Piece. func NewPiece(taskID, node, dstCid, pieceRange string, result, status int, cdnSource apiTypes.CdnSource) *Piece { return &Piece{ @@ -85,7 +117,8 @@ func NewPiece(taskID, node, dstCid, pieceRange string, result, status int, cdnSo Range: pieceRange, Result: result, Status: status, - Content: &bytes.Buffer{}, + Content: nil, + autoReset: true, } } @@ -96,16 +129,14 @@ func NewPieceSimple(taskID string, node string, status int, cdnSource apiTypes.C SuperNode: node, Status: status, Result: constants.ResultInvalid, - Content: &bytes.Buffer{}, + Content: nil, + autoReset: true, } } // NewPieceContent creates a Piece with specified content. func NewPieceContent(taskID, node, dstCid, pieceRange string, - result, status int, contents *bytes.Buffer, cdnSource apiTypes.CdnSource) *Piece { - if contents == nil { - contents = &bytes.Buffer{} - } + result, status int, contents *pool.Buffer, cdnSource apiTypes.CdnSource) *Piece { return &Piece{ TaskID: taskID, SuperNode: node, @@ -114,5 +145,7 @@ func NewPieceContent(taskID, node, dstCid, pieceRange string, Result: result, Status: status, Content: contents, + length: int64(contents.Len()), + autoReset: true, } } diff --git a/dfget/core/downloader/p2p_downloader/piece_test.go b/dfget/core/downloader/p2p_downloader/piece_test.go index 68c943802..6cdea206a 100644 --- a/dfget/core/downloader/p2p_downloader/piece_test.go +++ b/dfget/core/downloader/p2p_downloader/piece_test.go @@ -20,6 +20,8 @@ import ( "bytes" "github.com/go-check/check" + + "github.com/dragonflyoss/Dragonfly/pkg/pool" ) type PieceTestSuite struct { @@ -35,9 +37,9 @@ func (s *PieceTestSuite) TestRawContent(c *check.C) { noWrapper bool expected *bytes.Buffer }{ - {piece: &Piece{Content: bytes.NewBufferString("")}, noWrapper: false, expected: nil}, - {piece: &Piece{Content: bytes.NewBufferString("000010")}, noWrapper: false, expected: bytes.NewBufferString("1")}, - {piece: &Piece{Content: bytes.NewBufferString("000020")}, noWrapper: true, expected: bytes.NewBufferString("000020")}, + {piece: &Piece{Content: pool.NewBufferString("")}, noWrapper: false, expected: nil}, + {piece: &Piece{Content: pool.NewBufferString("000010")}, noWrapper: false, expected: bytes.NewBufferString("1")}, + {piece: &Piece{Content: pool.NewBufferString("000020")}, noWrapper: true, expected: bytes.NewBufferString("000020")}, } for _, v := range cases { diff --git a/dfget/core/downloader/p2p_downloader/power_client.go b/dfget/core/downloader/p2p_downloader/power_client.go index 3b75b434c..17f5a30a3 100644 --- a/dfget/core/downloader/p2p_downloader/power_client.go +++ b/dfget/core/downloader/p2p_downloader/power_client.go @@ -33,6 +33,7 @@ import ( "github.com/dragonflyoss/Dragonfly/pkg/httputils" "github.com/dragonflyoss/Dragonfly/pkg/limitreader" "github.com/dragonflyoss/Dragonfly/pkg/netutils" + "github.com/dragonflyoss/Dragonfly/pkg/pool" "github.com/dragonflyoss/Dragonfly/pkg/queue" "github.com/dragonflyoss/Dragonfly/pkg/ratelimiter" @@ -113,7 +114,7 @@ func (pc *PowerClient) ClientError() *types.ClientErrorRequest { return pc.clientError } -func (pc *PowerClient) downloadPiece() (content *bytes.Buffer, e error) { +func (pc *PowerClient) downloadPiece() (content *pool.Buffer, e error) { dstIP := pc.pieceTask.PeerIP peerPort := pc.pieceTask.PeerPort @@ -149,7 +150,14 @@ func (pc *PowerClient) downloadPiece() (content *bytes.Buffer, e error) { // start to read data from resp // use limitReader to limit the download speed limitReader := limitreader.NewLimitReaderWithLimiter(pc.rateLimiter, resp.Body, pieceMD5 != "") - content = &bytes.Buffer{} + content = pool.AcquireBufferSize(int(pc.pieceTask.PieceSize)) + defer func() { + // if an error happened, the content cannot be released outside. + if e != nil { + pool.ReleaseBuffer(content) + content = nil + } + }() if pc.total, e = content.ReadFrom(limitReader); e != nil { return nil, e } @@ -193,7 +201,7 @@ func (pc *PowerClient) createDownloadRequest() *api.DownloadRequest { } } -func (pc *PowerClient) successPiece(content *bytes.Buffer) *Piece { +func (pc *PowerClient) successPiece(content *pool.Buffer) *Piece { piece := NewPieceContent(pc.taskID, pc.node, pc.pieceTask.Cid, pc.pieceTask.Range, constants.ResultSemiSuc, constants.TaskStatusRunning, content, pc.cdnSource) piece.PieceSize = pc.pieceTask.PieceSize diff --git a/dfget/core/downloader/p2p_downloader/power_client_test.go b/dfget/core/downloader/p2p_downloader/power_client_test.go index 71b13e5aa..55695214d 100644 --- a/dfget/core/downloader/p2p_downloader/power_client_test.go +++ b/dfget/core/downloader/p2p_downloader/power_client_test.go @@ -136,7 +136,8 @@ func (s *PowerClientTestSuite) TestDownloadPiece(c *check.C) { return resp, nil } content, err = s.powerClient.downloadPiece() - c.Check(content, check.DeepEquals, bytes.NewBufferString("hello")) + c.Check(content, check.NotNil) + c.Check(content.String(), check.Equals, "hello") c.Check(err, check.IsNil) } diff --git a/pkg/pool/buffer_pool.go b/pkg/pool/buffer_pool.go new file mode 100644 index 000000000..85d0827a1 --- /dev/null +++ b/pkg/pool/buffer_pool.go @@ -0,0 +1,174 @@ +/* + * Copyright The Dragonfly Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package pool + +import ( + "bytes" + "io" + "sync" +) + +const ( + // defaultAllocatedSize 2MB + defaultAllocatedSize int = 2 * 1024 * 1024 + + maxPoolCount = 8 + minPoolCount = 1 +) + +// bufferPools is the default pool which has 4 intervals of buffers' +// initial capacity: +// (0,2MB], (2MB, 4MB], (4MB, 8MB], (8MB, +∞) +var bufferPool = NewBufferPool(4, defaultAllocatedSize) + +// AcquireBufferSize returns an empty Buffer instance from buffer pool, +// whose capacity is greater than or equal to the giving size. +// +// The returned Buffer instance may be passed to ReleaseBuffer when it is +// no longer needed. This allows Buffer recycling, reduces GC pressure +// and usually improves performance. +// +// This function is recommended when you know how much memory you need +// actually. +func AcquireBufferSize(size int) *Buffer { + if buf := bufferPool.Get(size); buf != nil { + return buf + } + return NewBuffer(size) +} + +// AcquireBuffer returns an empty Buffer instance from buffer pool. +// +// The returned Buffer instance may be passed to ReleaseBuffer when it is +// no longer needed. This allows Buffer recycling, reduces GC pressure +// and usually improves performance. +func AcquireBuffer() *Buffer { + return AcquireBufferSize(defaultAllocatedSize) +} + +// ReleaseBuffer returns buf acquired via AcquireBuffer/AcquireBufferSize +// to buffer pool. +// +// It is forbidden accessing buf and/or its' members after returning +// it to buffer pool. +func ReleaseBuffer(buf *Buffer) { + if buf == nil { + return + } + buf.Reset() + bufferPool.Put(buf) +} + +// ---------------------------------------------------------------------------- +// struct BufferPool + +func NewBufferPool(count, base int) *BufferPool { + if count < minPoolCount { + count = minPoolCount + } else if count > maxPoolCount { + count = maxPoolCount + } + if base <= 0 { + base = defaultAllocatedSize + } + pool := &BufferPool{ + pools: make([]sync.Pool, count), + baseSize: base, + } + return pool +} + +// BufferPool stores several intervals of buffer's initial capacity, which can +// minimizing the allocation times by bytes.Buffer.grow(n int). +// +// It groups the scenarios of using buffer, tries to avoid a large buffer not +// being recycling because it's used by callers which only need a small one. +type BufferPool struct { + pools []sync.Pool + baseSize int +} + +// Get returns a buffer with a capacity from the buffer pool. +func (bp *BufferPool) Get(size int) *Buffer { + idx := bp.index(size) + if buf := bp.pools[idx].Get(); buf != nil { + return buf.(*Buffer) + } + return nil +} + +// Put puts the buf to the buffer pool. +func (bp *BufferPool) Put(buf *Buffer) { + if buf != nil { + idx := bp.index(buf.allocatedSize) + bp.pools[idx].Put(buf) + } +} + +// index finds the first index of pool whose buffer's capacity is greater than +// the giving capacity. +func (bp *BufferPool) index(allocatedSize int) int { + i, length := 0, len(bp.pools) + for c := bp.baseSize; i < length && c < allocatedSize; c *= 2 { + i++ + } + if i < length { + return i + } + return length - 1 +} + +// ---------------------------------------------------------------------------- +// struct Buffer + +// NewBuffer creates a new Buffer which initialized by empty content. +func NewBuffer(size int) *Buffer { + return &Buffer{ + Buffer: bytes.NewBuffer(make([]byte, 0, size)), + allocatedSize: size, + } +} + +// NewBufferString creates and initializes a new Buffer using string s as its +// content. +func NewBufferString(s string) *Buffer { + return &Buffer{ + Buffer: bytes.NewBufferString(s), + allocatedSize: len(s), + } +} + +var ( + _ io.ReaderFrom = &Buffer{} + _ io.WriterTo = &Buffer{} + _ io.ReadWriteCloser = &Buffer{} +) + +// Buffer provides byte buffer, which can be used for minimizing memory +// allocations and implements interfaces: io.ReaderFrom, io.WriterTo, +// io.ReadWriteCloser. +// +// The allocatedSize is the buffer's initial size. +type Buffer struct { + *bytes.Buffer + allocatedSize int +} + +func (b *Buffer) Close() error { + b.Reset() + return nil +} diff --git a/pkg/pool/buffer_pool_test.go b/pkg/pool/buffer_pool_test.go new file mode 100644 index 000000000..e499dcae1 --- /dev/null +++ b/pkg/pool/buffer_pool_test.go @@ -0,0 +1,129 @@ +/* + * Copyright The Dragonfly Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package pool + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/suite" +) + +func TestSuite(t *testing.T) { + suite.Run(t, new(BufferPoolTestSuite)) +} + +type BufferPoolTestSuite struct { + suite.Suite + tmpBufferPool *BufferPool +} + +func (s *BufferPoolTestSuite) SetupSuite() { + s.tmpBufferPool = bufferPool + bufferPool = NewBufferPool(8, defaultAllocatedSize) +} + +func (s *BufferPoolTestSuite) TearDownSuite() { + bufferPool = s.tmpBufferPool + s.tmpBufferPool = nil +} + +func (s *BufferPoolTestSuite) TestAcquireBuffer() { + buf := AcquireBuffer() + defer ReleaseBuffer(buf) + + s.NotNil(buf) + s.Equal(0, buf.Len(), "not empty") +} + +func (s *BufferPoolTestSuite) TestReleaseBuffer() { + buf1 := AcquireBuffer() + ReleaseBuffer(buf1) + ReleaseBuffer(nil) + + buf2 := AcquireBuffer() + buf3 := AcquireBuffer() + defer func() { + ReleaseBuffer(buf2) + ReleaseBuffer(buf3) + buf2, buf3 = nil, nil + }() + + s.NotNil(buf1) + s.NotNil(buf2) + s.NotNil(buf3) + s.True(buf1 == buf2, "should reuse an old buffer but got a new one") + s.True(buf1 != buf3, "should create a new buffer but got an old one") +} + +func (s *BufferPoolTestSuite) TestNewBufferPool() { + cases := []struct { + count int + base int + expectedCount int + expectedBase int + }{ + {0, 0, minPoolCount, defaultAllocatedSize}, + {maxPoolCount + 1, 0, maxPoolCount, defaultAllocatedSize}, + {2, 4, 2, 4}, + } + + for i, c := range cases { + msg := fmt.Sprintf("case %d: %v", i, c) + p := NewBufferPool(c.count, c.base) + s.Equal(c.expectedCount, len(p.pools), msg) + s.Equal(c.expectedBase, p.baseSize, msg) + } + +} + +func (s *BufferPoolTestSuite) TestBufferPool_index() { + count := 2 + base := 4 + p := NewBufferPool(count, base) + idx := func(i int) int { + if i < count { + return i + } + return count - 1 + } + + capacity := base + for i := 0; i <= count; i++ { + s.Equal(idx(i), p.index(capacity-1)) + s.Equal(idx(i), p.index(capacity)) + s.Equal(idx(i+1), p.index(capacity+1)) + capacity *= 2 + } + s.Equal(0, p.index(base)) +} + +func (s *BufferPoolTestSuite) TestBuffer_NewBufferString() { + buf := NewBufferString("hello") + s.Equal([]byte("hello"), buf.Bytes()) + buf.Close() + s.Equal(0, buf.Len()) +} + +func BenchmarkAcquireBuffer(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + buf := AcquireBuffer() + buf.WriteString("hello") + ReleaseBuffer(buf) + } +} diff --git a/pkg/pool/writer_pool.go b/pkg/pool/writer_pool.go new file mode 100644 index 000000000..31f5546c8 --- /dev/null +++ b/pkg/pool/writer_pool.go @@ -0,0 +1,48 @@ +/* + * Copyright The Dragonfly Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package pool + +import ( + "bufio" + "io" + "sync" +) + +var writerPool = &sync.Pool{} + +// defaultSize 1MB +const defaultSize = 1 * 1024 * 1024 + +// AcquireWriter returns an empty Writer instance from writer pool. +func AcquireWriter(w io.Writer) *bufio.Writer { + if writer := writerPool.Get(); writer != nil { + writer := writer.(*bufio.Writer) + writer.Reset(w) + return writer + } + return bufio.NewWriterSize(w, defaultSize) +} + +// ReleaseWriter returns buf acquired via AcquireWriter to writer pool. +// It will flush and reset the writer before putting to writer pool. +func ReleaseWriter(writer *bufio.Writer) { + if writer != nil { + _ = writer.Flush() + writer.Reset(nil) + writerPool.Put(writer) + } +} diff --git a/pkg/pool/writer_pool_test.go b/pkg/pool/writer_pool_test.go new file mode 100644 index 000000000..8d1fbb3b7 --- /dev/null +++ b/pkg/pool/writer_pool_test.go @@ -0,0 +1,61 @@ +/* + * Copyright The Dragonfly Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package pool + +import ( + "bytes" + "io/ioutil" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestWriter(t *testing.T) { + tmp := writerPool + writerPool = &sync.Pool{} + + buf := &bytes.Buffer{} + w1 := AcquireWriter(buf) + w1.WriteString("test") + w1.Flush() + require.Equal(t, "test", buf.String()) + + // get the old writer from pool + ReleaseWriter(w1) + w2 := AcquireWriter(buf) + require.True(t, w1 == w2) + + // get a new writer from pool + w3 := AcquireWriter(buf) + require.True(t, w1 != w3) + + ReleaseWriter(w2) + ReleaseWriter(w3) + + writerPool = tmp + tmp = nil +} + +func BenchmarkAcquireWriter(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + w := AcquireWriter(ioutil.Discard) + w.WriteString("test") + ReleaseWriter(w) + } +}