You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@thrift.apache.org by yu...@apache.org on 2025/05/28 17:55:44 UTC

(thrift) branch go-zlib-pool created (now f0cf77c9b)

This is an automated email from the ASF dual-hosted git repository.

yuxuan pushed a change to branch go-zlib-pool
in repository https://gitbox.apache.org/repos/asf/thrift.git


      at f0cf77c9b go: Add a zlib reader pool

This branch includes the following new commits:

     new f0cf77c9b go: Add a zlib reader pool

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.



(thrift) 01/01: go: Add a zlib reader pool

Posted by yu...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

yuxuan pushed a commit to branch go-zlib-pool
in repository https://gitbox.apache.org/repos/asf/thrift.git

commit f0cf77c9bef50b5effee232126cd6e1ffee3c078
Author: Yuxuan 'fishy' Wang <yu...@reddit.com>
AuthorDate: Wed May 28 10:54:04 2025 -0700

    go: Add a zlib reader pool
    
    We implemented a zlib writer pool for default level when implementing
    THeader, this change also add a zlib reader pool to help speed up things
    when zlib is used.
    
    Also make TZlibTransport to use the zlib writer pool when it's using the
    default compression level.
---
 lib/go/thrift/compact_protocol_test.go | 15 +++++--
 lib/go/thrift/header_transport.go      | 22 +---------
 lib/go/thrift/pool.go                  |  2 +-
 lib/go/thrift/zlib_pool.go             | 76 ++++++++++++++++++++++++++++++++++
 lib/go/thrift/zlib_transport.go        | 30 ++++++++++----
 5 files changed, 111 insertions(+), 34 deletions(-)

diff --git a/lib/go/thrift/compact_protocol_test.go b/lib/go/thrift/compact_protocol_test.go
index 65f77f2c4..0d9575992 100644
--- a/lib/go/thrift/compact_protocol_test.go
+++ b/lib/go/thrift/compact_protocol_test.go
@@ -33,9 +33,18 @@ func TestReadWriteCompactProtocol(t *testing.T) {
 		NewTFramedTransport(NewTMemoryBuffer()),
 	}
 
