1/*
2 *
3 * Copyright 2017 gRPC authors.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 *
17 */
18
19// Package gzip implements and registers the gzip compressor
20// during the initialization.
21// This package is EXPERIMENTAL.
22package gzip
23
24import (
25	"compress/gzip"
26	"fmt"
27	"io"
28	"io/ioutil"
29	"sync"
30
31	"google.golang.org/grpc/encoding"
32)
33
34// Name is the name registered for the gzip compressor.
35const Name = "gzip"
36
37func init() {
38	c := &compressor{}
39	c.poolCompressor.New = func() interface{} {
40		return &writer{Writer: gzip.NewWriter(ioutil.Discard), pool: &c.poolCompressor}
41	}
42	encoding.RegisterCompressor(c)
43}
44
45type writer struct {
46	*gzip.Writer
47	pool *sync.Pool
48}
49
50// SetLevel updates the registered gzip compressor to use the compression level specified (gzip.HuffmanOnly is not supported).
51// NOTE: this function must only be called during initialization time (i.e. in an init() function),
52// and is not thread-safe.
53//
54// The error returned will be nil if the specified level is valid.
55func SetLevel(level int) error {
56	if level < gzip.DefaultCompression || level > gzip.BestCompression {
57		return fmt.Errorf("grpc: invalid gzip compression level: %d", level)
58	}
59	c := encoding.GetCompressor(Name).(*compressor)
60	c.poolCompressor.New = func() interface{} {
61		w, err := gzip.NewWriterLevel(ioutil.Discard, level)
62		if err != nil {
63			panic(err)
64		}
65		return &writer{Writer: w, pool: &c.poolCompressor}
66	}
67	return nil
68}
69
70func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
71	z := c.poolCompressor.Get().(*writer)
72	z.Writer.Reset(w)
73	return z, nil
74}
75
76func (z *writer) Close() error {
77	defer z.pool.Put(z)
78	return z.Writer.Close()
79}
80
81type reader struct {
82	*gzip.Reader
83	pool *sync.Pool
84}
85
86func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
87	z, inPool := c.poolDecompressor.Get().(*reader)
88	if !inPool {
89		newZ, err := gzip.NewReader(r)
90		if err != nil {
91			return nil, err
92		}
93		return &reader{Reader: newZ, pool: &c.poolDecompressor}, nil
94	}
95	if err := z.Reset(r); err != nil {
96		c.poolDecompressor.Put(z)
97		return nil, err
98	}
99	return z, nil
100}
101
102func (z *reader) Read(p []byte) (n int, err error) {
103	n, err = z.Reader.Read(p)
104	if err == io.EOF {
105		z.pool.Put(z)
106	}
107	return n, err
108}
109
110func (c *compressor) Name() string {
111	return Name
112}
113
114type compressor struct {
115	poolCompressor   sync.Pool
116	poolDecompressor sync.Pool
117}
118