statistics: faster gzip (#47103)

close #47104
pingcap · Sep 20, 2023 · d569d19 · d569d19
1 parent 6271d1c
commit d569d19
Show file tree

Hide file tree

Showing 9 changed files with 74 additions and 31 deletions.
diff --git a/DEPS.bzl b/DEPS.bzl
@@ -4363,13 +4363,13 @@ def go_deps():
         name = "com_github_klauspost_compress",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/klauspost/compress",
-        sha256 = "7e004bb6b71535508bfa9c57256cfb2ca23f09ea281dbecafea217796b712fcd",
-        strip_prefix = "github.com/klauspost/compress@v1.16.5",
+        sha256 = "fa94794543608ad4f600c67994a317173b4e72c1159b8a84ab46a846c7643587",
+        strip_prefix = "github.com/klauspost/compress@v1.17.0",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/klauspost/compress/com_github_klauspost_compress-v1.16.5.zip",
-            "http://ats.apps.svc/gomod/github.com/klauspost/compress/com_github_klauspost_compress-v1.16.5.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/klauspost/compress/com_github_klauspost_compress-v1.16.5.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/klauspost/compress/com_github_klauspost_compress-v1.16.5.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/klauspost/compress/com_github_klauspost_compress-v1.17.0.zip",
+            "http://ats.apps.svc/gomod/github.com/klauspost/compress/com_github_klauspost_compress-v1.17.0.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/klauspost/compress/com_github_klauspost_compress-v1.17.0.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/klauspost/compress/com_github_klauspost_compress-v1.17.0.zip",
         ],
     )
     go_repository(

diff --git a/br/pkg/lightning/backend/local/BUILD.bazel b/br/pkg/lightning/backend/local/BUILD.bazel
@@ -55,6 +55,7 @@ go_library(
         "//tablecodec",
         "//types",
         "//util/codec",
+        "//util/compress",
         "//util/engine",
         "//util/hack",
         "//util/mathutil",

diff --git a/br/pkg/lightning/backend/local/compress.go b/br/pkg/lightning/backend/local/compress.go
@@ -16,9 +16,9 @@ package local
 
 import (
 	"io"
-	"sync"
 
 	"github.com/klauspost/compress/gzip" // faster than stdlib
+	"github.com/pingcap/tidb/util/compress"
 	"google.golang.org/grpc"
 )
 
@@ -29,15 +29,9 @@ var (
 
 type gzipCompressor struct{}
 
-var gzipWriterPool = sync.Pool{
-	New: func() any {
-		return gzip.NewWriter(io.Discard)
-	},
-}
-
 func (*gzipCompressor) Do(w io.Writer, p []byte) error {
-	z := gzipWriterPool.Get().(*gzip.Writer)
-	defer gzipWriterPool.Put(z)
+	z := compress.GzipWriterPool.Get().(*gzip.Writer)
+	defer compress.GzipWriterPool.Put(z)
 	z.Reset(w)
 	if _, err := z.Write(p); err != nil {
 		return err
@@ -51,22 +45,16 @@ func (*gzipCompressor) Type() string {
 
 type gzipDecompressor struct{}
 
-var gzipReaderPool = sync.Pool{
-	New: func() any {
-		return &gzip.Reader{}
-	},
-}
-
 func (*gzipDecompressor) Do(r io.Reader) ([]byte, error) {
-	z := gzipReaderPool.Get().(*gzip.Reader)
+	z := compress.GzipReaderPool.Get().(*gzip.Reader)
 	if err := z.Reset(r); err != nil {
-		gzipReaderPool.Put(z)
+		compress.GzipReaderPool.Put(z)
 		return nil, err
 	}
 
 	defer func() {
 		_ = z.Close()
-		gzipReaderPool.Put(z)
+		compress.GzipReaderPool.Put(z)
 	}()
 	return io.ReadAll(z)
 }

diff --git a/go.mod b/go.mod
@@ -64,7 +64,7 @@ require (
 	github.com/johannesboyne/gofakes3 v0.0.0-20230506070712-04da935ef877
 	github.com/joho/sqltocsv v0.0.0-20210428211105-a6d6801d59df
 	github.com/kisielk/errcheck v1.6.3
-	github.com/klauspost/compress v1.16.5
+	github.com/klauspost/compress v1.17.0
 	github.com/kyoh86/exportloopref v0.1.11
 	github.com/lestrrat-go/jwx/v2 v2.0.11
 	github.com/mgechev/revive v1.3.2

diff --git a/go.sum b/go.sum
@@ -645,8 +645,8 @@ github.com/klauspost/compress v1.9.7/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0
 github.com/klauspost/compress v1.10.5/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 github.com/klauspost/compress v1.11.7/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 github.com/klauspost/compress v1.12.2/go.mod h1:8dP1Hq4DHOhN9w426knH3Rhby4rFm6D8eO+e+Dq5Gzg=
-github.com/klauspost/compress v1.16.5 h1:IFV2oUNUzZaz+XyusxpLzpzS8Pt5rh0Z16For/djlyI=
-github.com/klauspost/compress v1.16.5/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
+github.com/klauspost/compress v1.17.0 h1:Rnbp4K9EjcDuVuHtd0dgA4qNuv9yKDYKK1ulpJwgrqM=
+github.com/klauspost/compress v1.17.0/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
 github.com/klauspost/cpuid v1.2.1/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
 github.com/klauspost/cpuid v1.3.1 h1:5JNjFYYQrZeKRJ0734q51WCEEn2huer72Dc7K+R/b6s=
 github.com/klauspost/cpuid v1.3.1/go.mod h1:bYW4mA6ZgKPob1/Dlai2LviZJO7KGI3uoWLd42rAQw4=

diff --git a/statistics/handle/BUILD.bazel b/statistics/handle/BUILD.bazel
@@ -38,10 +38,12 @@ go_library(
         "//types",
         "//util",
         "//util/chunk",
+        "//util/compress",
         "//util/logutil",
         "//util/mathutil",
         "//util/sqlexec",
         "//util/timeutil",
+        "@com_github_klauspost_compress//gzip",
         "@com_github_ngaut_pools//:pools",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",

diff --git a/statistics/handle/dump.go b/statistics/handle/dump.go
@@ -16,12 +16,12 @@ package handle
 
 import (
 	"bytes"
-	"compress/gzip"
 	"encoding/json"
 	"fmt"
 	"io"
 	"time"
 
+	"github.com/klauspost/compress/gzip"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/parser/model"
@@ -33,6 +33,7 @@ import (
 	"github.com/pingcap/tidb/statistics/handle/globalstats"
 	handle_metrics "github.com/pingcap/tidb/statistics/handle/metrics"
 	"github.com/pingcap/tidb/types"
+	compressutil "github.com/pingcap/tidb/util/compress"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/sqlexec"
 	"github.com/pingcap/tipb/go-tipb"
@@ -539,7 +540,9 @@ func JSONTableToBlocks(jsTable *JSONTable, blockSize int) ([][]byte, error) {
 		return nil, errors.Trace(err)
 	}
 	var gzippedData bytes.Buffer
-	gzipWriter := gzip.NewWriter(&gzippedData)
+	gzipWriter := compressutil.GzipWriterPool.Get().(*gzip.Writer)
+	defer compressutil.GzipWriterPool.Put(gzipWriter)
+	gzipWriter.Reset(&gzippedData)
 	if _, err := gzipWriter.Write(data); err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -568,10 +571,14 @@ func BlocksToJSONTable(blocks [][]byte) (*JSONTable, error) {
 		data = append(data, blocks[i]...)
 	}
 	gzippedData := bytes.NewReader(data)
-	gzipReader, err := gzip.NewReader(gzippedData)
-	if err != nil {
+	gzipReader := compressutil.GzipReaderPool.Get().(*gzip.Reader)
+	if err := gzipReader.Reset(gzippedData); err != nil {
+		compressutil.GzipReaderPool.Put(gzipReader)
 		return nil, err
 	}
+	defer func() {
+		compressutil.GzipReaderPool.Put(gzipReader)
+	}()
 	if err := gzipReader.Close(); err != nil {
 		return nil, err
 	}

diff --git a/util/compress/BUILD.bazel b/util/compress/BUILD.bazel
@@ -0,0 +1,9 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "compress",
+    srcs = ["gzip.go"],
+    importpath = "github.com/pingcap/tidb/util/compress",
+    visibility = ["//visibility:public"],
+    deps = ["@com_github_klauspost_compress//gzip"],
+)
diff --git a/util/compress/gzip.go b/util/compress/gzip.go
@@ -0,0 +1,36 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// 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 compress
+
+import (
+	"io"
+	"sync"
+
+	"github.com/klauspost/compress/gzip"
+)
+
+// GzipWriterPool is a sync.Pool of gzip.Writer.
+var GzipWriterPool = sync.Pool{
+	New: func() any {
+		return gzip.NewWriter(io.Discard)
+	},
+}
+
+// GzipReaderPool is a sync.Pool of gzip.Reader.
+var GzipReaderPool = sync.Pool{
+	New: func() any {
+		return &gzip.Reader{}
+	},
+}