-	zlib0, _ := NewTZlibTransport(NewTMemoryBuffer(), 0)
-	zlib6, _ := NewTZlibTransport(NewTMemoryBuffer(), 6)
-	zlib9, _ := NewTZlibTransport(NewTFramedTransport(NewTMemoryBuffer()), 9)
+	newTZlibTransport := func(trans TTransport, level int) *TZlibTransport {
+		t.Helper()
+		zlibTrans, err := NewTZlibTransport(trans, level)
+		if err != nil {
+			t.Fatalf("NewTZlibTransport returned error: %v", err)
+		}
+		return zlibTrans
+	}
+
+	zlib0 := newTZlibTransport(NewTMemoryBuffer(), 0)
+	zlib6 := newTZlibTransport(NewTMemoryBuffer(), 6)
+	zlib9 := newTZlibTransport(NewTFramedTransport(NewTMemoryBuffer()), 9)
 	transports = append(transports, zlib0, zlib6, zlib9)
 
 	for _, trans := range transports {
diff --git a/lib/go/thrift/header_transport.go b/lib/go/thrift/header_transport.go
index d6d64160a..0e5c7ec1d 100644
--- a/lib/go/thrift/header_transport.go
+++ b/lib/go/thrift/header_transport.go
@@ -22,7 +22,6 @@ package thrift
 import (
 	"bufio"
 	"bytes"
-	"compress/zlib"
 	"context"
 	"encoding/binary"
 	"errors"
@@ -166,7 +165,7 @@ func (tr *TransformReader) AddTransform(id THeaderTransformID) error {
 	case TransformNone:
 		// no-op
 	case TransformZlib:
-		readCloser, err := zlib.NewReader(tr.Reader)
+		readCloser, err := newZlibReader(tr.Reader)
 		if err != nil {
 			return err
 		}
@@ -211,25 +210,6 @@ func (tw *TransformWriter) Close() error {
 	return nil
 }
 
-var zlibDefaultLevelWriterPool = newPool(
-	func() *zlib.Writer {
-		return zlib.NewWriter(nil)
-	},
-	nil,
-)
-
-type zlibPoolCloser struct {
-	writer *zlib.Writer
-}
-
-func (z *zlibPoolCloser) Close() error {
-	defer func() {
-		z.writer.Reset(nil)
-		zlibDefaultLevelWriterPool.put(&z.writer)
-	}()
-	return z.writer.Close()
-}
-
 // AddTransform adds a transform.
 func (tw *TransformWriter) AddTransform(id THeaderTransformID) error {
 	switch id {
diff --git a/lib/go/thrift/pool.go b/lib/go/thrift/pool.go
index 1d623d422..6912f3ea5 100644
--- a/lib/go/thrift/pool.go
+++ b/lib/go/thrift/pool.go
@@ -43,7 +43,7 @@ func newPool[T any](generate func() *T, reset func(*T)) *pool[T] {
 	}
 	return &pool[T]{
 		pool: sync.Pool{
-			New: func() interface{} {
+			New: func() any {
 				return generate()
 			},
 		},
diff --git a/lib/go/thrift/zlib_pool.go b/lib/go/thrift/zlib_pool.go
new file mode 100644
index 000000000..c38e56850
--- /dev/null
+++ b/lib/go/thrift/zlib_pool.go
@@ -0,0 +1,76 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one
+* or more contributor license agreements. See the NOTICE file
+* distributed with this work for additional information
+* regarding copyright ownership. The ASF licenses this file
+* to you 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 thrift
+
+import (
+	"compress/zlib"
+	"io"
+	"sync"
+)
+
+type zlibReader interface {
+	io.ReadCloser
+	zlib.Resetter
+}
+
+var zlibReaderPool sync.Pool
+
+func newZlibReader(r io.Reader) (io.ReadCloser, error) {
+	if reader, _ := zlibReaderPool.Get().(*wrappedZlibReader); reader != nil {
+		if err := reader.Reset(r, nil); err == nil {
+			return reader, nil
+		}
+	}
+	reader, err := zlib.NewReader(r)
+	if err != nil {
+		return nil, err
+	}
+	return &wrappedZlibReader{reader.(zlibReader)}, nil
+}
+
+type wrappedZlibReader struct {
+	zlibReader
+}
+
+func (wr *wrappedZlibReader) Close() error {
+	defer func() {
+		zlibReaderPool.Put(wr)
+	}()
+	return wr.zlibReader.Close()
+}
+
+var zlibDefaultLevelWriterPool = newPool(
+	func() *zlib.Writer {
+		return zlib.NewWriter(nil)
+	},
+	nil,
+)
+
+type zlibPoolCloser struct {
+	writer *zlib.Writer
+}
+
+func (z *zlibPoolCloser) Close() error {
+	defer func() {
+		z.writer.Reset(nil)
+		zlibDefaultLevelWriterPool.put(&z.writer)
+	}()
+	return z.writer.Close()
+}
diff --git a/lib/go/thrift/zlib_transport.go b/lib/go/thrift/zlib_transport.go
index cefe1f994..c3863696d 100644
--- a/lib/go/thrift/zlib_transport.go
+++ b/lib/go/thrift/zlib_transport.go
@@ -33,9 +33,10 @@ type TZlibTransportFactory struct {
 
 // TZlibTransport is a TTransport implementation that makes use of zlib compression.
 type TZlibTransport struct {
-	reader    io.ReadCloser
-	transport TTransport
-	writer    *zlib.Writer
+	reader      io.ReadCloser
+	transport   TTransport
+	writer      *zlib.Writer
+	writeCloser io.Closer
 }
 
 // GetTransport constructs a new instance of NewTZlibTransport
@@ -64,14 +65,25 @@ func NewTZlibTransportFactoryWithFactory(level int, factory TTransportFactory) *
 
 // NewTZlibTransport constructs a new instance of TZlibTransport
 func NewTZlibTransport(trans TTransport, level int) (*TZlibTransport, error) {
-	w, err := zlib.NewWriterLevel(trans, level)
-	if err != nil {
-		return nil, err
+	var w *zlib.Writer
+	var writeCloser io.Closer
+	if level == zlib.DefaultCompression {
+		w = zlibDefaultLevelWriterPool.get()
+		w.Reset(trans)
+		writeCloser = &zlibPoolCloser{writer: w}
+	} else {
+		writer, err := zlib.NewWriterLevel(trans, level)
+		if err != nil {
+			return nil, err
+		}
+		w = writer
+		writeCloser = writer
 	}
 
 	return &TZlibTransport{
-		writer:    w,
-		transport: trans,
+		writer:      w,
+		writeCloser: writeCloser,
+		transport:   trans,
 	}, nil
 }
 
@@ -109,7 +121,7 @@ func (z *TZlibTransport) Open() error {
 
 func (z *TZlibTransport) Read(p []byte) (int, error) {
 	if z.reader == nil {
-		r, err := zlib.NewReader(z.transport)
+		r, err := newZlibReader(z.transport)
 		if err != nil {
 			return 0, NewTTransportExceptionFromError(err)
 		}