From b6f7de00c5e924d212adb0d49fcb11b4950867e8 Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Tue, 31 May 2016 16:41:01 +0500
Subject: [PATCH] Add compression support

---
 output_file.go      | 38 +++++++++++++++++++++++++-------------
 output_file_test.go | 29 ++++++++++++++++++++++++++---
 2 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/output_file.go b/output_file.go
index fec6bf4..8dcc11a 100644
--- a/output_file.go
+++ b/output_file.go
@@ -2,7 +2,9 @@ package main
 
 import (
 	"bufio"
+	"compress/gzip"
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"strings"
@@ -24,7 +26,7 @@ type FileOutput struct {
 	pathTemplate string
 	currentName  string
 	file         *os.File
-	writer       *bufio.Writer
+	writer       io.Writer
 }
 
 // NewFileOutput constructor for FileOutput, accepts path
@@ -37,9 +39,7 @@ func NewFileOutput(pathTemplate string, flushInterval time.Duration) *FileOutput
 	go func() {
 		for {
 			time.Sleep(flushInterval)
-			if err := o.writer.Flush(); err != nil {
-				break
-			}
+			o.flush()
 		}
 	}()
 
@@ -73,13 +73,15 @@ func (o *FileOutput) Write(data []byte) (n int, err error) {
 	}
 
 	if o.file == nil || o.currentName != o.file.Name() {
-		if o.file != nil {
-			o.writer.Flush()
-			o.file.Close()
-		}
+		o.Close()
 
 		o.file, err = os.OpenFile(o.currentName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
-		o.writer = bufio.NewWriter(o.file)
+
+		if strings.HasSuffix(o.currentName, ".gz") {
+			o.writer = gzip.NewWriter(o.file)
+		} else {
+			o.writer = bufio.NewWriter(o.file)
+		}
 
 		if err != nil {
 			log.Fatal(o, "Cannot open file %q. Error: %s", o.currentName, err)
@@ -92,9 +94,13 @@ func (o *FileOutput) Write(data []byte) (n int, err error) {
 	return len(data), nil
 }
 
-func (o *FileOutput) Flush() {
+func (o *FileOutput) flush() {
 	if o.file != nil {
-		o.writer.Flush()
+		if strings.HasSuffix(o.currentName, ".gz") {
+			o.writer.(*gzip.Writer).Flush()
+		} else {
+			o.writer.(*bufio.Writer).Flush()
+		}
 	}
 }
 
@@ -103,6 +109,12 @@ func (o *FileOutput) String() string {
 }
 
 func (o *FileOutput) Close() {
-	o.writer.Flush()
-	o.file.Close()
+	if o.file != nil {
+		if strings.HasSuffix(o.currentName, ".gz") {
+			o.writer.(*gzip.Writer).Close()
+		} else {
+			o.writer.(*bufio.Writer).Flush()
+		}
+		o.file.Close()
+	}
 }
diff --git a/output_file_test.go b/output_file_test.go
index 51a5109..7bda58a 100644
--- a/output_file_test.go
+++ b/output_file_test.go
@@ -29,7 +29,7 @@ func TestFileOutput(t *testing.T) {
 		input.EmitPOST()
 	}
 	time.Sleep(100 * time.Millisecond)
-	output.Flush()
+	output.flush()
 
 	close(quit)
 
@@ -77,18 +77,41 @@ func TestFileOutputMultipleFiles(t *testing.T) {
 	name2 := output.file.Name()
 
 	time.Sleep(time.Second)
+	output.updateName()
 
 	output.Write([]byte("1 1 1\r\ntest"))
 	name3 := output.file.Name()
 
 	if name2 != name1 {
-		t.Errorf("Fast changes should happen in same file:", name1, name2)
+		t.Error("Fast changes should happen in same file:", name1, name2, name3)
 	}
 
 	if name3 == name1 {
-		t.Errorf("File name should change:", name1, name3)
+		t.Error("File name should change:", name1, name2, name3)
 	}
 
 	os.Remove(name1)
 	os.Remove(name3)
 }
+
+func TestFileOutputCompression(t *testing.T) {
+	output := NewFileOutput("/tmp/log-%Y-%m-%d-%S.gz", time.Minute)
+
+	if output.file != nil {
+		t.Error("Should not initialize file if no writes")
+	}
+
+	for i := 0; i < 1000; i++ {
+		output.Write([]byte("1 1 1\r\ntest"))
+	}
+
+	name := output.file.Name()
+	output.Close()
+
+	s, _ := os.Stat(name)
+	if s.Size() == 12*1000 {
+		t.Error("Should be compressed file:", s.Size())
+	}
+
+	// os.Remove(name)/
+}