mirror of
https://github.com/buger/goreplay.git
synced 2024-04-21 12:32:02 +00:00
Refactor #308
This commit is contained in:
Leonid Bugaev
+127
-116
@@ -9,35 +9,109 @@ import (
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// FileInput can read requests generated by FileOutput
|
||||
type FileInput struct {
|
||||
data chan []byte
|
||||
exit chan bool
|
||||
path string
|
||||
fileInputReaders []*fileInputReader
|
||||
speedFactor float64
|
||||
loop bool
|
||||
}
|
||||
|
||||
|
||||
type fileInputReader struct {
|
||||
reader *bufio.Reader
|
||||
meta [][]byte
|
||||
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 {
|
||||
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.exit = make(chan bool)
|
||||
i.data = make(chan []byte, 1000)
|
||||
i.exit = make(chan bool, 1)
|
||||
i.path = path
|
||||
i.speedFactor = 1
|
||||
i.loop = loop
|
||||
@@ -58,6 +132,9 @@ func (_ *NextFileNotFound) Error() string {
|
||||
}
|
||||
|
||||
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 {
|
||||
@@ -70,24 +147,10 @@ func (i *FileInput) init() (err error) {
|
||||
return errors.New("No matching files")
|
||||
}
|
||||
|
||||
i.fileInputReaders = make([]*fileInputReader, len(matches))
|
||||
i.readers = make([]*fileInputReader, len(matches))
|
||||
|
||||
for idx, p := range matches {
|
||||
file, _ := os.Open(p)
|
||||
fileInputReader := &fileInputReader{}
|
||||
fileInputReader.file = file
|
||||
if strings.HasSuffix(p, ".gz") {
|
||||
gzReader, err := gzip.NewReader(file)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
fileInputReader.reader = bufio.NewReader(gzReader)
|
||||
} else {
|
||||
fileInputReader.reader = bufio.NewReader(file)
|
||||
}
|
||||
|
||||
fileInputReader.readNextInput()
|
||||
i.fileInputReaders[idx] = fileInputReader
|
||||
i.readers[idx] = NewFileInputReader(p)
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -104,123 +167,71 @@ func (i *FileInput) String() string {
|
||||
return "File input: " + i.path
|
||||
}
|
||||
|
||||
func (f *fileInputReader) readNextInput() {
|
||||
nextInput := f.nextInput()
|
||||
f.parseNextInput(nextInput)
|
||||
}
|
||||
|
||||
func (f *fileInputReader) parseNextInput(input []byte) {
|
||||
if (input != nil) {
|
||||
f.meta = payloadMeta(input)
|
||||
f.timestamp, _ = strconv.ParseInt(string(f.meta[2]), 10, 64)
|
||||
f.data = input
|
||||
}
|
||||
}
|
||||
|
||||
func (f *fileInputReader) nextInput() []byte {
|
||||
payloadSeparatorAsBytes := []byte(payloadSeparator)
|
||||
var buffer bytes.Buffer
|
||||
|
||||
for {
|
||||
line, err := f.reader.ReadBytes('\n')
|
||||
|
||||
if err != nil {
|
||||
if err != io.EOF {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
if err == io.EOF {
|
||||
f.file.Close()
|
||||
f.file = nil
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
if bytes.Equal(payloadSeparatorAsBytes[1:], line) {
|
||||
asBytes := buffer.Bytes()
|
||||
|
||||
// Bytes() returns only pointer, so to remove data-race copy the data to an array
|
||||
newBuf := make([]byte, len(asBytes) - 1)
|
||||
copy(newBuf, asBytes)
|
||||
return newBuf
|
||||
}
|
||||
|
||||
buffer.Write(line)
|
||||
}
|
||||
}
|
||||
|
||||
func (i *FileInput) nextInputReader() *fileInputReader {
|
||||
var nextFileInputReader *fileInputReader
|
||||
for _, fileInputReader := range i.fileInputReaders {
|
||||
if fileInputReader.file == nil {
|
||||
// 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
|
||||
}
|
||||
|
||||
if fileInputReader.meta[0][0] == ResponsePayload {
|
||||
return fileInputReader
|
||||
}
|
||||
|
||||
if nextFileInputReader == nil || nextFileInputReader.timestamp > fileInputReader.timestamp {
|
||||
nextFileInputReader = fileInputReader
|
||||
if next == nil || r.timestamp < next.timestamp {
|
||||
next = r
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
return nextFileInputReader;
|
||||
return
|
||||
}
|
||||
|
||||
func (i *FileInput) emit() {
|
||||
var lastTime int64 = -1
|
||||
|
||||
for {
|
||||
fileInputReader := i.nextInputReader()
|
||||
select {
|
||||
case <-i.exit:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
if fileInputReader == nil {
|
||||
reader := i.nextReader()
|
||||
|
||||
if reader == nil {
|
||||
if i.loop {
|
||||
i.init()
|
||||
lastTime = -1
|
||||
continue
|
||||
} else {
|
||||
break;
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if fileInputReader.meta[0][0] == RequestPayload {
|
||||
lastTime = i.simulateRequestDelay(fileInputReader, lastTime)
|
||||
if lastTime != -1 {
|
||||
diff := reader.timestamp - lastTime
|
||||
lastTime = reader.timestamp
|
||||
|
||||
if i.speedFactor != 1 {
|
||||
diff = int64(float64(diff) / i.speedFactor)
|
||||
}
|
||||
|
||||
time.Sleep(time.Duration(diff))
|
||||
} else {
|
||||
lastTime = reader.timestamp
|
||||
}
|
||||
|
||||
select {
|
||||
case <-i.exit:
|
||||
for _, fileInputReader := range i.fileInputReaders {
|
||||
if fileInputReader.file != nil {
|
||||
fileInputReader.file.Close()
|
||||
}
|
||||
}
|
||||
break
|
||||
case i.data <- fileInputReader.data:
|
||||
fileInputReader.readNextInput()
|
||||
}
|
||||
i.data <- reader.ReadPayload()
|
||||
}
|
||||
|
||||
log.Printf("FileInput: end of file '%s'\n", i.path)
|
||||
}
|
||||
|
||||
func (i*FileInput) simulateRequestDelay(fileInputReader *fileInputReader, lastTime int64) int64 {
|
||||
if lastTime != -1 {
|
||||
timeDiff := fileInputReader.timestamp - lastTime
|
||||
func (i *FileInput) Close() error {
|
||||
defer i.mu.Unlock()
|
||||
i.mu.Lock()
|
||||
|
||||
if i.speedFactor != 1 {
|
||||
timeDiff = int64(float64(timeDiff) / i.speedFactor)
|
||||
}
|
||||
i.exit <- true
|
||||
|
||||
time.Sleep(time.Duration(timeDiff))
|
||||
r.Close()
|
||||
}
|
||||
|
||||
return fileInputReader.timestamp
|
||||
}
|
||||
|
||||
func (i *FileInput) Close() error {
|
||||
i.exit <- true
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user