diff --git a/Makefile b/Makefile index ac0cd62..0835fd3 100644 --- a/Makefile +++ b/Makefile @@ -2,7 +2,7 @@ SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go SOURCE_PATH = /go/src/github.com/buger/gor/ PORT = 8000 FADDR = :8000 -RUN = docker run -v `pwd`:$(SOURCE_PATH) -p 0.0.0.0:$(PORT):$(PORT) -t -i gor +RUN = docker run -v `pwd`:$(SOURCE_PATH) -e AWS_ACCESS_KEY_ID=AKIAIOJPOSYCQCWU4YHQ -e AWS_SECRET_ACCESS_KEY=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 -p 0.0.0.0:$(PORT):$(PORT) -t -i gor BENCHMARK = BenchmarkRAWInput TEST = TestRawListenerBench VERSION = DEV-$(shell date +%s) diff --git a/input_file.go b/input_file.go index b562a99..99ca58e 100644 --- a/input_file.go +++ b/input_file.go @@ -9,31 +9,114 @@ import ( "log" "os" "path/filepath" - "sort" "strconv" "strings" + "sync" "time" ) +type fileInputReader struct { + reader *bufio.Reader + data []byte + file *os.File + timestamp int64 +} + +func (f *fileInputReader) parseNext() error { + payloadSeparatorAsBytes := []byte(payloadSeparator) + var buffer bytes.Buffer + + for { + line, err := f.reader.ReadBytes('\n') + + if err != nil { + if err != io.EOF { + log.Println(err) + return err + } + + if err == io.EOF { + f.file.Close() + f.file = nil + return err + } + } + + if bytes.Equal(payloadSeparatorAsBytes[1:], line) { + asBytes := buffer.Bytes() + meta := payloadMeta(asBytes) + + f.timestamp, _ = strconv.ParseInt(string(meta[2]), 10, 64) + f.data = asBytes[:len(asBytes)-1] + + return nil + } + + buffer.Write(line) + } + + return nil +} + +func (f *fileInputReader) ReadPayload() []byte { + defer f.parseNext() + + return f.data +} +func (f *fileInputReader) Close() error { + if f.file != nil { + f.file.Close() + } + + return nil +} + +func NewFileInputReader(path string) *fileInputReader { + file, err := os.Open(path) + + if err != nil { + log.Println(err) + return nil + } + + r := &fileInputReader{file: file} + if strings.HasSuffix(path, ".gz") { + gzReader, err := gzip.NewReader(file) + if err != nil { + log.Println(err) + return nil + } + r.reader = bufio.NewReader(gzReader) + } else { + r.reader = bufio.NewReader(file) + } + + r.parseNext() + + return r +} + // FileInput can read requests generated by FileOutput type FileInput struct { - data chan []byte - path string - currentFile *os.File - currentReader *bufio.Reader - speedFactor float64 - loop bool + mu sync.Mutex + data chan []byte + exit chan bool + path string + readers []*fileInputReader + speedFactor float64 + loop bool } // NewFileInput constructor for FileInput. Accepts file path as argument. func NewFileInput(path string, loop bool) (i *FileInput) { i = new(FileInput) - i.data = make(chan []byte) + i.data = make(chan []byte, 1000) + i.exit = make(chan bool, 1) i.path = path i.speedFactor = 1 i.loop = loop - if err := i.updateFile(); err != nil { + if err := i.init(); err != nil { return } @@ -48,9 +131,10 @@ 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) { +func (i *FileInput) init() (err error) { + defer i.mu.Unlock() + i.mu.Lock() + var matches []string if matches, err = filepath.Glob(i.path); err != nil { @@ -63,40 +147,10 @@ func (i *FileInput) updateFile() (err error) { return errors.New("No matching files") } - sort.Sort(sortByFileIndex(matches)) + i.readers = make([]*fileInputReader, len(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) + for idx, p := range matches { + i.readers[idx] = NewFileInputReader(p) } return nil @@ -113,82 +167,72 @@ func (i *FileInput) String() string { return "File input: " + i.path } -func (i *FileInput) emit() { - var lastTime int64 +// Find reader with smallest timestamp e.g next payload in row +func (i *FileInput) nextReader() (next *fileInputReader) { + for _, r := range i.readers { + if r == nil || r.file == nil { + continue + } - payloadSeparatorAsBytes := []byte(payloadSeparator) - - var buffer bytes.Buffer - - if i.currentReader == nil { - return + if next == nil || r.timestamp < next.timestamp { + next = r + continue + } } + return +} + +func (i *FileInput) emit() { + var lastTime int64 = -1 + for { - line, err := i.currentReader.ReadBytes('\n') + select { + case <-i.exit: + return + default: + } - if err != nil { - if err != io.EOF { - log.Fatal(err) - } - - // 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.Close() - i.currentFile = nil - i.currentReader = nil - lastTime = 0 - i.updateFile() - - continue - } else { - break - } - } + reader := i.nextReader() + if reader == nil { + if i.loop { + i.init() + lastTime = -1 continue + } else { + break } } - if bytes.Equal(payloadSeparatorAsBytes[1:], line) { - asBytes := buffer.Bytes() - buffer.Reset() + if lastTime != -1 { + diff := reader.timestamp - lastTime + lastTime = reader.timestamp - meta := payloadMeta(asBytes) - - if len(meta) > 2 && meta[0][0] == RequestPayload { - ts, _ := strconv.ParseInt(string(meta[2]), 10, 64) - - if lastTime != 0 { - timeDiff := ts - lastTime - - if i.speedFactor != 1 { - timeDiff = int64(float64(timeDiff) / i.speedFactor) - } - - time.Sleep(time.Duration(timeDiff)) - } - - lastTime = ts + if i.speedFactor != 1 { + diff = int64(float64(diff) / i.speedFactor) } - // Bytes() returns only pointer, so to remove data-race copy the data to an array - newBuf := make([]byte, len(asBytes)-1) - copy(newBuf, asBytes) - - i.data <- newBuf + time.Sleep(time.Duration(diff)) } else { - buffer.Write(line) + lastTime = reader.timestamp } + i.data <- reader.ReadPayload() } log.Printf("FileInput: end of file '%s'\n", i.path) } -func (i *FileInput) Close() { - i.currentFile.Close() -} +func (i *FileInput) Close() error { + defer i.mu.Unlock() + i.mu.Lock() + + i.exit <- true + + for _, r := range i.readers { + r.Close() + } + + return nil +} \ No newline at end of file diff --git a/input_file_test.go b/input_file_test.go index e5857cf..61e9212 100644 --- a/input_file_test.go +++ b/input_file_test.go @@ -97,35 +97,109 @@ func TestInputFileWithGETAndPOST(t *testing.T) { } -func TestInputFileMultipleFiles(t *testing.T) { +func TestInputFileMultipleFilesWithRequestsOnly(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("1 1 3\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("1 1 2\ntest3")) + file2.Write([]byte(payloadSeparator)) + file2.Write([]byte("1 1 4\ntest4")) 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])) + + for i := '1'; i <= '4'; i++ { + n, _ := input.Read(buf) + if buf[4] != byte(i) { + t.Error("Should 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 TestInputFileRequestsWithLatency(t *testing.T) { + rnd := rand.Int63() + + file, _ := os.OpenFile(fmt.Sprintf("/tmp/%d", rnd), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660) + defer file.Close() + + file.Write([]byte("1 1 100000000\nrequest1")) + file.Write([]byte(payloadSeparator)) + file.Write([]byte("1 2 150000000\nrequest2")) + file.Write([]byte(payloadSeparator)) + file.Write([]byte("1 3 250000000\nrequest3")) + file.Write([]byte(payloadSeparator)) + + input := NewFileInput(fmt.Sprintf("/tmp/%d", rnd), false) + buf := make([]byte, 1000) + + start := time.Now().UnixNano() + for i := 0; i < 3; i++ { + input.Read(buf) + } + end := time.Now().UnixNano() + + var expectedLatency int64 = 250000000 - 100000000 + realLatency := end - start + if realLatency < expectedLatency { + t.Errorf("Should emit requests respecting latency. Expected: %v, real: %v", expectedLatency, realLatency) + } + + if realLatency > expectedLatency+10000000 { + t.Errorf("Should emit requests respecting latency. Expected: %v, real: %v", expectedLatency, realLatency) + + } +} + +func TestInputFileMultipleFilesWithRequestsAndResponses(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\nrequest1")) + file1.Write([]byte(payloadSeparator)) + file1.Write([]byte("2 1 1\nresponse1")) + file1.Write([]byte(payloadSeparator)) + file1.Write([]byte("1 2 3\nrequest2")) + file1.Write([]byte(payloadSeparator)) + file1.Write([]byte("2 2 3\nresponse2")) + 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 3 2\nrequest3")) + file2.Write([]byte(payloadSeparator)) + file2.Write([]byte("2 3 2\nresponse3")) + file2.Write([]byte(payloadSeparator)) + file2.Write([]byte("1 4 4\nrequest4")) + file2.Write([]byte(payloadSeparator)) + file2.Write([]byte("2 4 4\nresponse4")) + file2.Write([]byte(payloadSeparator)) + file2.Close() + + input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false) + buf := make([]byte, 1000) + + for i := '1'; i <= '4'; i++ { + n, _ := input.Read(buf) + if buf[0] != '1' && buf[4] != byte(i) { + t.Error("Shound emit requests in right order", string(buf[:n])) + } + + n, _ = input.Read(buf) + if buf[0] != '2' && buf[4] != byte(i) { + t.Error("Shound emit responses in right order", string(buf[:n])) + } } os.Remove(file1.Name()) @@ -149,8 +223,8 @@ func TestInputFileLoop(t *testing.T) { for i := 0; i < 1000; i++ { input.Read(buf) } - input.Close() + input.Close() os.Remove(file.Name()) } @@ -221,11 +295,9 @@ func (expectedCaptureFile *CaptureFile) PayloadsEqual(other [][]byte) bool { } for i, payload := range other { - if !bytes.Equal(expectedCaptureFile.data[i], payload) { return false } - } return true @@ -307,4 +379,4 @@ func Duplicate(data []byte) (duplicate []byte) { copy(duplicate, data) return -} +} \ No newline at end of file diff --git a/input_s3.go b/input_s3.go index 960708c..9f73ba5 100644 --- a/input_s3.go +++ b/input_s3.go @@ -1,185 +1,12 @@ package main import ( - _ "bufio" - "fmt" - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/session" - "github.com/aws/aws-sdk-go/service/s3" - _ "github.com/aws/aws-sdk-go/service/s3/s3manager" - "io" - "log" - "math/rand" - "os" - "path/filepath" - "strings" - "sort" - "crypto/sha1" - "encoding/hex" ) type S3InputConfig struct { - bufferConfig FileInputConfig + // bufferConfig FileInputConfig bufferPath string region string endpoint string } - -// FileOutput output plugin -type S3Output struct { - pathTemplate string - - buffer *FileInput - session *session.Session - config *S3InputConfig -} - -// NewFileOutput constructor for FileOutput, accepts path -func NewS3Input(pathTemplate string, config *S3InputConfig) *S3Input { - o := new(S3Input) - o.pathTemplate = pathTemplate - o.config = config - - if config.region == "" { - config.region = "us-east-1" - } - - if config.bufferPath == "" { - config.bufferPath = "/tmp" - } - - o.connect() - - if !strings.HasPrefix(pathTemplate, "s3://") { - log.Fatal("S3 path format should be: s3:///") - } - - return o -} - -func (o *S3Output) connect() { - if o.session == nil { - o.session = session.New(&aws.Config{Region: aws.String(o.config.region)}) - } -} - -type sortByS3FileIndex []*s3.Object - -func (s sortByS3FileIndex) Len() int { - return len(s) -} - -func (s sortByS3FileIndex) Swap(i, j int) { - s[i], s[j] = s[j], s[i] -} - -func (s sortByS3FileIndex) Less(i, j int) bool { - if withoutIndex(*s[i].Key) == withoutIndex(*s[j].Key) { - return getFileIndex(*s[i].Key) < getFileIndex(*s[j].Key) - } - - return s[i] < s[j] -} - -func (o *S3Input) updateBuffer() (err error) { - path := o.pathTemplate[5:] // stripping `s3://` - sep := strings.IndexByte(path, '/') - - bucket = path[:sep] - key = path[sep+1:] - - params := &s3.ListObjectsInput{ - Bucket: bucket, - Prefix: key, - } - resp, err := svc.ListObjects(params) - sort.Sort(sortByS3FileIndex(resp.Contents)) - - if err != nil { - return err - } - - if o.buffer.currentFile == nil { - fileToDownload := resp.Contents[0] - } else { - found := false - - bufName := filepath.Base(i.buffer.currentFile.Name()) - bufHex := strings.TrimSuffix(bufName, filepath.Ext(bufName)) - bufSha, _ := hex.DecodeString(bufHex) - - for idx, c := range resp.Contents { - sha := sha1.Sum(*c.Key)[0:] - - if bytes.Equal(bufSha, sha) && idx != len(matches)-1 { - if i.buffer.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(o.pathTemplate, ".gz") { - buffer_name += ".gz" - } - - buffer_path := filepath.Join(o.config.bufferPath, buffer_name) -} - -func (o *S3Output) Read(data []byte) (int, error) { - return o.buffer.Read(data) -} - -func (o *S3Output) String() string { - return "S3 Input: " + o.file.Name() -} - -func (o *S3Output) Close() { - o.buffer.Close() -} - -func (o *S3Output) keyPath(idx int) (bucket, key string) { - path := o.pathTemplate[5:] // stripping `s3://` - sep := strings.IndexByte(path, '/') - - bucket = path[:sep] - key = path[sep+1:] - - for name, fn := range dateFileNameFuncs { - key = strings.Replace(key, name, fn(), -1) - } - - key = setFileIndex(key, idx) - - return -} - -func (o *S3Output) onBufferUpdate(path string) { - svc := s3.New(o.session) - idx := getFileIndex(path) - bucket, key := o.keyPath(idx) - - file, _ := os.Open(path) - // reader := bufio.NewReader(file) - - _, err := svc.PutObject(&s3.PutObjectInput{ - Body: file, - Bucket: aws.String(bucket), - Key: aws.String(key), - }) - if err != nil { - log.Printf("Failed to upload data to %s/%s, %s\n", bucket, key, err) - return - } - - os.Remove(path) -} diff --git a/output_s3.go b/output_s3.go index 11527df..752eff1 100644 --- a/output_s3.go +++ b/output_s3.go @@ -7,7 +7,7 @@ import ( "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/s3" _ "github.com/aws/aws-sdk-go/service/s3/s3manager" - "io" + _ "io" "log" "math/rand" "os" @@ -77,7 +77,7 @@ func (o *S3Output) Write(data []byte) (n int, err error) { } func (o *S3Output) String() string { - return "File output: " + o.file.Name() + return "S3 output: " + o.pathTemplate } func (o *S3Output) Close() { diff --git a/output_s3_test.go b/output_s3_test.go index bc9247b..15693b5 100644 --- a/output_s3_test.go +++ b/output_s3_test.go @@ -8,6 +8,7 @@ import ( "os" "path/filepath" "testing" + "time" ) func TestS3Output(t *testing.T) { @@ -31,6 +32,8 @@ func TestS3Output(t *testing.T) { output.buffer.updateName() output.Write([]byte("1 1 1\ntest")) + time.Sleep(time.Second) + params := &s3.ListObjectsInput{ Bucket: bucket, Prefix: aws.String(fmt.Sprintf("%d", rnd)),