github.com/iDigitalFlame/xmt@v0.5.4/c2/wrapper/compress.go (about)

     1  // Copyright (C) 2020 - 2023 iDigitalFlame
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU General Public License as published by
     5  // the Free Software Foundation, either version 3 of the License, or
     6  // any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU General Public License
    14  // along with this program.  If not, see <https://www.gnu.org/licenses/>.
    15  //
    16  
    17  package wrapper
    18  
    19  import (
    20  	"compress/gzip"
    21  	"compress/zlib"
    22  	"io"
    23  	"sync"
    24  )
    25  
    26  const compLevel = zlib.BestSpeed
    27  
    28  const (
    29  	// Zlib is the default Zlib Wrapper. This wrapper uses the default compression level. Use the 'NewZlib'
    30  	// function to create a wrapper with a different level.
    31  	Zlib = compress(0x0)
    32  	// Gzip is the default Gzip Wrapper. This wrapper uses the default compression level. Use the 'NewGzip'
    33  	// function to create a wrapper with a different level.
    34  	Gzip = compress(0x1)
    35  )
    36  
    37  var (
    38  	zlibWriterPool = sync.Pool{
    39  		New: func() interface{} {
    40  			w, _ := zlib.NewWriterLevel(nil, compLevel)
    41  			return w
    42  		},
    43  	}
    44  	gzipWriterPool = sync.Pool{
    45  		New: func() interface{} {
    46  			w, _ := gzip.NewWriterLevel(nil, compLevel)
    47  			return w
    48  		},
    49  	}
    50  	zlibReaderPool, gzipReaderPool sync.Pool
    51  )
    52  
    53  type compress uint8
    54  type reader struct {
    55  	_ [0]func()
    56  	p *sync.Pool
    57  	io.ReadCloser
    58  }
    59  type writer struct {
    60  	_ [0]func()
    61  	p *sync.Pool
    62  	io.WriteCloser
    63  }
    64  
    65  func (r *reader) Close() error {
    66  	if r.ReadCloser == nil {
    67  		return nil
    68  	}
    69  	err := r.ReadCloser.Close()
    70  	r.p.Put(r.ReadCloser)
    71  	r.ReadCloser, r.p = nil, nil
    72  	return err
    73  }
    74  func (w *writer) Close() error {
    75  	if w.WriteCloser == nil {
    76  		return nil
    77  	}
    78  	err := w.WriteCloser.Close()
    79  	w.p.Put(w.WriteCloser)
    80  	w.WriteCloser, w.p = nil, nil
    81  	return err
    82  }
    83  func (c compress) Unwrap(r io.Reader) (io.Reader, error) {
    84  	switch c {
    85  	case Zlib:
    86  		c := zlibReaderPool.Get()
    87  		if c == nil {
    88  			n, err := zlib.NewReader(r)
    89  			if err != nil {
    90  				return nil, err
    91  			}
    92  			return &reader{ReadCloser: n, p: &zlibReaderPool}, nil
    93  		}
    94  		var ()
    95  		if err := c.(zlib.Resetter).Reset(r, nil); err != nil {
    96  			zlibReaderPool.Put(c)
    97  			return nil, err
    98  		}
    99  		return &reader{ReadCloser: c.(io.ReadCloser), p: &zlibReaderPool}, nil
   100  	case Gzip:
   101  		c := gzipReaderPool.Get()
   102  		if c == nil {
   103  			n, err := gzip.NewReader(r)
   104  			if err != nil {
   105  				return nil, err
   106  			}
   107  			return &reader{ReadCloser: n, p: &gzipReaderPool}, nil
   108  		}
   109  		var (
   110  			n   = c.(*gzip.Reader)
   111  			err = n.Reset(r)
   112  		)
   113  		if err == nil {
   114  			return &reader{ReadCloser: n, p: &gzipReaderPool}, nil
   115  		}
   116  		gzipReaderPool.Put(c)
   117  		return nil, err
   118  	}
   119  	return r, nil
   120  }
   121  func (c compress) Wrap(w io.WriteCloser) (io.WriteCloser, error) {
   122  	switch c {
   123  	case Zlib:
   124  		c := zlibWriterPool.Get().(*zlib.Writer)
   125  		c.Reset(w)
   126  		return &writer{WriteCloser: c, p: &zlibWriterPool}, nil
   127  	case Gzip:
   128  		c := gzipWriterPool.Get().(*gzip.Writer)
   129  		c.Reset(w)
   130  		return &writer{WriteCloser: c, p: &gzipWriterPool}, nil
   131  	}
   132  	return w, nil
   133  }