Fix influxdb output serialization on connection closed (#6621)

internal/internal.go

@@ -16,6 +16,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
 	"unicode"
@@ -50,6 +51,11 @@ type Number struct {
 	Value float64
 }
 
+type ReadWaitCloser struct {
+	pipeReader *io.PipeReader
+	wg         sync.WaitGroup
+}
+
 // SetVersion sets the telegraf agent version
 func SetVersion(v string) error {
 	if version != "" {
@@ -281,14 +287,25 @@ func ExitStatus(err error) (int, bool) {
 	return 0, false
 }
 
+func (r *ReadWaitCloser) Close() error {
+	err := r.pipeReader.Close()
+	r.wg.Wait() // wait for the gzip goroutine finish
+	return err
+}
+
 // CompressWithGzip takes an io.Reader as input and pipes
 // it through a gzip.Writer returning an io.Reader containing
 // the gzipped data.
 // An error is returned if passing data to the gzip.Writer fails
-func CompressWithGzip(data io.Reader) (io.Reader, error) {
+func CompressWithGzip(data io.Reader) (io.ReadCloser, error) {
 	pipeReader, pipeWriter := io.Pipe()
 	gzipWriter := gzip.NewWriter(pipeWriter)
 
+	rc := &ReadWaitCloser{
+		pipeReader: pipeReader,
+	}
+
+	rc.wg.Add(1)
 	var err error
 	go func() {
 		_, err = io.Copy(gzipWriter, data)
@@ -296,6 +313,7 @@ func CompressWithGzip(data io.Reader) (io.Reader, error) {
 		// subsequent reads from the read half of the pipe will
 		// return no bytes and the error err, or EOF if err is nil.
 		pipeWriter.CloseWithError(err)
+		rc.wg.Done()
 	}()
 
 	return pipeReader, err

internal/internal_test.go

@@ -3,6 +3,8 @@ package internal
 import (
 	"bytes"
 	"compress/gzip"
+	"crypto/rand"
+	"io"
 	"io/ioutil"
 	"log"
 	"os/exec"
@@ -232,6 +234,38 @@ func TestCompressWithGzip(t *testing.T) {
 	assert.Equal(t, testData, string(output))
 }
 
+type mockReader struct {
+	readN uint64 // record the number of calls to Read
+}
+
+func (r *mockReader) Read(p []byte) (n int, err error) {
+	r.readN++
+	return rand.Read(p)
+}
+
+func TestCompressWithGzipEarlyClose(t *testing.T) {
+	mr := &mockReader{}
+
+	rc, err := CompressWithGzip(mr)
+	assert.NoError(t, err)
+
+	n, err := io.CopyN(ioutil.Discard, rc, 10000)
+	assert.NoError(t, err)
+	assert.Equal(t, int64(10000), n)
+
+	r1 := mr.readN
+	err = rc.Close()
+	assert.NoError(t, err)
+
+	n, err = io.CopyN(ioutil.Discard, rc, 10000)
+	assert.Error(t, io.EOF, err)
+	assert.Equal(t, int64(0), n)
+
+	r2 := mr.readN
+	// no more read to the source after closing
+	assert.Equal(t, r1, r2)
+}
+
 func TestVersionAlreadySet(t *testing.T) {
 	err := SetVersion("foo")
 	assert.Nil(t, err)