optimize: release allocated ByteBuffer explicitly

use sync.Pool to manager the bytes.Buffer and bufio.Writer to minimizing the memory allocation and usage Signed-off-by: lowzj <[email protected]>
dragonflyoss · Jul 1, 2020 · 0122fc9 · 0122fc9
1 parent 129d562
commit 0122fc9
Show file tree

Hide file tree

Showing 12 changed files with 487 additions and 31 deletions.
diff --git 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
@@ -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
@@ -17,14 +17,14 @@
 package downloader
 
 import (
-	"bytes"
 	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
 
 	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
@@ -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
@@ -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,14 +52,25 @@ 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,
 // If the piece has wrapper, and the piece content will remove the head and tail.
 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,13 +81,33 @@ 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)
 	}
 	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
@@ -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
@@ -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.cdnDource)
 	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
@@ -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)
 }