From cefa2b5cc204d03b9a2a3e6a98920df071f4a46a Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Tue, 7 Jun 2016 18:46:13 +0600 Subject: [PATCH] Write files in chunks (#293) * Chunked output strategy * Fix sorting with indexes >= 10 --- Makefile | 2 +- input_file.go | 6 +- input_file_test.go | 6 +- output_file.go | 114 ++++++++++++++++++++++++++- output_file_settings.go | 39 ++++++++++ output_file_test.go | 166 +++++++++++++++++++++++++++++++++++++++- plugins.go | 2 +- proto/proto.go | 2 +- settings.go | 16 ++-- 9 files changed, 334 insertions(+), 19 deletions(-) create mode 100644 output_file_settings.go diff --git a/Makefile b/Makefile index 4eb8a86..81feb0e 100644 --- a/Makefile +++ b/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 diff --git a/input_file.go b/input_file.go index 281a1a1..7f59a49 100644 --- a/input_file.go +++ b/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') diff --git a/input_file_test.go b/input_file_test.go index 78e0d56..e5857cf 100644 --- a/input_file_test.go +++ b/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} diff --git a/output_file.go b/output_file.go index 8dcc11a..9e99c94 100644 --- a/output_file.go +++ b/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() } -} +} \ No newline at end of file diff --git a/output_file_settings.go b/output_file_settings.go new file mode 100644 index 0000000..e011f61 --- /dev/null +++ b/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 +} diff --git a/output_file_test.go b/output_file_test.go index 5db0f8d..765045c 100644 --- a/output_file_test.go +++ b/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) + } +} \ No newline at end of file diff --git a/plugins.go b/plugins.go index 805a086..e85dee9 100644 --- a/plugins.go +++ b/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 { diff --git a/proto/proto.go b/proto/proto.go index 41768e2..0e74cc8 100644 --- a/proto/proto.go +++ b/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 { diff --git a/settings.go b/settings.go index a49bd26..72a8c77 100644 --- a/settings.go +++ b/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")