libgo/go/archive/zip/register.go

   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         compressors   sync.Map // map[uint16]Compressor
 107         decompressors sync.Map // map[uint16]Decompressor
 108 )
 109
 110 func init() {
 111         compressors.Store(Store, Compressor(func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil }))
 112         compressors.Store(Deflate, Compressor(func(w io.Writer) (io.WriteCloser, error) { return newFlateWriter(w), nil }))
 113
 114         decompressors.Store(Store, Decompressor(ioutil.NopCloser))
 115         decompressors.Store(Deflate, Decompressor(newFlateReader))
 116 }
 117
 118 // RegisterDecompressor allows custom decompressors for a specified method ID.
 119 // The common methods Store and Deflate are built in.
 120 func RegisterDecompressor(method uint16, dcomp Decompressor) {
 121         if _, dup := decompressors.LoadOrStore(method, dcomp); dup {
 122                 panic("decompressor already registered")
 123         }
 124 }
 125
 126 // RegisterCompressor registers custom compressors for a specified method ID.
 127 // The common methods Store and Deflate are built in.
 128 func RegisterCompressor(method uint16, comp Compressor) {
 129         if _, dup := compressors.LoadOrStore(method, comp); dup {
 130                 panic("compressor already registered")
 131         }
 132 }
 133
 134 func compressor(method uint16) Compressor {
 135         ci, ok := compressors.Load(method)
 136         if !ok {
 137                 return nil
 138         }
 139         return ci.(Compressor)
 140 }
 141
 142 func decompressor(method uint16) Decompressor {
 143         di, ok := decompressors.Load(method)
 144         if !ok {
 145                 return nil
 146         }
 147         return di.(Decompressor)
 148 }