Source file src/pkg/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	}
   149	

View as plain text