From 651bd82d3d6ba2acebf354988c2ef36f9e22e212 Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Wed, 1 Jun 2016 10:41:29 +0500
Subject: [PATCH] input-file support for gzip files

---
 input_file.go      | 39 +++++++++++++++++++++++++++------------
 input_file_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/input_file.go b/input_file.go
index 19f4dc1..626dc27 100644
--- a/input_file.go
+++ b/input_file.go
@@ -3,6 +3,7 @@ package main
 import (
 	"bufio"
 	"bytes"
+	"compress/gzip"
 	"errors"
 	"io"
 	"log"
@@ -10,15 +11,17 @@ import (
 	"path/filepath"
 	"sort"
 	"strconv"
+	"strings"
 	"time"
 )
 
 // FileInput can read requests generated by FileOutput
 type FileInput struct {
-	data        chan []byte
-	path        string
-	currentFile *os.File
-	speedFactor float64
+	data          chan []byte
+	path          string
+	currentFile   *os.File
+	currentReader *bufio.Reader
+	speedFactor   float64
 }
 
 // NewFileInput constructor for FileInput. Accepts file path as argument.
@@ -59,22 +62,35 @@ func (i *FileInput) updateFile() (err error) {
 			log.Println("Can't read file ", matches[0], err)
 			return
 		}
-
-		return
 	} else {
+		found := false
 		for idx, p := range matches {
 			if p == i.currentFile.Name() && idx != len(matches)-1 {
 				if i.currentFile, err = os.Open(matches[idx+1]); err != nil {
 					log.Println("Can't read file ", matches[idx+1], err)
 					return
-				} else {
-					return nil
 				}
+
+				found = true
 			}
 		}
 
-		return errors.New("There is no new files")
+		if !found {
+			return errors.New("There is no new files")
+		}
 	}
+
+	if strings.HasSuffix(i.currentFile.Name(), ".gz") {
+		gzReader, err := gzip.NewReader(i.currentFile)
+		if err != nil {
+			log.Fatal(err)
+		}
+		i.currentReader = bufio.NewReader(gzReader)
+	} else {
+		i.currentReader = bufio.NewReader(i.currentFile)
+	}
+
+	return nil
 }
 
 func (i *FileInput) Read(data []byte) (int, error) {
@@ -92,11 +108,11 @@ func (i *FileInput) emit() {
 	var lastTime int64
 
 	payloadSeparatorAsBytes := []byte(payloadSeparator)
-	reader := bufio.NewReader(i.currentFile)
+
 	var buffer bytes.Buffer
 
 	for {
-		line, err := reader.ReadBytes('\n')
+		line, err := i.currentReader.ReadBytes('\n')
 
 		if err != nil {
 			if err != io.EOF {
@@ -109,7 +125,6 @@ func (i *FileInput) emit() {
 					break
 				}
 
-				reader = bufio.NewReader(i.currentFile)
 				continue
 			}
 		}
diff --git a/input_file_test.go b/input_file_test.go
index 95b60a6..bcef3a7 100644
--- a/input_file_test.go
+++ b/input_file_test.go
@@ -132,6 +132,33 @@ func TestInputFileMultipleFiles(t *testing.T) {
 	os.Remove(file2.Name())
 }
 
+func TestInputFileCompressed(t *testing.T) {
+	rnd := rand.Int63()
+
+	output := NewFileOutput(fmt.Sprintf("/tmp/%d_0.gz", rnd), time.Minute)
+	for i := 0; i < 1000; i++ {
+		output.Write([]byte("1 1 1\r\ntest"))
+	}
+	name1 := output.file.Name()
+	output.Close()
+
+	output2 := NewFileOutput(fmt.Sprintf("/tmp/%d_1.gz", rnd), time.Minute)
+	for i := 0; i < 1000; i++ {
+		output2.Write([]byte("1 1 1\r\ntest"))
+	}
+	name2 := output2.file.Name()
+	output2.Close()
+
+	input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd))
+	buf := make([]byte, 1000)
+	for i := 0; i < 2000; i++ {
+		input.Read(buf)
+	}
+
+	os.Remove(name1)
+	os.Remove(name2)
+}
+
 type CaptureFile struct {
 	data [][]byte
 	file *os.File