Write files in chunks (#293)

* Chunked output strategy
* Fix sorting with indexes >= 10

Makefile

@@ -1,4 +1,4 @@
-SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go middleware.go protocol.go
+SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go middleware.go protocol.go output_file_settings.go
 SOURCE_PATH = /go/src/github.com/buger/gor/
 RUN = docker run -v `pwd`:$(SOURCE_PATH) -p 0.0.0.0:8000:8000 -t -i gor
 BENCHMARK = BenchmarkRAWInput

input_file.go

@@ -63,7 +63,7 @@ func (i *FileInput) updateFile() (err error) {
 		return errors.New("No matching files")
 	}
 
-	sort.Strings(matches)
+	sort.Sort(sortByFileIndex(matches))
 
 	// Just pick first file, if there is many, and we are just started
 	if i.currentFile == nil {
@@ -120,6 +120,10 @@ func (i *FileInput) emit() {
 
 	var buffer bytes.Buffer
 
+	if i.currentReader == nil {
+		return
+	}
+
 	for {
 		line, err := i.currentReader.ReadBytes('\n')
 

input_file_test.go

@@ -157,14 +157,14 @@ func TestInputFileLoop(t *testing.T) {
 func TestInputFileCompressed(t *testing.T) {
 	rnd := rand.Int63()
 
-	output := NewFileOutput(fmt.Sprintf("/tmp/%d_0.gz", rnd), time.Minute)
+	output := NewFileOutput(fmt.Sprintf("/tmp/%d_0.gz", rnd), &FileOutputConfig{flushInterval: time.Minute, append: true})
 	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)
+	output2 := NewFileOutput(fmt.Sprintf("/tmp/%d_1.gz", rnd), &FileOutputConfig{flushInterval: time.Minute, append: true})
 	for i := 0; i < 1000; i++ {
 		output2.Write([]byte("1 1 1\r\ntest"))
 	}
@@ -248,7 +248,7 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
 		requestGenerator.wg.Done()
 	})
 
-	outputFile := NewFileOutput(f.Name(), time.Minute)
+	outputFile := NewFileOutput(f.Name(), &FileOutputConfig{flushInterval: time.Minute, append: true})
 
 	Plugins.Inputs = requestGenerator.inputs
 	Plugins.Outputs = []io.Writer{output, outputFile}

output_file.go

@@ -7,6 +7,9 @@ import (
 	"io"
 	"log"
 	"os"
+	"path/filepath"
+	"sort"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -21,24 +24,36 @@ var dateFileNameFuncs = map[string]func() string{
 	"%NS": func() string { return fmt.Sprint(time.Now().Nanosecond()) },
 }
 
+type FileOutputConfig struct {
+	flushInterval time.Duration
+	sizeLimit     unitSizeVar
+	queueLimit    int
+	append        bool
+}
+
 // FileOutput output plugin
 type FileOutput struct {
 	pathTemplate string
 	currentName  string
 	file         *os.File
+	queueLength  int
+	chunkSize    int
 	writer       io.Writer
+
+	config *FileOutputConfig
 }
 
 // NewFileOutput constructor for FileOutput, accepts path
-func NewFileOutput(pathTemplate string, flushInterval time.Duration) *FileOutput {
+func NewFileOutput(pathTemplate string, config *FileOutputConfig) *FileOutput {
 	o := new(FileOutput)
 	o.pathTemplate = pathTemplate
+	o.config = config
 	o.updateName()
 
 	// Force flushing every minute
 	go func() {
 		for {
-			time.Sleep(flushInterval)
+			time.Sleep(o.config.flushInterval)
 			o.flush()
 		}
 	}()
@@ -53,6 +68,57 @@ func NewFileOutput(pathTemplate string, flushInterval time.Duration) *FileOutput
 	return o
 }
 
+func getFileIndex(name string) int {
+	ext := filepath.Ext(name)
+	withoutExt := strings.TrimSuffix(name, ext)
+
+	if idx := strings.LastIndex(withoutExt, "_"); idx != -1 {
+		if i, err := strconv.Atoi(withoutExt[idx+1:]); err == nil {
+			return i
+		}
+	}
+
+	return -1
+}
+
+func setFileIndex(name string, idx int) string {
+	idxS := strconv.Itoa(idx)
+	ext := filepath.Ext(name)
+	withoutExt := strings.TrimSuffix(name, ext)
+
+	if i := strings.LastIndex(withoutExt, "_"); i != -1 {
+		withoutExt = withoutExt[:i]
+	}
+
+	return withoutExt + "_" + idxS + ext
+}
+
+func withoutIndex(s string) string {
+	if i := strings.LastIndex(s, "_"); i != -1 {
+		return s[:i]
+	}
+
+	return s
+}
+
+type sortByFileIndex []string
+
+func (s sortByFileIndex) Len() int {
+    return len(s)
+}
+
+func (s sortByFileIndex) Swap(i, j int) {
+    s[i], s[j] = s[j], s[i]
+}
+
+func (s sortByFileIndex) Less(i, j int) bool {
+    if withoutIndex(s[i]) == withoutIndex(s[j]) {
+    	return getFileIndex(s[i]) < getFileIndex(s[j])
+    }
+
+    return s[i] < s[j]
+}
+
 func (o *FileOutput) filename() string {
 	path := o.pathTemplate
 
@@ -60,6 +126,39 @@ func (o *FileOutput) filename() string {
 		path = strings.Replace(path, name, fn(), -1)
 	}
 
+	if !o.config.append {
+		nextChunk := false
+
+		if o.currentName == "" ||
+			((o.config.queueLimit > 0 && o.queueLength >= o.config.queueLimit) ||
+				(o.config.sizeLimit > 0 && o.chunkSize >= int(o.config.sizeLimit))) {
+			nextChunk = true
+		}
+
+		ext := filepath.Ext(path)
+		withoutExt := strings.TrimSuffix(path, ext)
+
+		if matches, err := filepath.Glob(withoutExt + "*" + ext); err == nil {
+			if len(matches) == 0 {
+				return setFileIndex(path, 0)
+			}
+			sort.Sort(sortByFileIndex(matches))
+
+			last := matches[len(matches)-1]
+
+			fileIndex := 0
+			if idx := getFileIndex(last); idx != -1 {
+				fileIndex = idx
+
+				if nextChunk {
+					fileIndex++
+				}
+			}
+
+			return setFileIndex(last, fileIndex)
+		}
+	}
+
 	return path
 }
 
@@ -76,6 +175,7 @@ func (o *FileOutput) Write(data []byte) (n int, err error) {
 		o.Close()
 
 		o.file, err = os.OpenFile(o.currentName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
+		o.file.Sync()
 
 		if strings.HasSuffix(o.currentName, ".gz") {
 			o.writer = gzip.NewWriter(o.file)
@@ -86,11 +186,15 @@ func (o *FileOutput) Write(data []byte) (n int, err error) {
 		if err != nil {
 			log.Fatal(o, "Cannot open file %q. Error: %s", o.currentName, err)
 		}
+
+		o.queueLength = 0
 	}
 
 	o.writer.Write(data)
 	o.writer.Write([]byte(payloadSeparator))
 
+	o.queueLength++
+
 	return len(data), nil
 }
 
@@ -102,6 +206,10 @@ func (o *FileOutput) flush() {
 			o.writer.(*bufio.Writer).Flush()
 		}
 	}
+
+	if stat, err := o.file.Stat(); err != nil {
+		o.chunkSize = int(stat.Size())
+	}
 }
 
 func (o *FileOutput) String() string {
@@ -117,4 +225,4 @@ func (o *FileOutput) Close() {
 		}
 		o.file.Close()
 	}
-}
+}

output_file_settings.go

@@ -0,0 +1,39 @@
+package main
+
+import (
+	"strconv"
+	"strings"
+)
+
+var dataUnitMap = map[byte]int64{
+	'k': 1024,
+	'm': 1024 * 1024,
+	'g': 1024 * 1024 * 1024,
+}
+
+func parseDataUnit(s string) int64 {
+	// Allow kb, mb, gb
+	if strings.HasSuffix(s, "b") {
+		s = s[:len(s)-1]
+	}
+
+	if unit, ok := dataUnitMap[s[len(s)-1]]; ok {
+		size, _ := strconv.ParseInt(s[:len(s)-1], 10, 64)
+		return unit * size
+	} else {
+		// If no unit specified use bytes
+		size, _ := strconv.ParseInt(s, 10, 64)
+		return size
+	}
+}
+
+type unitSizeVar int64
+
+func (u unitSizeVar) String() string {
+	return strconv.Itoa(int(u))
+}
+
+func (u unitSizeVar) Set(s string) error {
+	u = unitSizeVar(parseDataUnit(s))
+	return nil
+}

output_file_test.go

@@ -3,11 +3,14 @@ package main
 import (
 	"fmt"
 	"io"
+	"math/rand"
 	"os"
 	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
+	"sort"
+	"reflect"
 )
 
 func TestFileOutput(t *testing.T) {
@@ -15,7 +18,7 @@ func TestFileOutput(t *testing.T) {
 	quit := make(chan int)
 
 	input := NewTestInput()
-	output := NewFileOutput("/tmp/test_requests.gor", time.Minute)
+	output := NewFileOutput("/tmp/test_requests.gor", &FileOutputConfig{flushInterval: time.Minute, append: true})
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -51,7 +54,7 @@ func TestFileOutput(t *testing.T) {
 }
 
 func TestFileOutputPathTemplate(t *testing.T) {
-	output := &FileOutput{pathTemplate: "/tmp/log-%Y-%m-%d-%S"}
+	output := &FileOutput{pathTemplate: "/tmp/log-%Y-%m-%d-%S", config: &FileOutputConfig{flushInterval: time.Minute, append: true}}
 	now := time.Now()
 	expectedPath := fmt.Sprintf("/tmp/log-%s-%s-%s-%s", now.Format("2006"), now.Format("01"), now.Format("02"), now.Format("05"))
 	path := output.filename()
@@ -62,7 +65,7 @@ func TestFileOutputPathTemplate(t *testing.T) {
 }
 
 func TestFileOutputMultipleFiles(t *testing.T) {
-	output := NewFileOutput("/tmp/log-%Y-%m-%d-%S", time.Minute)
+	output := NewFileOutput("/tmp/log-%Y-%m-%d-%S", &FileOutputConfig{append: true, flushInterval: time.Minute})
 
 	if output.file != nil {
 		t.Error("Should not initialize file if no writes")
@@ -93,7 +96,7 @@ func TestFileOutputMultipleFiles(t *testing.T) {
 }
 
 func TestFileOutputCompression(t *testing.T) {
-	output := NewFileOutput("/tmp/log-%Y-%m-%d-%S.gz", time.Minute)
+	output := NewFileOutput("/tmp/log-%Y-%m-%d-%S.gz", &FileOutputConfig{append: true, flushInterval: time.Minute})
 
 	if output.file != nil {
 		t.Error("Should not initialize file if no writes")
@@ -113,3 +116,158 @@ func TestFileOutputCompression(t *testing.T) {
 
 	os.Remove(name)
 }
+
+func TestParseDataUnit(t *testing.T) {
+	var tests = []struct {
+		value string
+		size  int64
+	}{
+		{"100kb", dataUnitMap['k'] * 100},
+		{"100k", dataUnitMap['k'] * 100},
+		{"1kb", dataUnitMap['k']},
+		{"1g", dataUnitMap['g']},
+		{"10m", dataUnitMap['m'] * 10},
+		{"zsaa312", 0},
+	}
+
+	for _, c := range tests {
+		if parseDataUnit(c.value) != c.size {
+			t.Error(c.value, "should be", c.size, "instead", parseDataUnit(c.value))
+		}
+	}
+}
+
+func TestGetFileIndex(t *testing.T) {
+	var tests = []struct {
+		path  string
+		index int
+	}{
+		{"/tmp/logs", -1},
+		{"/tmp/logs_1", 1},
+		{"/tmp/logs_2.gz", 2},
+		{"/tmp/logs_0.gz", 0},
+	}
+
+	for _, c := range tests {
+		if getFileIndex(c.path) != c.index {
+			t.Error(c.path, "should be", c.index, "instead", getFileIndex(c.path))
+		}
+	}
+}
+
+func TestSetFileIndex(t *testing.T) {
+	var tests = []struct {
+		path    string
+		index   int
+		newPath string
+	}{
+		{"/tmp/logs", 0, "/tmp/logs_0"},
+		{"/tmp/logs.gz", 1, "/tmp/logs_1.gz"},
+		{"/tmp/logs_1", 0, "/tmp/logs_0"},
+		{"/tmp/logs_0", 10, "/tmp/logs_10"},
+		{"/tmp/logs_0.gz", 10, "/tmp/logs_10.gz"},
+	}
+
+	for _, c := range tests {
+		if setFileIndex(c.path, c.index) != c.newPath {
+			t.Error(c.path, "should be", c.newPath, "instead", setFileIndex(c.path, c.index))
+		}
+	}
+}
+
+func TestFileOutputAppendQueueLimitOverflow(t *testing.T) {
+	rnd := rand.Int63()
+	name := fmt.Sprintf("/tmp/%d", rnd)
+
+	output := NewFileOutput(name, &FileOutputConfig{append: false, flushInterval: time.Minute, queueLimit: 2})
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name1 := output.file.Name()
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name2 := output.file.Name()
+
+	output.updateName()
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name3 := output.file.Name()
+
+	if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0", rnd) {
+		t.Error("Fast changes should happen in same file:", name1, name2, name3)
+	}
+
+	if name3 == name1 || name3 != fmt.Sprintf("/tmp/%d_1", rnd) {
+		t.Error("File name should change:", name1, name2, name3)
+	}
+
+	os.Remove(name1)
+	os.Remove(name3)
+}
+
+func TestFileOutputAppendQueueLimitNoOverflow(t *testing.T) {
+	rnd := rand.Int63()
+	name := fmt.Sprintf("/tmp/%d", rnd)
+
+	output := NewFileOutput(name, &FileOutputConfig{append: false, flushInterval: time.Minute, queueLimit: 3})
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name1 := output.file.Name()
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name2 := output.file.Name()
+
+	output.updateName()
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name3 := output.file.Name()
+
+	if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0", rnd) {
+		t.Error("Fast changes should happen in same file:", name1, name2, name3)
+	}
+
+	if name3 != name1 || name3 != fmt.Sprintf("/tmp/%d_0", rnd) {
+		t.Error("File name should not change:", name1, name2, name3)
+	}
+
+	os.Remove(name1)
+	os.Remove(name3)
+}
+
+func TestFileOutputAppendQueueLimitGzips(t *testing.T) {
+	rnd := rand.Int63()
+	name := fmt.Sprintf("/tmp/%d.gz", rnd)
+
+	output := NewFileOutput(name, &FileOutputConfig{append: false, flushInterval: time.Minute, queueLimit: 2})
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name1 := output.file.Name()
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name2 := output.file.Name()
+
+	output.updateName()
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name3 := output.file.Name()
+
+	if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0.gz", rnd) {
+		t.Error("Fast changes should happen in same file:", name1, name2, name3)
+	}
+
+	if name3 == name1 || name3 != fmt.Sprintf("/tmp/%d_1.gz", rnd) {
+		t.Error("File name should change:", name1, name2, name3)
+	}
+
+	os.Remove(name1)
+	os.Remove(name3)
+}
+
+func TestFileOutputSort(t *testing.T) {
+	var files = []string{"2016_0", "2014_10", "2015_0", "2015_10", "2015_2"}
+	var expected = []string{"2014_10", "2015_0", "2015_2", "2015_10", "2016_0"}
+	sort.Sort(sortByFileIndex(files))
+
+	if !reflect.DeepEqual(files, expected) {
+		t.Error("Should properly sort file names using indexes", files, expected)
+	}
+}

plugins.go

@@ -117,7 +117,7 @@ func InitPlugins() {
 	}
 
 	for _, options := range Settings.outputFile {
-		registerPlugin(NewFileOutput, options, Settings.outputFileFlushInterval)
+		registerPlugin(NewFileOutput, options, Settings.outputFileConfig)
 	}
 
 	for _, options := range Settings.inputHTTP {

proto/proto.go

@@ -295,7 +295,7 @@ func Status(payload []byte) []byte {
 }
 
 var httpMethods []string = []string{
-	"GET ", "OPTI", "HEAD", "POST", "PUT ", "DELE", "TRAC", "CONN", /* custom methods */"BAN", "PURG"
+	"GET ", "OPTI", "HEAD", "POST", "PUT ", "DELE", "TRAC", "CONN", /* custom methods */"BAN", "PURG",
 }
 
 func IsHTTPPayload(payload []byte) bool {

settings.go

@@ -39,10 +39,10 @@ type AppSettings struct {
 	outputTCP      MultiOption
 	outputTCPStats bool
 
-	inputFile               MultiOption
-	inputFileLoop           bool
-	outputFile              MultiOption
-	outputFileFlushInterval time.Duration
+	inputFile        MultiOption
+	inputFileLoop    bool
+	outputFile       MultiOption
+	outputFileConfig FileOutputConfig
 
 	inputRAW              MultiOption
 	inputRAWEngine        string
@@ -88,7 +88,13 @@ func init() {
 	flag.BoolVar(&Settings.inputFileLoop, "input-file-loop", false, "Loop input files, useful for performance testing.")
 
 	flag.Var(&Settings.outputFile, "output-file", "Write incoming requests to file: \n\tgor --input-raw :80 --output-file ./requests.gor")
-	flag.DurationVar(&Settings.outputFileFlushInterval, "output-file-flush-interval", time.Minute, "Interval for forcing buffer flush to the file, default: 60s.")
+	flag.DurationVar(&Settings.outputFileConfig.flushInterval, "output-file-flush-interval", time.Minute, "Interval for forcing buffer flush to the file, default: 60s.")
+	flag.BoolVar(&Settings.outputFileConfig.append, "output-file-append", false, "The flushed chunk is appended to existence file or not. ")
+
+	// Set default
+	Settings.outputFileConfig.sizeLimit.Set("32mb")
+	flag.Var(&Settings.outputFileConfig.sizeLimit, "output-file-size-limit", "Size of each chunk. Default: 32mb")
+	flag.IntVar(&Settings.outputFileConfig.queueLimit, "output-file-queue-limit", 256, "The length of the chunk queue. Default: 256")
 
 	flag.Var(&Settings.inputRAW, "input-raw", "Capture traffic from given port (use RAW sockets and require *sudo* access):\n\t# Capture traffic from 8080 port\n\tgor --input-raw :8080 --output-http staging.com")