github.com/muesli/go@v0.0.0-20170208044820-e410d2a81ef2/src/archive/zip/register.go (about)

     1  // Copyright 2010 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 zip
     6  
     7  import (
     8  	"compress/flate"
     9  	"errors"
    10  	"io"
    11  	"io/ioutil"
    12  	"sync"
    13  )
    14  
    15  // A Compressor returns a new compressing writer, writing to w.
    16  // The WriteCloser's Close method must be used to flush pending data to w.
    17  // The Compressor itself must be safe to invoke from multiple goroutines
    18  // simultaneously, but each returned writer will be used only by
    19  // one goroutine at a time.
    20  type Compressor func(w io.Writer) (io.WriteCloser, error)
    21  
    22  // A Decompressor returns a new decompressing reader, reading from r.
    23  // The ReadCloser's Close method must be used to release associated resources.
    24  // The Decompressor itself must be safe to invoke from multiple goroutines
    25  // simultaneously, but each returned reader will be used only by
    26  // one goroutine at a time.
    27  type Decompressor func(r io.Reader) io.ReadCloser
    28  
    29  var flateWriterPool sync.Pool
    30  
    31  func newFlateWriter(w io.Writer) io.WriteCloser {
    32  	fw, ok := flateWriterPool.Get().(*flate.Writer)
    33  	if ok {
    34  		fw.Reset(w)
    35  	} else {
    36  		fw, _ = flate.NewWriter(w, 5)
    37  	}
    38  	return &pooledFlateWriter{fw: fw}
    39  }
    40  
    41  type pooledFlateWriter struct {
    42  	mu sync.Mutex // guards Close and Write
    43  	fw *flate.Writer
    44  }
    45  
    46  func (w *pooledFlateWriter) Write(p []byte) (n int, err error) {
    47  	w.mu.Lock()
    48  	defer w.mu.Unlock()
    49  	if w.fw == nil {
    50  		return 0, errors.New("Write after Close")
    51  	}
    52  	return w.fw.Write(p)
    53  }
    54  
    55  func (w *pooledFlateWriter) Close() error {
    56  	w.mu.Lock()
    57  	defer w.mu.Unlock()
    58  	var err error
    59  	if w.fw != nil {
    60  		err = w.fw.Close()
    61  		flateWriterPool.Put(w.fw)
    62  		w.fw = nil
    63  	}
    64  	return err
    65  }
    66  
    67  var flateReaderPool sync.Pool
    68  
    69  func newFlateReader(r io.Reader) io.ReadCloser {
    70  	fr, ok := flateReaderPool.Get().(io.ReadCloser)
    71  	if ok {
    72  		fr.(flate.Resetter).Reset(r, nil)
    73  	} else {
    74  		fr = flate.NewReader(r)
    75  	}
    76  	return &pooledFlateReader{fr: fr}
    77  }
    78  
    79  type pooledFlateReader struct {
    80  	mu sync.Mutex // guards Close and Read
    81  	fr io.ReadCloser
    82  }
    83  
    84  func (r *pooledFlateReader) Read(p []byte) (n int, err error) {
    85  	r.mu.Lock()
    86  	defer r.mu.Unlock()
    87  	if r.fr == nil {
    88  		return 0, errors.New("Read after Close")
    89  	}
    90  	return r.fr.Read(p)
    91  }
    92  
    93  func (r *pooledFlateReader) Close() error {
    94  	r.mu.Lock()
    95  	defer r.mu.Unlock()
    96  	var err error
    97  	if r.fr != nil {
    98  		err = r.fr.Close()
    99  		flateReaderPool.Put(r.fr)
   100  		r.fr = nil
   101  	}
   102  	return err
   103  }
   104  
   105  var (
   106  	mu sync.RWMutex // guards compressor and decompressor maps
   107  
   108  	compressors = map[uint16]Compressor{
   109  		Store:   func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil },
   110  		Deflate: func(w io.Writer) (io.WriteCloser, error) { return newFlateWriter(w), nil },
   111  	}
   112  
   113  	decompressors = map[uint16]Decompressor{
   114  		Store:   ioutil.NopCloser,
   115  		Deflate: newFlateReader,
   116  	}
   117  )
   118  
   119  // RegisterDecompressor allows custom decompressors for a specified method ID.
   120  // The common methods Store and Deflate are built in.
   121  func RegisterDecompressor(method uint16, dcomp Decompressor) {
   122  	mu.Lock()
   123  	defer mu.Unlock()
   124  
   125  	if _, ok := decompressors[method]; ok {
   126  		panic("decompressor already registered")
   127  	}
   128  	decompressors[method] = dcomp
   129  }
   130  
   131  // RegisterCompressor registers custom compressors for a specified method ID.
   132  // The common methods Store and Deflate are built in.
   133  func RegisterCompressor(method uint16, comp Compressor) {
   134  	mu.Lock()
   135  	defer mu.Unlock()
   136  
   137  	if _, ok := compressors[method]; ok {
   138  		panic("compressor already registered")
   139  	}
   140  	compressors[method] = comp
   141  }
   142  
   143  func compressor(method uint16) Compressor {
   144  	mu.RLock()
   145  	defer mu.RUnlock()
   146  	return compressors[method]
   147  }
   148  
   149  func decompressor(method uint16) Decompressor {
   150  	mu.RLock()
   151  	defer mu.RUnlock()
   152  	return decompressors[method]
   153  }