github.com/wfusion/gofusion@v1.1.14/common/utils/compress/encode.go (about)

     1  package compress
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/wfusion/gofusion/common/utils"
     7  )
     8  
     9  type encodable interface {
    10  	io.WriteCloser
    11  	Flush() error
    12  	Reset(w io.Writer)
    13  }
    14  
    15  func EncodeBytesFunc(algo Algorithm) func(src []byte) (dst []byte, err error) {
    16  	return func(src []byte) (dst []byte, err error) {
    17  		var (
    18  			enc      encodable
    19  			recycles []func()
    20  		)
    21  		defer func() {
    22  			for _, cb := range recycles {
    23  				if cb != nil {
    24  					cb()
    25  				}
    26  			}
    27  		}()
    28  
    29  		dstBuffer, bufferCb := utils.BytesBufferPool.Get(nil)
    30  		recycles = append(recycles, bufferCb)
    31  
    32  		_, err = utils.Catch(func() {
    33  			var cb func()
    34  			enc, cb = getEncoder(algo, dstBuffer)
    35  			recycles = append(recycles, cb)
    36  		})
    37  		if err != nil || enc == nil {
    38  			return
    39  		}
    40  
    41  		srcBuffer, bufferCb := utils.BytesBufferPool.Get(nil)
    42  		recycles = append(recycles, bufferCb)
    43  		srcBuffer.Write(src)
    44  		if _, err = io.Copy(enc, srcBuffer); err != nil {
    45  			if err = utils.ErrIgnore(err, eofErrs...); err != nil {
    46  				return
    47  			}
    48  		}
    49  		if err = enc.Flush(); err != nil {
    50  			if err = utils.ErrIgnore(err, eofErrs...); err != nil {
    51  				return
    52  			}
    53  		}
    54  		dst = make([]byte, dstBuffer.Len())
    55  		copy(dst, dstBuffer.Bytes())
    56  		return
    57  	}
    58  }
    59  
    60  func EncodeStreamFunc(algo Algorithm) func(dst io.Writer, src io.Reader) (err error) {
    61  	return func(dst io.Writer, src io.Reader) (err error) {
    62  		var (
    63  			enc encodable
    64  			cb  func()
    65  		)
    66  		if _, err = utils.Catch(func() { enc, cb = getEncoder(algo, dst) }); err != nil {
    67  			return
    68  		}
    69  		defer func() {
    70  			if cb != nil {
    71  				cb()
    72  			}
    73  		}()
    74  
    75  		if _, err = io.Copy(enc, src); err != nil {
    76  			if err = utils.ErrIgnore(err, eofErrs...); err != nil {
    77  				return
    78  			}
    79  		}
    80  		err = enc.Flush()
    81  		return utils.ErrIgnore(err, eofErrs...)
    82  	}
    83  }
    84  
    85  func getEncoder(algo Algorithm, dst io.Writer) (enc encodable, recycle func()) {
    86  	p, ok := encoderPools[algo]
    87  	if !algo.IsValid() || !ok {
    88  		panic(ErrUnknownAlgorithm)
    89  	}
    90  
    91  	sealer, recycle := p.Get(dst)
    92  	enc = sealer.encodable
    93  	return
    94  }
    95  
    96  type encoder struct {
    97  	enc encodable
    98  	w   io.Writer
    99  	cb  func()
   100  }
   101  
   102  func NewEncFunc(algo Algorithm) func(w io.Writer) encodable {
   103  	return func(w io.Writer) encodable {
   104  		enc, cb := getEncoder(algo, w)
   105  		if enc == nil {
   106  			panic(ErrUnknownAlgorithm)
   107  		}
   108  		return &encoder{
   109  			enc: enc,
   110  			w:   w,
   111  			cb:  cb,
   112  		}
   113  	}
   114  }
   115  
   116  func (e *encoder) Write(p []byte) (n int, err error) { return e.enc.Write(p) }
   117  func (e *encoder) Reset(w io.Writer)                 { e.enc.Reset(w) }
   118  func (e *encoder) Flush() (err error)                { defer utils.FlushAnyway(e.w); return e.enc.Flush() }
   119  func (e *encoder) Close() (err error) {
   120  	defer e.cb()
   121  	defer utils.CloseAnyway(e.w)
   122  	return e.Flush()
   123  }