zip.go

Documentation: github.com/klauspost/compress/zstd

     1  // Copyright 2019+ Klaus Post. All rights reserved.
     2  // License information can be found in the LICENSE file.
     3  
     4  package zstd
     5  
     6  import (
     7  	"errors"
     8  	"io"
     9  	"sync"
    10  )
    11  
    12  // ZipMethodWinZip is the method for Zstandard compressed data inside Zip files for WinZip.
    13  // See https://www.winzip.com/win/en/comp_info.html
    14  const ZipMethodWinZip = 93
    15  
    16  // ZipMethodPKWare is the original method number used by PKWARE to indicate Zstandard compression.
    17  // Deprecated: This has been deprecated by PKWARE, use ZipMethodWinZip instead for compression.
    18  // See https://pkware.cachefly.net/webdocs/APPNOTE/APPNOTE-6.3.9.TXT
    19  const ZipMethodPKWare = 20
    20  
    21  // zipReaderPool is the default reader pool.
    22  var zipReaderPool = sync.Pool{New: func() interface{} {
    23  	z, err := NewReader(nil, WithDecoderLowmem(true), WithDecoderMaxWindow(128<<20), WithDecoderConcurrency(1))
    24  	if err != nil {
    25  		panic(err)
    26  	}
    27  	return z
    28  }}
    29  
    30  // newZipReader creates a pooled zip decompressor.
    31  func newZipReader(opts ...DOption) func(r io.Reader) io.ReadCloser {
    32  	pool := &zipReaderPool
    33  	if len(opts) > 0 {
    34  		opts = append([]DOption{WithDecoderLowmem(true), WithDecoderMaxWindow(128 << 20)}, opts...)
    35  		// Force concurrency 1
    36  		opts = append(opts, WithDecoderConcurrency(1))
    37  		// Create our own pool
    38  		pool = &sync.Pool{}
    39  	}
    40  	return func(r io.Reader) io.ReadCloser {
    41  		dec, ok := pool.Get().(*Decoder)
    42  		if ok {
    43  			dec.Reset(r)
    44  		} else {
    45  			d, err := NewReader(r, opts...)
    46  			if err != nil {
    47  				panic(err)
    48  			}
    49  			dec = d
    50  		}
    51  		return &pooledZipReader{dec: dec, pool: pool}
    52  	}
    53  }
    54  
    55  type pooledZipReader struct {
    56  	mu   sync.Mutex // guards Close and Read
    57  	pool *sync.Pool
    58  	dec  *Decoder
    59  }
    60  
    61  func (r *pooledZipReader) Read(p []byte) (n int, err error) {
    62  	r.mu.Lock()
    63  	defer r.mu.Unlock()
    64  	if r.dec == nil {
    65  		return 0, errors.New("read after close or EOF")
    66  	}
    67  	dec, err := r.dec.Read(p)
    68  	if err == io.EOF {
    69  		r.dec.Reset(nil)
    70  		r.pool.Put(r.dec)
    71  		r.dec = nil
    72  	}
    73  	return dec, err
    74  }
    75  
    76  func (r *pooledZipReader) Close() error {
    77  	r.mu.Lock()
    78  	defer r.mu.Unlock()
    79  	var err error
    80  	if r.dec != nil {
    81  		err = r.dec.Reset(nil)
    82  		r.pool.Put(r.dec)
    83  		r.dec = nil
    84  	}
    85  	return err
    86  }
    87  
    88  type pooledZipWriter struct {
    89  	mu   sync.Mutex // guards Close and Read
    90  	enc  *Encoder
    91  	pool *sync.Pool
    92  }
    93  
    94  func (w *pooledZipWriter) Write(p []byte) (n int, err error) {
    95  	w.mu.Lock()
    96  	defer w.mu.Unlock()
    97  	if w.enc == nil {
    98  		return 0, errors.New("Write after Close")
    99  	}
   100  	return w.enc.Write(p)
   101  }
   102  
   103  func (w *pooledZipWriter) Close() error {
   104  	w.mu.Lock()
   105  	defer w.mu.Unlock()
   106  	var err error
   107  	if w.enc != nil {
   108  		err = w.enc.Close()
   109  		w.pool.Put(w.enc)
   110  		w.enc = nil
   111  	}
   112  	return err
   113  }
   114  
   115  // ZipCompressor returns a compressor that can be registered with zip libraries.
   116  // The provided encoder options will be used on all encodes.
   117  func ZipCompressor(opts ...EOption) func(w io.Writer) (io.WriteCloser, error) {
   118  	var pool sync.Pool
   119  	return func(w io.Writer) (io.WriteCloser, error) {
   120  		enc, ok := pool.Get().(*Encoder)
   121  		if ok {
   122  			enc.Reset(w)
   123  		} else {
   124  			var err error
   125  			enc, err = NewWriter(w, opts...)
   126  			if err != nil {
   127  				return nil, err
   128  			}
   129  		}
   130  		return &pooledZipWriter{enc: enc, pool: &pool}, nil
   131  	}
   132  }
   133  
   134  // ZipDecompressor returns a decompressor that can be registered with zip libraries.
   135  // See ZipCompressor for example.
   136  // Options can be specified. WithDecoderConcurrency(1) is forced,
   137  // and by default a 128MB maximum decompression window is specified.
   138  // The window size can be overridden if required.
   139  func ZipDecompressor(opts ...DOption) func(r io.Reader) io.ReadCloser {
   140  	return newZipReader(opts...)
   141  }
   142
View as plain text