Add support for gzip in csv output format

Refactor:

- Add closeFn for CSV collector
- Made writeToFile protected

Signed-off-by: thejas <[email protected]>

stats/csv/collector.go

@@ -22,12 +22,13 @@ package csv
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/csv"
 	"fmt"
-	"io"
 	"os"
 	"sort"
+	"strings"
 	"sync"
 	"time"
 
@@ -40,7 +41,7 @@ import (
 
 // Collector saving output to csv implements the lib.Collector interface
 type Collector struct {
-	outfile      io.WriteCloser
+	closeFn      func() error
 	fname        string
 	resTags      []string
 	ignoredTags  []string
@@ -55,13 +56,6 @@ type Collector struct {
 // Verify that Collector implements lib.Collector
 var _ lib.Collector = &Collector{}
 
-// Similar to ioutil.NopCloser, but for writers
-type nopCloser struct {
-	io.Writer
-}
-
-func (nopCloser) Close() error { return nil }
-
 // New Creates new instance of CSV collector
 func New(fs afero.Fs, tags stats.TagSet, config Config) (*Collector, error) {
 	resTags := []string{}
@@ -80,32 +74,43 @@ func New(fs afero.Fs, tags stats.TagSet, config Config) (*Collector, error) {
 	fname := config.FileName.String
 
 	if fname == "" || fname == "-" {
-		logfile := nopCloser{os.Stdout}
 		return &Collector{
-			outfile:      logfile,
 			fname:        "-",
 			resTags:      resTags,
 			ignoredTags:  ignoredTags,
-			csvWriter:    csv.NewWriter(logfile),
+			csvWriter:    csv.NewWriter(os.Stdout),
 			row:          make([]string, 3+len(resTags)+1),
 			saveInterval: saveInterval,
+			closeFn:      func() error { return nil },
 		}, nil
 	}
 
-	logfile, err := fs.Create(fname)
+	logFile, err := fs.Create(fname)
 	if err != nil {
 		return nil, err
 	}
 
-	return &Collector{
-		outfile:      logfile,
+	c := Collector{
 		fname:        fname,
 		resTags:      resTags,
 		ignoredTags:  ignoredTags,
-		csvWriter:    csv.NewWriter(logfile),
 		row:          make([]string, 3+len(resTags)+1),
 		saveInterval: saveInterval,
-	}, nil
+	}
+
+	if strings.HasSuffix(fname, ".gz") {
+		outfile := gzip.NewWriter(logFile)
+		c.csvWriter = csv.NewWriter(outfile)
+		c.closeFn = func() error {
+			_ = outfile.Close()
+			return logFile.Close()
+		}
+	} else {
+		c.csvWriter = csv.NewWriter(logFile)
+		c.closeFn = logFile.Close
+	}
+
+	return &c, nil
 }
 
 // Init writes column names to csv file
@@ -125,16 +130,19 @@ func (c *Collector) SetRunStatus(status lib.RunStatus) {}
 // Run just blocks until the context is done
 func (c *Collector) Run(ctx context.Context) {
 	ticker := time.NewTicker(c.saveInterval)
+	defer func() {
+		err := c.closeFn()
+		if err != nil {
+			logrus.WithField("filename", c.fname).Errorf("CSV: Error closing the file: %v", err)
+		}
+	}()
+
 	for {
 		select {
 		case <-ticker.C:
-			c.WriteToFile()
+			c.writeToFile()
 		case <-ctx.Done():
-			c.WriteToFile()
-			err := c.outfile.Close()
-			if err != nil {
-				logrus.WithField("filename", c.fname).Error("CSV: Error closing the file")
-			}
+			c.writeToFile()
 			return
 		}
 	}
@@ -149,8 +157,8 @@ func (c *Collector) Collect(scs []stats.SampleContainer) {
 	}
 }
 
-// WriteToFile Writes samples to the csv file
-func (c *Collector) WriteToFile() {
+// writeToFile Writes samples to the csv file
+func (c *Collector) writeToFile() {
 	c.bufferLock.Lock()
 	samples := c.buffer
 	c.buffer = nil

stats/csv/collector_test.go

@@ -326,6 +326,7 @@ func TestNew(t *testing.T) {
 		fname       string
 		resTags     []string
 		ignoredTags []string
+		closeFn     func() error
 	}{
 		{
 			fname: "name",
@@ -365,6 +366,7 @@ func TestNew(t *testing.T) {
 			sort.Strings(expected.ignoredTags)
 			sort.Strings(collector.ignoredTags)
 			assert.Equal(t, expected.ignoredTags, collector.ignoredTags)
+			assert.NoError(t, collector.closeFn())
 		})
 	}
 }