Merge pull request #290 from buger/output-file-improvements

File replay and recording improvements

Makefile

@@ -32,7 +32,7 @@ race:
 	$(RUN) go test ./... $(ARGS) -v -race -timeout 15s
 
 test:
-	$(RUN) go test ./. -race -timeout 60s $(LDFLAGS) $(ARGS)  -v
+	$(RUN) go test ./. -timeout 60s $(LDFLAGS) $(ARGS)  -v
 
 test_all:
 	$(RUN) go test ./... -timeout 60s $(LDFLAGS) $(ARGS) -v

input_file.go

@@ -3,42 +3,103 @@ package main
 import (
 	"bufio"
 	"bytes"
+	"compress/gzip"
+	"errors"
 	"io"
 	"log"
 	"os"
+	"path/filepath"
+	"sort"
 	"strconv"
+	"strings"
 	"time"
 )
 
 // FileInput can read requests generated by FileOutput
 type FileInput struct {
-	data        chan []byte
-	path        string
-	file        *os.File
-	speedFactor float64
+	data          chan []byte
+	path          string
+	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.init(path)
+	i.loop = loop
+
+	if err := i.updateFile(); err != nil {
+		return
+	}
 
 	go i.emit()
 
 	return
 }
 
-func (i *FileInput) init(path string) {
-	file, err := os.Open(path)
+type NextFileNotFound struct{}
 
-	if err != nil {
-		log.Fatal(i, "Cannot open file %q. Error: %s", path, err)
+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) {
+	var matches []string
+
+	if matches, err = filepath.Glob(i.path); err != nil {
+		log.Println("Wrong file pattern", i.path, err)
+		return
 	}
 
-	i.file = file
+	if len(matches) == 0 {
+		log.Println("No files match pattern: ", i.path)
+		return errors.New("No matching files")
+	}
+
+	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)
+			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
+				}
+
+				found = true
+			}
+		}
+
+		if !found {
+			return new(NextFileNotFound)
+		}
+	}
+
+	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) {
@@ -56,18 +117,35 @@ func (i *FileInput) emit() {
 	var lastTime int64
 
 	payloadSeparatorAsBytes := []byte(payloadSeparator)
-	reader := bufio.NewReader(i.file)
+
 	var buffer bytes.Buffer
 
 	for {
-		line, err := reader.ReadBytes('\n')
+		line, err := i.currentReader.ReadBytes('\n')
 
 		if err != nil {
 			if err != io.EOF {
 				log.Fatal(err)
 			}
-			break
 
+			// If our path pattern match multiple files, try to find them
+			if err == io.EOF {
+				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
+			}
 		}
 
 		if bytes.Equal(payloadSeparatorAsBytes[1:], line) {
@@ -105,3 +183,7 @@ func (i *FileInput) emit() {
 
 	log.Printf("FileInput: end of file '%s'\n", i.path)
 }
+
+func (i *FileInput) Close() {
+	i.currentFile.Close()
+}

input_file_test.go

@@ -3,8 +3,11 @@ package main
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"io"
 	"io/ioutil"
+	"log"
+	"math/rand"
 	"os"
 	"sync"
 	"syscall"
@@ -12,6 +15,8 @@ import (
 	"time"
 )
 
+var _ = log.Println
+
 func TestInputFileWithGET(t *testing.T) {
 
 	input := NewTestInput()
@@ -35,7 +40,6 @@ func TestInputFileWithGET(t *testing.T) {
 			t.Error("Request read back from file should match")
 		}
 	}
-
 }
 
 func TestInputFileWithPayloadLargerThan64Kb(t *testing.T) {
@@ -93,6 +97,90 @@ func TestInputFileWithGETAndPOST(t *testing.T) {
 
 }
 
+func TestInputFileMultipleFiles(t *testing.T) {
+	rnd := rand.Int63()
+
+	file1, _ := os.OpenFile(fmt.Sprintf("/tmp/%d_0", rnd), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
+	file1.Write([]byte("1 1 1\ntest1"))
+	file1.Write([]byte(payloadSeparator))
+	file1.Write([]byte("1 1 2\ntest2"))
+	file1.Write([]byte(payloadSeparator))
+	file1.Close()
+
+	file2, _ := os.OpenFile(fmt.Sprintf("/tmp/%d_1", rnd), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
+	file2.Write([]byte("1 1 3\ntest3"))
+	file2.Write([]byte(payloadSeparator))
+	file2.Close()
+
+	input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false)
+	buf := make([]byte, 1000)
+	n, _ := input.Read(buf)
+	if buf[10] != '1' {
+		t.Error("Shound emit requests in right order", string(buf[:n]))
+	}
+	input.Read(buf)
+	if buf[10] != '2' {
+		t.Error("Shound emit requests in right order", string(buf[:n]))
+	}
+
+	input.Read(buf)
+	if buf[10] != '3' {
+		t.Error("Shound emit requests from second file", string(buf[:n]))
+	}
+
+	os.Remove(file1.Name())
+	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()
+
+	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), false)
+	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
@@ -155,13 +243,12 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
 
 	readPayloads := [][]byte{}
 	output := NewTestOutput(func(data []byte) {
-
 		readPayloads = append(readPayloads, Duplicate(data))
 
 		requestGenerator.wg.Done()
 	})
 
-	outputFile := NewFileOutput(f.Name())
+	outputFile := NewFileOutput(f.Name(), time.Minute)
 
 	Plugins.Inputs = requestGenerator.inputs
 	Plugins.Outputs = []io.Writer{output, outputFile}
@@ -171,6 +258,9 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
 	requestGenerator.emit()
 	requestGenerator.wg.Wait()
 
+	time.Sleep(100 * time.Millisecond)
+	outputFile.Close()
+
 	close(quit)
 
 	return NewExpectedCaptureFile(readPayloads, f)
@@ -182,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()

output_file.go

@@ -1,34 +1,70 @@
 package main
 
 import (
+	"bufio"
+	"compress/gzip"
+	"fmt"
 	"io"
 	"log"
 	"os"
+	"strings"
+	"time"
 )
 
+var dateFileNameFuncs = map[string]func() string{
+	"%Y":  func() string { return time.Now().Format("2006") },
+	"%m":  func() string { return time.Now().Format("01") },
+	"%d":  func() string { return time.Now().Format("02") },
+	"%H":  func() string { return time.Now().Format("15") },
+	"%M":  func() string { return time.Now().Format("04") },
+	"%S":  func() string { return time.Now().Format("05") },
+	"%NS": func() string { return fmt.Sprint(time.Now().Nanosecond()) },
+}
+
 // FileOutput output plugin
 type FileOutput struct {
-	path string
-	file *os.File
+	pathTemplate string
+	currentName  string
+	file         *os.File
+	writer       io.Writer
 }
 
 // NewFileOutput constructor for FileOutput, accepts path
-func NewFileOutput(path string) io.Writer {
+func NewFileOutput(pathTemplate string, flushInterval time.Duration) *FileOutput {
 	o := new(FileOutput)
-	o.path = path
-	o.init(path)
+	o.pathTemplate = pathTemplate
+	o.updateName()
+
+	// Force flushing every minute
+	go func() {
+		for {
+			time.Sleep(flushInterval)
+			o.flush()
+		}
+	}()
+
+	go func() {
+		for {
+			time.Sleep(time.Second)
+			o.updateName()
+		}
+	}()
 
 	return o
 }
 
-func (o *FileOutput) init(path string) {
-	var err error
+func (o *FileOutput) filename() string {
+	path := o.pathTemplate
 
-	o.file, err = os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
-
-	if err != nil {
-		log.Fatal(o, "Cannot open file %q. Error: %s", path, err)
+	for name, fn := range dateFileNameFuncs {
+		path = strings.Replace(path, name, fn(), -1)
 	}
+
+	return path
+}
+
+func (o *FileOutput) updateName() {
+	o.currentName = o.filename()
 }
 
 func (o *FileOutput) Write(data []byte) (n int, err error) {
@@ -36,12 +72,49 @@ func (o *FileOutput) Write(data []byte) (n int, err error) {
 		return len(data), nil
 	}
 
-	o.file.Write(data)
-	o.file.Write([]byte(payloadSeparator))
+	if o.file == nil || o.currentName != o.file.Name() {
+		o.Close()
+
+		o.file, err = os.OpenFile(o.currentName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
+
+		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)
+		}
+	}
+
+	o.writer.Write(data)
+	o.writer.Write([]byte(payloadSeparator))
 
 	return len(data), nil
 }
 
+func (o *FileOutput) flush() {
+	if o.file != nil {
+		if strings.HasSuffix(o.currentName, ".gz") {
+			o.writer.(*gzip.Writer).Flush()
+		} else {
+			o.writer.(*bufio.Writer).Flush()
+		}
+	}
+}
+
 func (o *FileOutput) String() string {
-	return "File output: " + o.path
+	return "File output: " + o.file.Name()
+}
+
+func (o *FileOutput) 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()
+	}
 }

output_file_test.go

@@ -1,9 +1,13 @@
 package main
 
 import (
+	"fmt"
 	"io"
+	"os"
 	"sync"
+	"sync/atomic"
 	"testing"
+	"time"
 )
 
 func TestFileOutput(t *testing.T) {
@@ -11,7 +15,7 @@ func TestFileOutput(t *testing.T) {
 	quit := make(chan int)
 
 	input := NewTestInput()
-	output := NewFileOutput("/tmp/test_requests.gor")
+	output := NewFileOutput("/tmp/test_requests.gor", time.Minute)
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -23,12 +27,17 @@ func TestFileOutput(t *testing.T) {
 		input.EmitGET()
 		input.EmitPOST()
 	}
+	time.Sleep(100 * time.Millisecond)
+	output.flush()
+
 	close(quit)
 
 	quit = make(chan int)
 
-	input2 := NewFileInput("/tmp/test_requests.gor")
+	var counter int64
+	input2 := NewFileInput("/tmp/test_requests.gor", false)
 	output2 := NewTestOutput(func(data []byte) {
+		atomic.AddInt64(&counter, 1)
 		wg.Done()
 	})
 
@@ -40,3 +49,67 @@ func TestFileOutput(t *testing.T) {
 	wg.Wait()
 	close(quit)
 }
+
+func TestFileOutputPathTemplate(t *testing.T) {
+	output := &FileOutput{pathTemplate: "/tmp/log-%Y-%m-%d-%S"}
+	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()
+
+	if expectedPath != path {
+		t.Errorf("Expected path %s but got %s", expectedPath, path)
+	}
+}
+
+func TestFileOutputMultipleFiles(t *testing.T) {
+	output := NewFileOutput("/tmp/log-%Y-%m-%d-%S", time.Minute)
+
+	if output.file != nil {
+		t.Error("Should not initialize file if no writes")
+	}
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name1 := output.file.Name()
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	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.Error("Fast changes should happen in same file:", name1, name2, name3)
+	}
+
+	if name3 == name1 {
+		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)
+}

plugins.go

@@ -113,11 +113,11 @@ func InitPlugins() {
 	}
 
 	for _, options := range Settings.inputFile {
-		registerPlugin(NewFileInput, options)
+		registerPlugin(NewFileInput, options, Settings.inputFileLoop)
 	}
 
 	for _, options := range Settings.outputFile {
-		registerPlugin(NewFileOutput, options)
+		registerPlugin(NewFileOutput, options, Settings.outputFileFlushInterval)
 	}
 
 	for _, options := range Settings.inputHTTP {

settings.go

@@ -39,8 +39,10 @@ type AppSettings struct {
 	outputTCP      MultiOption
 	outputTCPStats bool
 
-	inputFile  MultiOption
-	outputFile MultiOption
+	inputFile               MultiOption
+	inputFileLoop           bool
+	outputFile              MultiOption
+	outputFileFlushInterval time.Duration
 
 	inputRAW              MultiOption
 	inputRAWEngine        string
@@ -83,7 +85,10 @@ 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.")
 
 	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")