Browse Source

Grow download buffer if Content-Length is provided

DarthSim 6 years ago
parent
commit
19cac3dd98
2 changed files with 11 additions and 4 deletions
  1. 4 4
      bufpool.go
  2. 7 0
      download.go

+ 4 - 4
bufpool.go

@@ -29,10 +29,10 @@ func newBufPool(n int, size int) *bufPool {
 func (p *bufPool) grow() {
 	var buf *bytes.Buffer
 
-	if p.size == 0 {
-		buf = new(bytes.Buffer)
-	} else {
-		buf = bytes.NewBuffer(make([]byte, p.size, p.size))
+	buf = new(bytes.Buffer)
+
+	if p.size > 0 {
+		buf.Grow(p.size)
 	}
 
 	p.top = &bufPoolEntry{buf: buf, next: p.top}

+ 7 - 0
download.go

@@ -9,6 +9,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"strconv"
 	"time"
 
 	_ "image/gif"
@@ -102,6 +103,12 @@ func readAndCheckImage(ctx context.Context, res *http.Response) (context.Context
 		return ctx, cancel, err
 	}
 
+	if cls := res.Header.Get("Content-Length"); len(cls) > 0 {
+		if cl, err := strconv.Atoi(cls); err == nil && cl > buf.Len() && cl > buf.Cap() {
+			buf.Grow(cl - buf.Len())
+		}
+	}
+
 	if _, err = buf.ReadFrom(res.Body); err != nil {
 		return ctx, cancel, newError(404, err.Error(), msgSourceImageIsUnreachable)
 	}