This commit is contained in:
Leonid Bugaev
2016-12-19 20:16:11 +03:00
parent 55fe5b9c11
commit ddf2dc2876
6 changed files with 243 additions and 297 deletions
+1 -1
View File
@@ -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)
+147 -103
View File
@@ -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
}
+89 -17
View File
@@ -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
}
}
+1 -174
View File
@@ -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://<bucket>/<path_format>")
}
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)
}
+2 -2
View File
@@ -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() {
+3
View File
@@ -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)),