Remove double buffering in S3 reader (#17)

Author: murphm8

pkg/stream/s3/s3_reader.go

@@ -35,9 +35,7 @@ func newS3ReaderConfig() s3ReaderConfig {
 //Implements io.ReadSeeker
 type s3Reader struct {
 	config        s3ReaderConfig
-	buffer        []byte
-	bufferStart   int64
-	bufferEnd     int64
+	resp          *s3.GetObjectOutput
 	bucket        string
 	key           string
 	s3            s3iface.S3API
@@ -48,9 +46,6 @@ type s3Reader struct {
 
 func newS3Reader(s3Api s3iface.S3API, bucket string, key string, contentLength int64, etag string, config s3ReaderConfig) *s3Reader {
 	return &s3Reader{
-		buffer:        make([]byte, config.BufferSize),
-		bufferStart:   0,
-		bufferEnd:     0,
 		bucket:        bucket,
 		key:           key,
 		s3:            s3Api,
@@ -115,62 +110,59 @@ func (r *s3Reader) Read(p []byte) (n int, err error) {
 }
 
 func (r *s3Reader) read(p []byte) (n int, err error) {
-	if r.bufferEnd != 0 && r.bufferStart != r.bufferEnd {
-		return r.copyInto(p)
+	if r.resp == nil {
+		getObjectErr := r.makeNewS3Request()
+		if getObjectErr != nil {
+			return 0, getObjectErr
+		}
 	}
 
-	r.resetBuffer()
+	bytesRead, readErr := r.resp.Body.Read(p)
 
-	bufferBytes := int64(len(r.buffer))
-	bytesLeft := r.ContentLength - r.offset
-	bytesToRead := min(bufferBytes, bytesLeft)
+	if readErr != nil {
+		// Error throwing from S3, close the current s3 socket
+		defer r.closeS3Socket()
+		if readErr == io.EOF {
+			r.offset += int64(bytesRead)
+			return bytesRead, io.EOF
+		} else {
+			return bytesRead, &ReadError{readErr}
+		}
+	} else {
+		r.offset += int64(bytesRead)
+		return bytesRead, nil
 
+	}
+}
+
+func (r *s3Reader) makeNewS3Request() (err error) {
 	resp, getObjectErr := r.s3.GetObjectWithContext(
 		aws.BackgroundContext(),
 		&s3.GetObjectInput{
 			Bucket:  aws.String(r.bucket),
 			Key:     aws.String(r.key),
 			IfMatch: aws.String(r.Etag),
-			Range:   aws.String(fmt.Sprintf("bytes=%v-%v", r.offset, r.offset+bytesToRead-1)),
+			// Always open a connection read from current position to end of file
+			Range: aws.String(fmt.Sprintf("bytes=%v-%v", r.offset, r.ContentLength-1)),
 		},
-		request.WithResponseReadTimeout(5*time.Second)) //Sets a timeout on the underlying Body.Read() calls. By default, there is no timeout on this read
+		// Sets a timeout on the underlying Body.Read() calls to avoid a S3 connection leaking
+		// since this S3Reader keep connection open until read end of file. This is a client config,
+		// and S3 Service will also terminate the idle connections, which we will rely on the retry
+		// if two reads duration is too long and s3 service terminate the socket.
+		request.WithResponseReadTimeout(10*time.Second))
 
 	if getObjectErr != nil {
-		return 0, getObjectErr
-	}
-
-	defer resp.Body.Close()
-	for {
-		bytesRead, readErr := resp.Body.Read(r.buffer[r.bufferEnd:])
-		r.bufferEnd += int64(bytesRead)
-
-		if readErr == io.EOF {
-			break
-		}
-
-		if readErr != nil {
-			return n, &ReadError{readErr}
-		}
+		return getObjectErr
 	}
-
-	return r.copyInto(p)
+	r.resp = resp
+	return nil
 }
 
-func (r *s3Reader) copyInto(p []byte) (n int, err error) {
-	n = copy(p, r.buffer[r.bufferStart:r.bufferEnd])
-	r.bufferStart += int64(n)
-
-	r.offset += int64(n)
-	if r.offset == r.ContentLength {
-		return n, io.EOF
+func (r *s3Reader) closeS3Socket() {
+	if r.resp != nil {
+		r.resp.Body.Close()
+		r.resp = nil
 	}
-
-	return n, nil
-}
-
-func (r *s3Reader) resetBuffer() {
-	r.bufferStart = 0
-	r.bufferEnd = 0
 }
 
 func (r *s3Reader) Seek(offset int64, whence int) (newOffset int64, err error) {
@@ -191,10 +183,10 @@ func (r *s3Reader) Seek(offset int64, whence int) (newOffset int64, err error) {
 		r.offset = r.ContentLength
 	}
 
-	//Reset buffer when seeking
-	//Special case: Dont reset if the position hasn't changed
+	//Close the socket when seeking
+	//Special case: Dont close if the position hasn't changed
 	if oldPos != r.offset {
-		r.resetBuffer()
+		r.closeS3Socket()
 	}
 
 	return r.offset, nil

pkg/stream/s3/s3_reader_test.go

@@ -115,44 +115,53 @@ func TestS3Reader_Read_CallFails_ReturnsError(t *testing.T) {
 	assert.True(t, err != nil)
 }
 
-func TestS3Reader_Read_ShouldBuffer(t *testing.T) {
+func TestS3Reader_Seek_SeeksToCorrectPosition(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 	mockS3Client := NewMockS3API(ctrl)
-	setupS3MockExpects(mockS3Client)
 
-	mockS3Client.EXPECT().GetObjectWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(&s3.GetObjectOutput{
-		Body: ioutil.NopCloser(strings.NewReader(testBodyContent[2:4])),
+	mockS3Client.EXPECT().HeadObject(gomock.Any()).Return(&s3.HeadObjectOutput{
+		ContentLength: aws.Int64(int64(len(testBodyContent))),
+		ETag:          aws.String(testEtag),
 	}, nil).Times(1)
 
-	config := newS3ReaderConfig()
-	config.BufferSize = 2
-	s3Reader, _ := newS3ReaderWithConfig(mockS3Client, testBucket, testKey, config)
+	s3Reader, _ := newS3ReaderBucketAndKey(mockS3Client, testBucket, testKey)
+	s3Reader.Seek(5, io.SeekCurrent)
+	assert.True(t, s3Reader.offset == 5)
 
-	//First call will read from s3 since the buffer is empty.
-	//Second call should read from internal buffer
-	//Third call should read from s3 again
-	s3Reader.Read(make([]byte, 1))
-	s3Reader.Read(make([]byte, 1))
-	s3Reader.Read(make([]byte, 1))
+	s3Reader.Seek(1, io.SeekCurrent)
+	assert.True(t, s3Reader.offset == 6)
+
+	s3Reader.Seek(0, io.SeekStart)
+	assert.True(t, s3Reader.offset == 0)
+
+	s3Reader.Seek(0, io.SeekEnd)
+	assert.True(t, s3Reader.offset == s3Reader.ContentLength)
+
+	s3Reader.Seek(100000, io.SeekCurrent)
+	assert.True(t, s3Reader.offset == s3Reader.ContentLength)
 }
 
-func TestS3Reader_Seek_ResetsBuffer(t *testing.T) {
+func TestS3Reader_Read_ShouldKeepS3SocketOpen(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 	mockS3Client := NewMockS3API(ctrl)
 	setupS3MockExpects(mockS3Client)
-	s3Reader, _ := newS3ReaderBucketAndKey(mockS3Client, testBucket, testKey)
+	config := newS3ReaderConfig()
+	s3Reader, _ := newS3ReaderWithConfig(mockS3Client, testBucket, testKey, config)
 
-	//Read to fill the buffer
+	//First call will read from s3 since the buffer is empty.
+	//Second call should reuse the existing s3 socket
+	//Third call should reuse the existing s3 socket
+	s3Reader.Read(make([]byte, 1))
+	s3Reader.Read(make([]byte, 1))
 	s3Reader.Read(make([]byte, 1))
-	assert.True(t, s3Reader.bufferEnd != 0)
 
-	s3Reader.Seek(1, io.SeekCurrent)
-	assert.True(t, s3Reader.bufferEnd == 0)
+	//After finish reading, socket closed
+	assert.Nil(t, s3Reader.resp)
 }
 
-func TestS3Reader_SeekNoop_DoesNotResetBuffer(t *testing.T) {
+func TestS3Reader_Seek_S3SocketClose(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 	mockS3Client := NewMockS3API(ctrl)
@@ -161,34 +170,25 @@ func TestS3Reader_SeekNoop_DoesNotResetBuffer(t *testing.T) {
 
 	//Read to fill the buffer
 	s3Reader.Read(make([]byte, 1))
-	assert.True(t, s3Reader.bufferEnd != 0)
+	assert.NotNil(t, s3Reader.resp)
 
-	s3Reader.Seek(0, io.SeekCurrent)
-	assert.True(t, s3Reader.bufferEnd != 0)
+	//seek and s3 socket close
+	s3Reader.Seek(0, io.SeekStart)
+	assert.Nil(t, s3Reader.resp)
 }
 
-func TestS3Reader_Seek_SeeksToCorrectPosition(t *testing.T) {
+func TestS3Reader_SeekNoop_DoesNotCloseSocket(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 	mockS3Client := NewMockS3API(ctrl)
-	mockS3Client.EXPECT().HeadObject(gomock.Any()).Return(&s3.HeadObjectOutput{
-		ContentLength: aws.Int64(int64(len(testBodyContent))),
-		ETag:          aws.String(testEtag),
-	}, nil).Times(1)
+	setupS3MockExpects(mockS3Client)
 	s3Reader, _ := newS3ReaderBucketAndKey(mockS3Client, testBucket, testKey)
 
-	s3Reader.Seek(5, io.SeekCurrent)
-	assert.True(t, s3Reader.offset == 5)
-
-	s3Reader.Seek(1, io.SeekCurrent)
-	assert.True(t, s3Reader.offset == 6)
-
-	s3Reader.Seek(0, io.SeekStart)
-	assert.True(t, s3Reader.offset == 0)
-
-	s3Reader.Seek(0, io.SeekEnd)
-	assert.True(t, s3Reader.offset == s3Reader.ContentLength)
+	//Read to fill the buffer
+	s3Reader.Read(make([]byte, 1))
+	assert.NotNil(t, s3Reader.resp)
 
-	s3Reader.Seek(100000, io.SeekCurrent)
-	assert.True(t, s3Reader.offset == s3Reader.ContentLength)
+	//seek and s3 socket still open
+	s3Reader.Seek(1, io.SeekStart)
+	assert.NotNil(t, s3Reader.resp)
 }