gzip.go 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. // Copyright 2020-2021 InfluxData, Inc.. All rights reserved.
  2. // Use of this source code is governed by MIT
  3. // license that can be found in the LICENSE file.
  4. // Package gzip provides GZip related functionality
  5. package gzip
  6. import (
  7. "compress/gzip"
  8. "io"
  9. "sync"
  10. )
  11. // ReadWaitCloser is ReadCloser that waits for finishing underlying reader
  12. type ReadWaitCloser struct {
  13. pipeReader *io.PipeReader
  14. wg sync.WaitGroup
  15. }
  16. // Close closes underlying reader and waits for finishing operations
  17. func (r *ReadWaitCloser) Close() error {
  18. err := r.pipeReader.Close()
  19. r.wg.Wait() // wait for the gzip goroutine finish
  20. return err
  21. }
  22. // CompressWithGzip takes an io.Reader as input and pipes
  23. // it through a gzip.Writer returning an io.Reader containing
  24. // the gzipped data.
  25. // An error is returned if passing data to the gzip.Writer fails
  26. // this is shamelessly stolen from https://github.com/influxdata/telegraf
  27. func CompressWithGzip(data io.Reader) (io.ReadCloser, error) {
  28. pipeReader, pipeWriter := io.Pipe()
  29. gzipWriter := gzip.NewWriter(pipeWriter)
  30. rc := &ReadWaitCloser{
  31. pipeReader: pipeReader,
  32. }
  33. rc.wg.Add(1)
  34. var err error
  35. go func() {
  36. _, err = io.Copy(gzipWriter, data)
  37. gzipWriter.Close()
  38. // subsequent reads from the read half of the pipe will
  39. // return no bytes and the error err, or EOF if err is nil.
  40. pipeWriter.CloseWithError(err)
  41. rc.wg.Done()
  42. }()
  43. return pipeReader, err
  44. }