request zstd encoded files

related to #77
distr1 · May 30, 2020 · 46228cc · 46228cc
1 parent 6d0c16f
commit 46228cc
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 2 deletions.
diff --git a/go.mod b/go.mod
@@ -8,7 +8,7 @@ require (
 	github.com/google/go-github/v27 v27.0.6
 	github.com/google/renameio v0.1.0
 	github.com/jacobsa/fuse v0.0.0-20200423191118-1d001802f70a
-	github.com/klauspost/compress v1.10.5 // indirect
+	github.com/klauspost/compress v1.10.6
 	github.com/klauspost/pgzip v1.2.3
 	github.com/lib/pq v1.4.0
 	github.com/lpar/gzipped v1.1.0

diff --git a/go.sum b/go.sum
@@ -55,6 +55,8 @@ github.com/klauspost/compress v1.9.7 h1:hYW1gP94JUmAhBtJ+LNz5My+gBobDxPR1iVuKug2
 github.com/klauspost/compress v1.9.7/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
 github.com/klauspost/compress v1.10.5 h1:7q6vHIqubShURwQz8cQK6yIe/xC3IF0Vm7TGfqjewrc=
 github.com/klauspost/compress v1.10.5/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
+github.com/klauspost/compress v1.10.6 h1:SP6zavvTG3YjOosWePXFDlExpKIWMTO4SE/Y8MZB2vI=
+github.com/klauspost/compress v1.10.6/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 github.com/klauspost/pgzip v1.2.1 h1:oIPZROsWuPHpOdMVWLuJZXwgjhrW8r1yEX8UqMyeNHM=
 github.com/klauspost/pgzip v1.2.1/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs=
 github.com/klauspost/pgzip v1.2.3 h1:Ce2to9wvs/cuJ2b86/CKQoTYr9VHfpanYosZ0UBJqdw=

diff --git a/internal/repo/reader.go b/internal/repo/reader.go
@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/distr1/distri"
+	"github.com/klauspost/compress/zstd"
 )
 
 type ErrNotFound struct {
@@ -41,6 +42,20 @@ func (r *gzipReader) Close() error {
 	return r.body.Close()
 }
 
+type zstdReader struct {
+	body io.ReadCloser
+	dec  *zstd.Decoder
+}
+
+func (r *zstdReader) Read(p []byte) (n int, err error) {
+	return r.dec.Read(p)
+}
+
+func (r *zstdReader) Close() error {
+	r.dec.Close()
+	return r.body.Close()
+}
+
 type closeFuncReadCloser struct {
 	reader    io.Reader
 	closeFunc func() error
@@ -102,7 +117,7 @@ func Reader(ctx context.Context, repo distri.Repo, fn string, cache bool) (io.Re
 	}
 	// good for typical links (≤ gigabit)
 	// performance bottleneck for faster links (10 gbit/s+)
-	req.Header.Set("Accept-Encoding", "gzip")
+	req.Header.Set("Accept-Encoding", "zstd, gzip")
 	resp, err := httpClient.Do(req.WithContext(ctx))
 	if err != nil {
 		return nil, err
@@ -123,7 +138,14 @@ func Reader(ctx context.Context, repo distri.Repo, fn string, cache bool) (io.Re
 			return nil, err
 		}
 		rdc = &gzipReader{body: resp.Body, zr: rd}
+	} else if strings.EqualFold(resp.Header.Get("Content-Encoding"), "zstd") {
+		dec, err := zstd.NewReader(resp.Body)
+		if err != nil {
+			return nil, err
+		}
+		rdc = &zstdReader{body: resp.Body, dec: dec}
 	}
+
 	var cacheFile *os.File
 	if cacheFn != "" {
 		cacheFile, err = os.Create(cacheFn)