@@ -22,6 +22,7 @@ import (
2222
2323	apiTypes "github.com/dragonflyoss/Dragonfly/apis/types" 
2424	"github.com/dragonflyoss/Dragonfly/pkg/constants" 
25+ 	"github.com/dragonflyoss/Dragonfly/pkg/pool" 
2526)
2627
2728// Piece contains all information of a piece. 
@@ -51,14 +52,25 @@ type Piece struct {
5152	PieceNum  int  `json:"pieceNum"` 
5253
5354	// Content uses a buffer to temporarily store the piece content. 
54- 	Content  * bytes.Buffer  `json:"-"` 
55+ 	Content  * pool.Buffer  `json:"-"` 
56+ 
57+ 	// length the length of the content. 
58+ 	length  int64 
59+ 
60+ 	// autoReset automatically reset content after reading. 
61+ 	autoReset  bool 
5562}
5663
5764// RawContent returns raw contents, 
5865// If the piece has wrapper, and the piece content will remove the head and tail. 
5966func  (p  * Piece ) RawContent (noWrapper  bool ) * bytes.Buffer  {
6067	contents  :=  p .Content .Bytes ()
6168	length  :=  len (contents )
69+ 	defer  func () {
70+ 		if  p .autoReset  {
71+ 			p .ResetContent ()
72+ 		}
73+ 	}()
6274
6375	if  noWrapper  {
6476		return  bytes .NewBuffer (contents [:])
@@ -69,13 +81,33 @@ func (p *Piece) RawContent(noWrapper bool) *bytes.Buffer {
6981	return  nil 
7082}
7183
84+ // ContentLength returns the content length. 
85+ func  (p  * Piece ) ContentLength () int64  {
86+ 	if  p .length  <=  0  &&  p .Content  !=  nil  {
87+ 		p .length  =  int64 (p .Content .Len ())
88+ 	}
89+ 	return  p .length 
90+ }
91+ 
7292func  (p  * Piece ) String () string  {
7393	if  b , e  :=  json .Marshal (p ); e  ==  nil  {
7494		return  string (b )
7595	}
7696	return  "" 
7797}
7898
99+ // ResetContent reset contents and returns it back to buffer pool. 
100+ func  (p  * Piece ) ResetContent () {
101+ 	if  p .Content  ==  nil  {
102+ 		return 
103+ 	}
104+ 	if  p .length  ==  0  {
105+ 		p .length  =  int64 (p .Content .Len ())
106+ 	}
107+ 	pool .ReleaseBuffer (p .Content )
108+ 	p .Content  =  nil 
109+ }
110+ 
79111// NewPiece creates a Piece. 
80112func  NewPiece (taskID , node , dstCid , pieceRange  string , result , status  int , cdnSource  apiTypes.CdnSource ) * Piece  {
81113	return  & Piece {
@@ -85,7 +117,8 @@ func NewPiece(taskID, node, dstCid, pieceRange string, result, status int, cdnSo
85117		Range :     pieceRange ,
86118		Result :    result ,
87119		Status :    status ,
88- 		Content :   & bytes.Buffer {},
120+ 		Content :   nil ,
121+ 		autoReset : true ,
89122	}
90123}
91124
@@ -96,16 +129,14 @@ func NewPieceSimple(taskID string, node string, status int, cdnSource apiTypes.C
96129		SuperNode : node ,
97130		Status :    status ,
98131		Result :    constants .ResultInvalid ,
99- 		Content :   & bytes.Buffer {},
132+ 		Content :   nil ,
133+ 		autoReset : true ,
100134	}
101135}
102136
103137// NewPieceContent creates a Piece with specified content. 
104138func  NewPieceContent (taskID , node , dstCid , pieceRange  string ,
105- 	result , status  int , contents  * bytes.Buffer , cdnSource  apiTypes.CdnSource ) * Piece  {
106- 	if  contents  ==  nil  {
107- 		contents  =  & bytes.Buffer {}
108- 	}
139+ 	result , status  int , contents  * pool.Buffer , cdnSource  apiTypes.CdnSource ) * Piece  {
109140	return  & Piece {
110141		TaskID :    taskID ,
111142		SuperNode : node ,
@@ -114,5 +145,7 @@ func NewPieceContent(taskID, node, dstCid, pieceRange string,
114145		Result :    result ,
115146		Status :    status ,
116147		Content :   contents ,
148+ 		length :    int64 (contents .Len ()),
149+ 		autoReset : true ,
117150	}
118151}
0 commit comments