github/goreplay
4
0
Fork 0
mirror of https://github.com/buger/goreplay.git synced 2024-04-21 12:32:02 +00:00
Code Issues Releases Wiki Activity

S3 fixes

Browse Source
This commit is contained in:
Leonid Bugaev
2019-01-22 23:08:23 +01:00
parent 166c8e8d87
commit 962087065a
5 changed files with 51 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gor.go
+2
View File
@@ -104,6 +104,8 @@ func finalize() {
cp.Close()
}
}
time.Sleep(100 * time.Millisecond)
}
func profileCPU(cpuprofile string) {
input_file.go
+33 -29
View File
@@ -19,10 +19,13 @@ import (
)
type S3ReadCloser struct {
bucket string
key string
offset int
sess *session.Session
bucket string
key string
offset int
totalSize int
readBytes int
sess *session.Session
buf *bytes.Buffer
}
func awsConfig() *aws.Config {
@@ -42,43 +45,44 @@ func awsConfig() *aws.Config {
func NewS3ReadCloser(path string) *S3ReadCloser {
bucket, key := parseS3Url(path)
sess := session.Must(session.NewSession(awsConfig()))
log.Println("[S3 Input] S3 connection succesfully initialized", path)
return &S3ReadCloser{
bucket: bucket,
key: key,
sess: session.New(awsConfig()),
sess: sess,
buf: &bytes.Buffer{},
}
}
func (s *S3ReadCloser) Read(b []byte) (n int, e error) {
svc := s3.New(s.sess)
if s.readBytes == 0 || s.readBytes+len(b) > s.offset {
svc := s3.New(s.sess)
objectRange := "bytes=" + strconv.Itoa(s.offset)
s.offset += 1000000 // Reading in chunks of 1 mb
objectRange += "-" + strconv.Itoa(s.offset-1)
objectRange := "bytes=" + strconv.Itoa(s.offset)
s.offset += 1000000 // Reading in chunks of 1 mb
objectRange += "-" + strconv.Itoa(s.offset-1)
params := &s3.GetObjectInput{
Bucket: aws.String(s.bucket),
Key: aws.String(s.key),
Range: aws.String(objectRange),
}
resp, err := svc.GetObject(params)
params := &s3.GetObjectInput{
Bucket: aws.String(s.bucket),
Key: aws.String(s.key),
Range: aws.String(objectRange),
}
resp, err := svc.GetObject(params)
if err != nil {
return 0, err
if err != nil {
log.Println("[S3 Input] Error during getting file", s.bucket, s.key, err)
} else {
s.totalSize, _ = strconv.Atoi(strings.Split(*resp.ContentRange, "/")[1])
s.buf.ReadFrom(resp.Body)
}
}
totalSize, _ := strconv.Atoi(strings.Split(*resp.ContentRange, "/")[1])
s.readBytes += len(b)
n, e = resp.Body.Read(b)
// S3 always return EOF when reading byte range
// It is not actually error, until we reached end of file
if n > 0 && e != nil && s.offset < totalSize {
e = nil
}
return n, e
return s.buf.Read(b)
}
func (s *S3ReadCloser) Close() error {
@@ -216,7 +220,7 @@ func (i *FileInput) init() (err error) {
var matches []string
if strings.HasPrefix(i.path, "s3://") {
sess := session.New(awsConfig())
sess := session.Must(session.NewSession(awsConfig()))
svc := s3.New(sess)
bucket, key := parseS3Url(i.path)
@@ -228,7 +232,7 @@ func (i *FileInput) init() (err error) {
resp, err := svc.ListObjects(params)
if err != nil {
log.Println("Error while retreiving list of lies from S3", i.path, err)
log.Println("Error while retreiving list of files from S3", i.path, err)
return err
}
input_file_test.go
+4 -3
View File
@@ -388,15 +388,16 @@ func TestInputFileFromS3(t *testing.T) {
path := fmt.Sprintf("s3://test-gor/%d/requests.gz", rnd)
output := NewS3Output(path, &FileOutputConfig{queueLimit: 5000})
output.closeC = make(chan struct{}, 2)
output.closeC = make(chan struct{}, 10)
for i := 0; i <= 10000; i++ {
for i := 0; i <= 20000; i++ {
output.Write([]byte("1 1 1\ntest"))
if i%5000 == 0 {
output.buffer.updateName()
}
}
output.Write([]byte("1 1 1\ntest"))
for i := 0; i < 2; i++ {
@@ -406,7 +407,7 @@ func TestInputFileFromS3(t *testing.T) {
input := NewFileInput(fmt.Sprintf("s3://test-gor/%d", rnd), false)
buf := make([]byte, 1000)
for i := 0; i <= 10000; i++ {
for i := 0; i <= 19999; i++ {
input.Read(buf)
}
middleware.go
+1 -1
View File
@@ -115,7 +115,7 @@ func (m *Middleware) read(from io.Reader) {
buf := make([]byte, len(line)/2)
if _, err := hex.Decode(buf, line[:len(line)-1]); err != nil {
fmt.Fprintln(os.Stderr, "Failed to decode input payload", err, len(line))
fmt.Fprintln(os.Stderr, "Failed to decode input payload", err, len(line), string(line[:len(line)-1]))
}
if Settings.debug {
output_s3.go
+11 -6
View File
@@ -15,7 +15,7 @@ import (
"strings"
)
// FileOutput output plugin
// S3Output output plugin
type S3Output struct {
pathTemplate string
@@ -25,7 +25,7 @@ type S3Output struct {
closeC chan struct{}
}
// NewFileOutput constructor for FileOutput, accepts path
// NewS3Output constructor for FileOutput, accepts path
func NewS3Output(pathTemplate string, config *FileOutputConfig) *S3Output {
o := new(S3Output)
o.pathTemplate = pathTemplate
@@ -37,7 +37,10 @@ func NewS3Output(pathTemplate string, config *FileOutputConfig) *S3Output {
}
rnd := rand.Int63()
buffer_name := fmt.Sprintf("gor_output_s3_%d_buf", rnd)
buffer_name := fmt.Sprintf("gor_output_s3_%d_buf_", rnd)
pathParts := strings.Split(pathTemplate, "/")
buffer_name += pathParts[len(pathParts)-1]
if strings.HasSuffix(o.pathTemplate, ".gz") {
buffer_name += ".gz"
@@ -54,6 +57,7 @@ func NewS3Output(pathTemplate string, config *FileOutputConfig) *S3Output {
func (o *S3Output) connect() {
if o.session == nil {
o.session = session.Must(session.NewSession(awsConfig()))
log.Println("[S3 Output] S3 connection succesfully initialized")
}
}
@@ -65,8 +69,8 @@ func (o *S3Output) String() string {
return "S3 output: " + o.pathTemplate
}
func (o *S3Output) Close() {
o.buffer.Close()
func (o *S3Output) Close() error {
return o.buffer.Close()
}
func parseS3Url(path string) (bucket, key string) {
@@ -105,7 +109,8 @@ func (o *S3Output) onBufferUpdate(path string) {
Key: aws.String(key),
})
if err != nil {
log.Printf("Failed to upload data to %s/%s, %s\n", bucket, key, err)
log.Printf("[S3 Output] Failed to upload data to %s/%s, %s\n", bucket, key, err)
os.Remove(path)
return
}