From 962087065a6397a0b097a2fa850da7efa273ad06 Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Tue, 22 Jan 2019 23:08:23 +0100 Subject: [PATCH] S3 fixes --- gor.go | 2 ++ input_file.go | 62 ++++++++++++++++++++++++---------------------- input_file_test.go | 7 +++--- middleware.go | 2 +- output_s3.go | 17 ++++++++----- 5 files changed, 51 insertions(+), 39 deletions(-) diff --git a/gor.go b/gor.go index 6823eb2..f035d4c 100644 --- a/gor.go +++ b/gor.go @@ -104,6 +104,8 @@ func finalize() { cp.Close() } } + + time.Sleep(100 * time.Millisecond) } func profileCPU(cpuprofile string) { diff --git a/input_file.go b/input_file.go index 8134bd6..a74a308 100644 --- a/input_file.go +++ b/input_file.go @@ -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 } diff --git a/input_file_test.go b/input_file_test.go index 9ad6de1..573b38f 100644 --- a/input_file_test.go +++ b/input_file_test.go @@ -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) } diff --git a/middleware.go b/middleware.go index 0fa0031..4d6d59a 100644 --- a/middleware.go +++ b/middleware.go @@ -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 { diff --git a/output_s3.go b/output_s3.go index 4e86013..5d43eff 100644 --- a/output_s3.go +++ b/output_s3.go @@ -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 }