github.com/miolini/go@v0.0.0-20160405192216-fca68c8cb408/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 (
    68  	mu sync.RWMutex // guards compressor and decompressor maps
    69  
    70  	compressors = map[uint16]Compressor{
    71  		Store:   func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil },
    72  		Deflate: func(w io.Writer) (io.WriteCloser, error) { return newFlateWriter(w), nil },
    73  	}
    74  
    75  	decompressors = map[uint16]Decompressor{
    76  		Store:   ioutil.NopCloser,
    77  		Deflate: flate.NewReader,
    78  	}
    79  )
    80  
    81  // RegisterDecompressor allows custom decompressors for a specified method ID.
    82  // The common methods Store and Deflate are built in.
    83  func RegisterDecompressor(method uint16, dcomp Decompressor) {
    84  	mu.Lock()
    85  	defer mu.Unlock()
    86  
    87  	if _, ok := decompressors[method]; ok {
    88  		panic("decompressor already registered")
    89  	}
    90  	decompressors[method] = dcomp
    91  }
    92  
    93  // RegisterCompressor registers custom compressors for a specified method ID.
    94  // The common methods Store and Deflate are built in.
    95  func RegisterCompressor(method uint16, comp Compressor) {
    96  	mu.Lock()
    97  	defer mu.Unlock()
    98  
    99  	if _, ok := compressors[method]; ok {
   100  		panic("compressor already registered")
   101  	}
   102  	compressors[method] = comp
   103  }
   104  
   105  func compressor(method uint16) Compressor {
   106  	mu.RLock()
   107  	defer mu.RUnlock()
   108  	return compressors[method]
   109  }
   110  
   111  func decompressor(method uint16) Decompressor {
   112  	mu.RLock()
   113  	defer mu.RUnlock()
   114  	return decompressors[method]
   115  }