github.com/google/cloudprober@v0.11.3/surfacers/common/compress/compress.go (about)

     1  // Copyright 2017-2018 The Cloudprober Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package compress implements compression related utilities.
    16  package compress
    17  
    18  import (
    19  	"bytes"
    20  	"compress/gzip"
    21  	"context"
    22  	"encoding/base64"
    23  	"sync"
    24  	"time"
    25  
    26  	"github.com/google/cloudprober/logger"
    27  )
    28  
    29  var (
    30  	compressionBufferFlushInterval = time.Second
    31  )
    32  
    33  // CompressionBuffer stores the data that is ready to be compressed.
    34  type CompressionBuffer struct {
    35  	mu        sync.Mutex
    36  	buf       *bytes.Buffer
    37  	lines     int
    38  	batchSize int
    39  	l         *logger.Logger
    40  
    41  	// Function to callback with compressed data.
    42  	callback func([]byte)
    43  
    44  	// Cancel function for the internal context, to stop the compression and
    45  	// flushing loop.
    46  	cancelCtx context.CancelFunc
    47  }
    48  
    49  // NewCompressionBuffer returns a new compression buffer.
    50  func NewCompressionBuffer(inctx context.Context, callback func([]byte), batchSize int, l *logger.Logger) *CompressionBuffer {
    51  	ctx, cancel := context.WithCancel(inctx)
    52  	c := &CompressionBuffer{
    53  		buf:       new(bytes.Buffer),
    54  		callback:  callback,
    55  		batchSize: batchSize,
    56  		l:         l,
    57  		cancelCtx: cancel,
    58  	}
    59  
    60  	// Start the flush loop: call flush every sec, until ctx.Done().
    61  	go func() {
    62  		ticker := time.NewTicker(compressionBufferFlushInterval)
    63  		defer ticker.Stop()
    64  
    65  		for range ticker.C {
    66  			select {
    67  			case <-ctx.Done():
    68  				return
    69  			default:
    70  			}
    71  
    72  			c.compressAndCallback()
    73  		}
    74  	}()
    75  
    76  	return c
    77  }
    78  
    79  // WriteLineToBuffer writes the given line to the buffer.
    80  func (c *CompressionBuffer) WriteLineToBuffer(line string) {
    81  	triggerFlush := false
    82  
    83  	c.mu.Lock()
    84  	c.buf.WriteString(line)
    85  	c.buf.WriteString("\n")
    86  	c.lines++
    87  	if c.lines >= c.batchSize {
    88  		triggerFlush = true
    89  	}
    90  	c.mu.Unlock()
    91  
    92  	// triggerFlush is decided within the locked section.
    93  	if triggerFlush {
    94  		c.compressAndCallback()
    95  	}
    96  }
    97  
    98  // Compress compresses the given bytes and encodes them into a string.
    99  func Compress(inBytes []byte) ([]byte, error) {
   100  	var outBuf bytes.Buffer
   101  	b64w := base64.NewEncoder(base64.StdEncoding, &outBuf)
   102  	gw := gzip.NewWriter(b64w)
   103  	if _, err := gw.Write(inBytes); err != nil {
   104  		return nil, err
   105  	}
   106  	gw.Close()
   107  	b64w.Close()
   108  
   109  	return outBuf.Bytes(), nil
   110  }
   111  
   112  // compressAndCallback compresses the data in buffer and writes it to outChan.
   113  func (c *CompressionBuffer) compressAndCallback() {
   114  	// Retrieve bytes from the buffer (c.buf) and get a new buffer for c.
   115  	c.mu.Lock()
   116  	inBytes := c.buf.Bytes()
   117  
   118  	// Start c's new buf with the same capacity as old buf.
   119  	c.buf = bytes.NewBuffer(make([]byte, 0, c.buf.Cap()))
   120  	c.lines = 0
   121  	c.mu.Unlock()
   122  
   123  	// Nothing to do.
   124  	if len(inBytes) == 0 {
   125  		return
   126  	}
   127  
   128  	compressed, err := Compress(inBytes)
   129  	if err != nil {
   130  		c.l.Errorf("Error while compressing bytes: %v, data: %s", err, string(inBytes))
   131  		return
   132  	}
   133  	c.callback(compressed)
   134  }
   135  
   136  // Close compresses the buffer and flushes it to the output channel.
   137  func (c *CompressionBuffer) Close() {
   138  	c.cancelCtx()
   139  	c.compressAndCallback()
   140  }