gzip.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package httpcommon
     6  
     7  import (
     8  	"compress/flate"
     9  	"compress/gzip"
    10  	"errors"
    11  	"io"
    12  	"io/fs"
    13  	"sync"
    14  )
    15  
    16  var errConcurrentRead = errors.New("http: concurrent read on response body")
    17  
    18  // incomparable is a zero-width, non-comparable type. Adding it to a struct
    19  // makes that struct also non-comparable, and generally doesn't add
    20  // any size (as long as it's first).
    21  type incomparable [0]func()
    22  
    23  // GzipReader wraps a response body so it can lazily
    24  // get gzip.Reader from the pool on the first call to Read.
    25  // After Close is called it puts gzip.Reader to the pool immediately
    26  // if there is no Read in progress or later when Read completes.
    27  type GzipReader struct {
    28  	_    incomparable
    29  	Body io.ReadCloser // underlying Response.Body
    30  	mu   sync.Mutex    // guards zr and zerr
    31  	zr   *gzip.Reader  // stores gzip reader from the pool between reads
    32  	zerr error         // sticky gzip reader init error or sentinel value to detect concurrent read and read after close
    33  }
    34  
    35  type eofReader struct{}
    36  
    37  func (eofReader) Read([]byte) (int, error) { return 0, io.EOF }
    38  func (eofReader) ReadByte() (byte, error)  { return 0, io.EOF }
    39  
    40  var gzipPool = sync.Pool{New: func() any { return new(gzip.Reader) }}
    41  
    42  // gzipPoolGet gets a gzip.Reader from the pool and resets it to read from r.
    43  func gzipPoolGet(r io.Reader) (*gzip.Reader, error) {
    44  	zr := gzipPool.Get().(*gzip.Reader)
    45  	if err := zr.Reset(r); err != nil {
    46  		gzipPoolPut(zr)
    47  		return nil, err
    48  	}
    49  	return zr, nil
    50  }
    51  
    52  // gzipPoolPut puts a gzip.Reader back into the pool.
    53  func gzipPoolPut(zr *gzip.Reader) {
    54  	// Reset will allocate bufio.Reader if we pass it anything
    55  	// other than a flate.Reader, so ensure that it's getting one.
    56  	var r flate.Reader = eofReader{}
    57  	zr.Reset(r)
    58  	gzipPool.Put(zr)
    59  }
    60  
    61  // acquire returns a gzip.Reader for reading response body.
    62  // The reader must be released after use.
    63  func (gz *GzipReader) acquire() (*gzip.Reader, error) {
    64  	gz.mu.Lock()
    65  	defer gz.mu.Unlock()
    66  	if gz.zerr != nil {
    67  		return nil, gz.zerr
    68  	}
    69  	if gz.zr == nil {
    70  		// gzipPoolGet might block indefinitely since it reads the gzip header.
    71  		// Therefore, drop mu temporarily when using gzipPoolGet.
    72  		// We set zerr to errConcurrentRead to prevent concurrent read
    73  		// even when mu is temporarily dropped.
    74  		gz.zerr = errConcurrentRead
    75  		gz.mu.Unlock()
    76  		zr, err := gzipPoolGet(gz.Body)
    77  		gz.mu.Lock()
    78  		// Guard against Close being called while gzipPoolGet is running.
    79  		if gz.zerr != errConcurrentRead {
    80  			if zr != nil {
    81  				gzipPoolPut(zr)
    82  			}
    83  			return nil, gz.zerr
    84  		}
    85  		gz.zr, gz.zerr = zr, err
    86  		if gz.zerr != nil {
    87  			return nil, gz.zerr
    88  		}
    89  	}
    90  	ret := gz.zr
    91  	gz.zr, gz.zerr = nil, errConcurrentRead
    92  	return ret, nil
    93  }
    94  
    95  // release returns the gzip.Reader to the pool if Close was called during Read.
    96  func (gz *GzipReader) release(zr *gzip.Reader) {
    97  	gz.mu.Lock()
    98  	defer gz.mu.Unlock()
    99  	if gz.zerr == errConcurrentRead {
   100  		gz.zr, gz.zerr = zr, nil
   101  	} else { // fs.ErrClosed
   102  		gzipPoolPut(zr)
   103  	}
   104  }
   105  
   106  // close returns the gzip.Reader to the pool immediately or
   107  // signals release to do so after Read completes.
   108  func (gz *GzipReader) close() {
   109  	gz.mu.Lock()
   110  	defer gz.mu.Unlock()
   111  	if gz.zerr == nil && gz.zr != nil {
   112  		gzipPoolPut(gz.zr)
   113  		gz.zr = nil
   114  	}
   115  	gz.zerr = fs.ErrClosed
   116  }
   117  
   118  func (gz *GzipReader) Read(p []byte) (n int, err error) {
   119  	zr, err := gz.acquire()
   120  	if err != nil {
   121  		return 0, err
   122  	}
   123  	defer gz.release(zr)
   124  
   125  	return zr.Read(p)
   126  }
   127  
   128  func (gz *GzipReader) Close() error {
   129  	gz.close()
   130  
   131  	return gz.Body.Close()
   132  }
   133
View as plain text