From ede0fb4c4ca4968486c3afad3c122e3aa4d8b59d Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Wed, 1 Jun 2016 18:49:20 +0500 Subject: [PATCH] Allow looping mode for --input-file --- input_file.go | 31 +++++++++++++++++++++++++++---- input_file_test.go | 28 +++++++++++++++++++++++++--- output_file_test.go | 2 +- plugins.go | 2 +- settings.go | 3 +++ 5 files changed, 57 insertions(+), 9 deletions(-) diff --git a/input_file.go b/input_file.go index 626dc27..281a1a1 100644 --- a/input_file.go +++ b/input_file.go @@ -22,14 +22,16 @@ type FileInput struct { currentFile *os.File currentReader *bufio.Reader speedFactor float64 + loop bool } // NewFileInput constructor for FileInput. Accepts file path as argument. -func NewFileInput(path string) (i *FileInput) { +func NewFileInput(path string, loop bool) (i *FileInput) { i = new(FileInput) i.data = make(chan []byte) i.path = path i.speedFactor = 1 + i.loop = loop if err := i.updateFile(); err != nil { return @@ -40,6 +42,12 @@ func NewFileInput(path string) (i *FileInput) { return } +type NextFileNotFound struct{} + +func (_ *NextFileNotFound) Error() string { + return "There is no new files" +} + // path can be a pattern // It sort paths lexicographically and tries to choose next one func (i *FileInput) updateFile() (err error) { @@ -57,6 +65,7 @@ func (i *FileInput) updateFile() (err error) { sort.Strings(matches) + // Just pick first file, if there is many, and we are just started if i.currentFile == nil { if i.currentFile, err = os.Open(matches[0]); err != nil { log.Println("Can't read file ", matches[0], err) @@ -76,7 +85,7 @@ func (i *FileInput) updateFile() (err error) { } if !found { - return errors.New("There is no new files") + return new(NextFileNotFound) } } @@ -121,8 +130,18 @@ func (i *FileInput) emit() { // If our path pattern match multiple files, try to find them if err == io.EOF { - if i.updateFile() != nil { - break + if e := i.updateFile(); e != nil { + if _, ok := e.(*NextFileNotFound); ok && i.loop { + // Start from the first file + i.currentFile = nil + i.currentReader = nil + lastTime = 0 + i.updateFile() + + continue + } else { + break + } } continue @@ -164,3 +183,7 @@ func (i *FileInput) emit() { log.Printf("FileInput: end of file '%s'\n", i.path) } + +func (i *FileInput) Close() { + i.currentFile.Close() +} diff --git a/input_file_test.go b/input_file_test.go index bcef3a7..78e0d56 100644 --- a/input_file_test.go +++ b/input_file_test.go @@ -112,7 +112,7 @@ func TestInputFileMultipleFiles(t *testing.T) { file2.Write([]byte(payloadSeparator)) file2.Close() - input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd)) + input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false) buf := make([]byte, 1000) n, _ := input.Read(buf) if buf[10] != '1' { @@ -132,6 +132,28 @@ func TestInputFileMultipleFiles(t *testing.T) { os.Remove(file2.Name()) } +func TestInputFileLoop(t *testing.T) { + rnd := rand.Int63() + + file, _ := os.OpenFile(fmt.Sprintf("/tmp/%d", rnd), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660) + file.Write([]byte("1 1 1\ntest1")) + file.Write([]byte(payloadSeparator)) + file.Write([]byte("1 1 2\ntest2")) + file.Write([]byte(payloadSeparator)) + file.Close() + + input := NewFileInput(fmt.Sprintf("/tmp/%d", rnd), true) + buf := make([]byte, 1000) + + // Even if we have just 2 requests in file, it should indifinitly loop + for i := 0; i < 1000; i++ { + input.Read(buf) + } + input.Close() + + os.Remove(file.Name()) +} + func TestInputFileCompressed(t *testing.T) { rnd := rand.Int63() @@ -149,7 +171,7 @@ func TestInputFileCompressed(t *testing.T) { name2 := output2.file.Name() output2.Close() - input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd)) + input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false) buf := make([]byte, 1000) for i := 0; i < 2000; i++ { input.Read(buf) @@ -250,7 +272,7 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback quit := make(chan int) wg := new(sync.WaitGroup) - input := NewFileInput(captureFile.Name()) + input := NewFileInput(captureFile.Name(), false) output := NewTestOutput(func(data []byte) { callback(data) wg.Done() diff --git a/output_file_test.go b/output_file_test.go index 51a27f7..5db0f8d 100644 --- a/output_file_test.go +++ b/output_file_test.go @@ -35,7 +35,7 @@ func TestFileOutput(t *testing.T) { quit = make(chan int) var counter int64 - input2 := NewFileInput("/tmp/test_requests.gor") + input2 := NewFileInput("/tmp/test_requests.gor", false) output2 := NewTestOutput(func(data []byte) { atomic.AddInt64(&counter, 1) wg.Done() diff --git a/plugins.go b/plugins.go index b593d78..805a086 100644 --- a/plugins.go +++ b/plugins.go @@ -113,7 +113,7 @@ func InitPlugins() { } for _, options := range Settings.inputFile { - registerPlugin(NewFileInput, options) + registerPlugin(NewFileInput, options, Settings.inputFileLoop) } for _, options := range Settings.outputFile { diff --git a/settings.go b/settings.go index 2296959..a49bd26 100644 --- a/settings.go +++ b/settings.go @@ -40,6 +40,7 @@ type AppSettings struct { outputTCPStats bool inputFile MultiOption + inputFileLoop bool outputFile MultiOption outputFileFlushInterval time.Duration @@ -84,6 +85,8 @@ func init() { flag.BoolVar(&Settings.outputTCPStats, "output-tcp-stats", false, "Report TCP output queue stats to console every 5 seconds.") flag.Var(&Settings.inputFile, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com") + 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.")