Repository: uber/cadence (go)

Issue:
Multipart Version of S3 History Archiver
Currently the S3 implementation of history archiver interface will always load the entire workflow history into memory and upload only one file. This approach has two drawbacks for large workflows:
1. Load full history into memory may cause a memory issue.
2. If the any error happens during the upload, the entire history needs to be uploaded again.

A better idea is split workflow history into multiple part, upload them piece by piece and checkpoint the progress.

Pull request description:
Implement multipart s3 history archiver.
Fixes #3002.
Remove MaxRetries, this was a left over from the old interface

Patch:
diff --git a/common/archiver/s3store/historyArchiver.go b/common/archiver/s3store/historyArchiver.go
index 0e710c44c..7e9facd4b 100644
--- a/common/archiver/s3store/historyArchiver.go
+++ b/common/archiver/s3store/historyArchiver.go
@@ -30,8 +30,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/uber/cadence/common/metrics"
-
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/request"
@@ -44,6 +42,8 @@ import (
 	"github.com/uber/cadence/common/archiver"
 	"github.com/uber/cadence/common/backoff"
 	"github.com/uber/cadence/common/log/tag"
+	"github.com/uber/cadence/common/metrics"
+	"github.com/uber/cadence/common/persistence"
 	"github.com/uber/cadence/common/service/config"
 )
 
@@ -73,7 +73,14 @@ type (
 
 	getHistoryToken struct {
 		CloseFailoverVersion int64
-		NextBatchIdx         int
+		BatchIdx             int
+	}
+
+	uploadProgress struct {
+		BatchIdx      int
+		IteratorState []byte
+		uploadedSize  int64
+		historySize   int64
 	}
 )
 
@@ -97,7 +104,6 @@ func newHistoryArchiver(
 		Endpoint:         config.Endpoint,
 		Region:           aws.String(config.Region),
 		S3ForcePathStyle: aws.Bool(config.S3ForcePathStyle),
-		MaxRetries:       aws.Int(0),
 	}
 	sess, err := session.NewSession(s3Config)
 	if err != nil {
@@ -145,12 +151,11 @@ func (h *historyArchiver) Archive(
 		return err
 	}
 
+	var progress uploadProgress
 	historyIterator := h.historyIterator
 	if historyIterator == nil { // will only be set by testing code
-		historyIterator = archiver.NewHistoryIterator(request, h.container.HistoryV2Manager, targetHistoryBlobSize)
+		historyIterator = loadHistoryIterator(ctx, request, h.container.HistoryV2Manager, featureCatalog, &progress)
 	}
-
-	historyBatches := []*shared.History{}
 	for historyIterator.HasNext() {
 		historyBlob, err := getNextHistoryBlob(ctx, historyIterator)
 		if err != nil {
@@ -168,29 +173,86 @@ func (h *historyArchiver) Archive(
 			return archiver.ErrHistoryMutated
 		}
 
-		historyBatches = append(historyBatches, historyBlob.Body...)
-	}
+		encodedHistoryBlob, err := encode(historyBlob)
+		if err != nil {
+			logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(errEncodeHistory), tag.Error(err))
+			return err
+		}
 
-	encodedHistoryBatches, err := encode(historyBatches)
-	if err != nil {
-		logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(errEncodeHistory), tag.Error(err))
-		return err
-	}
+		key := constructHistoryKey(URI.Path(), request.DomainID, request.WorkflowID, request.RunID, request.CloseFailoverVersion, progress.BatchIdx)
 
-	key := constructHistoryKey(URI.Path(), request.DomainID, request.WorkflowID, request.RunID, request.CloseFailoverVersion)
+		exists, err := keyExists(ctx, h.s3cli, URI, key)
+		if err != nil {
+			logger := logger.WithTags(tag.ArchivalArchiveFailReason(errWriteKey), tag.Error(err))
+			if isRetryableError(err) {
+				logger.Error(archiver.ArchiveTransientErrorMsg)
+			} else {
+				logger.Error(archiver.ArchiveNonRetriableErrorMsg)
+			}
+			return err
+		}
+		blobSize := int64(binary.Size(encodedHistoryBlob))
+		if exists {
+			scope.IncCounter(metrics.HistoryArchiverBlobExistsCount)
+		} else {
+			if err := upload(ctx, h.s3cli, URI, key, encodedHistoryBlob); err != nil {
+				logger := logger.WithTags(tag.ArchivalArchiveFailReason(errWriteKey), tag.Error(err))
+				if isRetryableError(err) {
+					logger.Error(archiver.ArchiveTransientErrorMsg)
+				} else {
+					logger.Error(archiver.ArchiveNonRetriableErrorMsg)
+				}
+				return err
+			}
+			progress.uploadedSize += blobSize
+			scope.RecordTimer(metrics.HistoryArchiverBlobSize, time.Duration(blobSize))
+		}
 
-	if err := upload(ctx, h.s3cli, URI, key, encodedHistoryBatches); err != nil {
-		logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(errWriteKey), tag.Error(err))
-		return err
+		progress.historySize += blobSize
+		progress.BatchIdx = progress.BatchIdx + 1
+		saveHistoryIteratorState(ctx, featureCatalog, historyIterator, &progress)
 	}
-	totalUploadSize := int64(binary.Size(encodedHistoryBatches))
-	scope.AddCounter(metrics.HistoryArchiverTotalUploadSize, totalUploadSize)
-	scope.AddCounter(metrics.HistoryArchiverHistorySize, totalUploadSize)
-	scope.IncCounter(metrics.HistoryArchiverArchiveSuccessCount)
 
+	scope.RecordTimer(metrics.HistoryArchiverTotalUploadSize, time.Duration(progress.uploadedSize))
+	scope.RecordTimer(metrics.HistoryArchiverHistorySize, time.Duration(progress.historySize))
+	scope.IncCounter(metrics.HistoryArchiverArchiveSuccessCount)
 	return nil
 }
 
+func loadHistoryIterator(ctx context.Context, request *archiver.ArchiveHistoryRequest, historyManager persistence.HistoryManager, featureCatalog *archiver.ArchiveFeatureCatalog, progress *uploadProgress) (historyIterator archiver.HistoryIterator) {
+	if featureCatalog.ProgressManager != nil {
+		if featureCatalog.ProgressManager.HasProgress(ctx) {
+			err := featureCatalog.ProgressManager.LoadProgress(ctx, progress)
+			if err == nil {
+				historyIterator, err := archiver.NewHistoryIteratorFromState(request, historyManager, targetHistoryBlobSize, progress.IteratorState)
+				if err == nil {
+					return historyIterator
+				}
+			}
+			progress.IteratorState = nil
+			progress.BatchIdx = 0
+			progress.historySize = 0
+			progress.uploadedSize = 0
+		}
+	}
+	return archiver.NewHistoryIterator(request, historyManager, targetHistoryBlobSize)
+}
+
+func saveHistoryIteratorState(ctx context.Context, featureCatalog *archiver.ArchiveFeatureCatalog, historyIterator archiver.HistoryIterator, progress *uploadProgress) {
+	// Saving history state is a best effort operation. Ignore errors and continue
+	if featureCatalog.ProgressManager != nil {
+		state, err := historyIterator.GetState()
+		if err != nil {
+			return
+		}
+		progress.IteratorState = state
+		err = featureCatalog.ProgressManager.RecordProgress(ctx, progress)
+		if err != nil {
+			return
+		}
+	}
+}
+
 func (h *historyArchiver) Get(
 	ctx context.Context,
 	URI archiver.URI,
@@ -214,7 +276,6 @@ func (h *historyArchiver) Get(
 	} else if request.CloseFailoverVersion != nil {
 		token = &getHistoryToken{
 			CloseFailoverVersion: *request.CloseFailoverVersion,
-			NextBatchIdx:         0,
 		}
 	} else {
 		highestVersion, err := h.getHighestVersion(ctx, URI, request)
@@ -223,35 +284,49 @@ func (h *historyArchiver) Get(
 		}
 		token = &getHistoryToken{
 			CloseFailoverVersion: *highestVersion,
-			NextBatchIdx:         0,
 		}
 	}
 
-	key := constructHistoryKey(URI.Path(), request.DomainID, request.WorkflowID, request.RunID, token.CloseFailoverVersion)
-	encodedHistoryBatches, err := download(ctx, h.s3cli, URI, key)
-	if err != nil {
-		return nil, err
-	}
-	historyBatches, err := decodeHistoryBatches(encodedHistoryBatches)
-	if err != nil {
-		return nil, &shared.InternalServiceError{Message: err.Error()}
-	}
-	historyBatches = historyBatches[token.NextBatchIdx:]
-
 	response := &archiver.GetHistoryResponse{}
 	numOfEvents := 0
-	numOfBatches := 0
-	for _, batch := range historyBatches {
-		response.HistoryBatches = append(response.HistoryBatches, batch)
-		numOfBatches++
-		numOfEvents += len(batch.Events)
+	isTruncated := false
+	for {
 		if numOfEvents >= request.PageSize {
+			isTruncated = true
+			break
+		}
+		key := constructHistoryKey(URI.Path(), request.DomainID, request.WorkflowID, request.RunID, token.CloseFailoverVersion, token.BatchIdx)
+
+		encodedRecord, err := download(ctx, h.s3cli, URI, key)
+		if err != nil {
+			if isRetryableError(err) {
+				return nil, &shared.InternalServiceError{Message: err.Error()}
+			}
+			switch err.(type) {
+			case *shared.BadRequestError, *shared.InternalServiceError, *shared.EntityNotExistsError:
+				return nil, err
+			default:
+				return nil, &shared.InternalServiceError{Message: err.Error()}
+			}
+		}
+
+		historyBlob, err := decodeHistoryBlob(encodedRecord)
+		if err != nil {
+			return nil, &shared.InternalServiceError{Message: err.Error()}
+		}
+
+		for _, batch := range historyBlob.Body {
+			response.HistoryBatches = append(response.HistoryBatches, batch)
+			numOfEvents += len(batch.Events)
+		}
+
+		if *historyBlob.Header.IsLast {
 			break
 		}
+		token.BatchIdx++
 	}
 
-	if numOfBatches < len(historyBatches) {
-		token.NextBatchIdx += numOfBatches
+	if isTruncated {
 		nextToken, err := serializeToken(token)
 		if err != nil {
 			return nil, &shared.InternalServiceError{Message: err.Error()}
@@ -293,8 +368,9 @@ func (h *historyArchiver) getHighestVersion(ctx context.Context, URI archiver.UR
 	defer cancel()
 	var prefix = constructHistoryKeyPrefix(URI.Path(), request.DomainID, request.WorkflowID, request.RunID) + "/"
 	results, err := h.s3cli.ListObjectsV2WithContext(ctx, &s3.ListObjectsV2Input{
-		Bucket: aws.String(URI.Hostname()),
-		Prefix: aws.String(prefix),
+		Bucket:    aws.String(URI.Hostname()),
+		Prefix:    aws.String(prefix),
+		Delimiter: aws.String("/"),
 	})
 	if err != nil {
 		if aerr, ok := err.(awserr.Error); ok && aerr.Code() == s3.ErrCodeNoSuchBucket {
@@ -304,9 +380,9 @@ func (h *historyArchiver) getHighestVersion(ctx context.Context, URI archiver.UR
 	}
 	var highestVersion *int64
 
-	for _, v := range results.Contents {
+	for _, v := range results.CommonPrefixes {
 		var version int64
-		version, err = strconv.ParseInt(strings.Replace(*v.Key, prefix, "", 1), 10, 64)
+		version, err = strconv.ParseInt(strings.Replace(strings.Replace(*v.Prefix, prefix, "", 1), "/", "", 1), 10, 64)
 		if err != nil {
 			continue
 		}
diff --git a/common/archiver/s3store/util.go b/common/archiver/s3store/util.go
index bd2c73fc4..176558476 100644
--- a/common/archiver/s3store/util.go
+++ b/common/archiver/s3store/util.go
@@ -47,13 +47,13 @@ func encode(v interface{}) ([]byte, error) {
 	return json.Marshal(v)
 }
 
-func decodeHistoryBatches(data []byte) ([]*shared.History, error) {
-	historyBatches := []*shared.History{}
-	err := json.Unmarshal(data, &historyBatches)
+func decodeHistoryBlob(data []byte) (*archiver.HistoryBlob, error) {
+	historyBlob := &archiver.HistoryBlob{}
+	err := json.Unmarshal(data, historyBlob)
 	if err != nil {
 		return nil, err
 	}
-	return historyBatches, nil
+	return historyBlob, nil
 }
 func decodeVisibilityRecord(data []byte) (*visibilityRecord, error) {
 	record := &visibilityRecord{}
@@ -105,18 +105,42 @@ func bucketExists(ctx context.Context, s3cli s3iface.S3API, URI archiver.URI) er
 	if err == nil {
 		return nil
 	}
-	if aerr, ok := err.(awserr.Error); ok {
-		if aerr.Code() == "NotFound" {
-			return errBucketNotExists
-		}
+	if isNotFoundError(err) {
+		return errBucketNotExists
 	}
 	return err
 }
 
+func keyExists(ctx context.Context, s3cli s3iface.S3API, URI archiver.URI, key string) (bool, error) {
+	ctx, cancel := ensureContextTimeout(ctx)
+	defer cancel()
+	_, err := s3cli.HeadObjectWithContext(ctx, &s3.HeadObjectInput{
+		Bucket: aws.String(URI.Hostname()),
+		Key:    aws.String(key),
+	})
+	if err != nil {
+		if isNotFoundError(err) {
+			return false, nil
+		}
+		return false, err
+	}
+	return true, nil
+}
+
+func isNotFoundError(err error) bool {
+	aerr, ok := err.(awserr.Error)
+	return ok && (aerr.Code() == "NotFound")
+}
+
 // Key construction
-func constructHistoryKey(path, domainID, workflowID, runID string, version int64) string {
+func constructHistoryKey(path, domainID, workflowID, runID string, version int64, batchIdx int) string {
+	prefix := constructHistoryKeyPrefixWithVersion(path, domainID, workflowID, runID, version)
+	return fmt.Sprintf("%s%d", prefix, batchIdx)
+}
+
+func constructHistoryKeyPrefixWithVersion(path, domainID, workflowID, runID string, version int64) string {
 	prefix := constructHistoryKeyPrefix(path, domainID, workflowID, runID)
-	return fmt.Sprintf("%s/%v", prefix, version)
+	return fmt.Sprintf("%s/%v/", prefix, version)
 }
 
 func constructHistoryKeyPrefix(path, domainID, workflowID, runID string) string {
@@ -193,7 +217,7 @@ func download(ctx context.Context, s3cli s3iface.S3API, URI archiver.URI, key st
 			}
 
 			if aerr.Code() == s3.ErrCodeNoSuchKey {
-				return nil, &shared.BadRequestError{Message: archiver.ErrHistoryNotExist.Error()}
+				return nil, &shared.EntityNotExistsError{Message: archiver.ErrHistoryNotExist.Error()}
 			}
 		}
 		return nil, err
diff --git a/common/archiver/s3store/visibilityArchiver.go b/common/archiver/s3store/visibilityArchiver.go
index 7b5962271..6b3505cfc 100644
--- a/common/archiver/s3store/visibilityArchiver.go
+++ b/common/archiver/s3store/visibilityArchiver.go
@@ -83,7 +83,6 @@ func newVisibilityArchiver(
 		Endpoint:         config.Endpoint,
 		Region:           aws.String(config.Region),
 		S3ForcePathStyle: aws.Bool(config.S3ForcePathStyle),
-		MaxRetries:       aws.Int(0),
 	}
 	sess, err := session.NewSession(s3Config)
 	if err != nil {


Test patch:
diff --git a/common/archiver/s3store/historyArchiver_test.go b/common/archiver/s3store/historyArchiver_test.go
index 9eceb2287..b9545050a 100644
--- a/common/archiver/s3store/historyArchiver_test.go
+++ b/common/archiver/s3store/historyArchiver_test.go
@@ -32,28 +32,24 @@ import (
 	"testing"
 	"time"
 
-	"github.com/uber-go/tally"
-
-	"github.com/uber/cadence/common/metrics"
-
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/s3"
-	"github.com/stretchr/testify/mock"
-
-	"github.com/uber/cadence/common/archiver/s3store/mocks"
-
 	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
+	"github.com/uber-go/tally"
 	"go.uber.org/zap"
 
 	"github.com/uber/cadence/.gen/go/shared"
 	"github.com/uber/cadence/common"
 	"github.com/uber/cadence/common/archiver"
+	"github.com/uber/cadence/common/archiver/s3store/mocks"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/loggerimpl"
+	"github.com/uber/cadence/common/metrics"
 )
 
 const (
@@ -79,8 +75,8 @@ type historyArchiverSuite struct {
 	container          *archiver.HistoryBootstrapContainer
 	logger             log.Logger
 	testArchivalURI    archiver.URI
-	historyBatchesV1   []*shared.History
-	historyBatchesV100 []*shared.History
+	historyBatchesV1   []*archiver.HistoryBlob
+	historyBatchesV100 []*archiver.HistoryBlob
 }
 
 func TestHistoryArchiverSuite(t *testing.T) {
@@ -127,13 +123,28 @@ func setupFsEmulation(s3cli *mocks.S3API) {
 	s3cli.On("ListObjectsV2WithContext", mock.Anything, mock.Anything).
 		Return(func(_ context.Context, input *s3.ListObjectsV2Input, opts ...request.Option) *s3.ListObjectsV2Output {
 			objects := make([]*s3.Object, 0)
+			commonPrefixMap := map[string]bool{}
 			for k := range fs {
 				if strings.HasPrefix(k, *input.Bucket+*input.Prefix) {
-					objects = append(objects, &s3.Object{
-						Key: aws.String(k[len(*input.Bucket):]),
-					})
+					key := k[len(*input.Bucket):]
+					keyWithoutPrefix := key[len(*input.Prefix):]
+					index := strings.Index(keyWithoutPrefix, "/")
+					if index == -1 || input.Delimiter == nil {
+						objects = append(objects, &s3.Object{
+							Key: aws.String(key),
+						})
+					} else {
+						commonPrefixMap[key[:len(*input.Prefix)+index]] = true
+					}
 				}
 			}
+			commonPrefixes := make([]*s3.CommonPrefix, 0)
+			for k := range commonPrefixMap {
+				commonPrefixes = append(commonPrefixes, &s3.CommonPrefix{
+					Prefix: aws.String(k),
+				})
+			}
+
 			sort.SliceStable(objects, func(i, j int) bool {
 				return *objects[i].Key < *objects[j].Key
 			})
@@ -146,6 +157,14 @@ func setupFsEmulation(s3cli *mocks.S3API) {
 				start, _ = strconv.Atoi(*input.ContinuationToken)
 			}
 
+			if input.StartAfter != nil {
+				for k, v := range objects {
+					if *input.StartAfter == *v.Key {
+						start = k + 1
+					}
+				}
+			}
+
 			isTruncated := false
 			var nextContinuationToken *string
 			if len(objects) > start+maxKeys {
@@ -156,7 +175,24 @@ func setupFsEmulation(s3cli *mocks.S3API) {
 				objects = objects[start:]
 			}
 
+			if input.StartAfter != nil {
+				for k, v := range commonPrefixes {
+					if *input.StartAfter == *v.Prefix {
+						start = k + 1
+					}
+				}
+			}
+
+			if len(commonPrefixes) > start+maxKeys {
+				isTruncated = true
+				nextContinuationToken = common.StringPtr(fmt.Sprintf("%d", start+maxKeys))
+				commonPrefixes = commonPrefixes[start : start+maxKeys]
+			} else if len(commonPrefixes) > 0 {
+				commonPrefixes = commonPrefixes[start:]
+			}
+
 			return &s3.ListObjectsV2Output{
+				CommonPrefixes:        commonPrefixes,
 				Contents:              objects,
 				IsTruncated:           &isTruncated,
 				NextContinuationToken: nextContinuationToken,
@@ -164,6 +200,12 @@ func setupFsEmulation(s3cli *mocks.S3API) {
 		}, nil)
 	s3cli.On("PutObjectWithContext", mock.Anything, mock.Anything).Return(putObjectFn, nil)
 
+	s3cli.On("HeadObjectWithContext", mock.Anything, mock.MatchedBy(func(input *s3.HeadObjectInput) bool {
+		_, ok := fs[*input.Bucket+*input.Key]
+		return !ok
+	})).Return(nil, awserr.New("NotFound", "", nil))
+	s3cli.On("HeadObjectWithContext", mock.Anything, mock.Anything).Return(&s3.HeadObjectOutput{}, nil)
+
 	s3cli.On("GetObjectWithContext", mock.Anything, mock.MatchedBy(func(input *s3.GetObjectInput) bool {
 		_, ok := fs[*input.Bucket+*input.Key]
 		return !ok
@@ -405,7 +447,7 @@ func (s *historyArchiverSuite) TestArchive_Success() {
 	err = historyArchiver.Archive(context.Background(), URI, request)
 	s.NoError(err)
 
-	expectedkey := constructHistoryKey("", testDomainID, testWorkflowID, testRunID, testCloseFailoverVersion)
+	expectedkey := constructHistoryKey("", testDomainID, testWorkflowID, testRunID, testCloseFailoverVersion, 0)
 	s.assertKeyExists(expectedkey)
 }
 
@@ -469,7 +511,7 @@ func (s *historyArchiverSuite) TestGet_Fail_KeyNotExist() {
 	response, err := historyArchiver.Get(context.Background(), URI, request)
 	s.Nil(response)
 	s.Error(err)
-	s.IsType(&shared.BadRequestError{}, err)
+	s.IsType(&shared.EntityNotExistsError{}, err)
 }
 
 func (s *historyArchiverSuite) TestGet_Success_PickHighestVersion() {
@@ -485,7 +527,7 @@ func (s *historyArchiverSuite) TestGet_Success_PickHighestVersion() {
 	response, err := historyArchiver.Get(context.Background(), URI, request)
 	s.NoError(err)
 	s.Nil(response.NextPageToken)
-	s.Equal(s.historyBatchesV100, response.HistoryBatches)
+	s.Equal(append(s.historyBatchesV100[0].Body, s.historyBatchesV100[1].Body...), response.HistoryBatches)
 }
 
 func (s *historyArchiverSuite) TestGet_Success_UseProvidedVersion() {
@@ -502,7 +544,7 @@ func (s *historyArchiverSuite) TestGet_Success_UseProvidedVersion() {
 	response, err := historyArchiver.Get(context.Background(), URI, request)
 	s.NoError(err)
 	s.Nil(response.NextPageToken)
-	s.Equal(s.historyBatchesV1, response.HistoryBatches)
+	s.Equal(s.historyBatchesV1[0].Body, response.HistoryBatches)
 }
 
 func (s *historyArchiverSuite) TestGet_Success_SmallPageSize() {
@@ -535,22 +577,18 @@ func (s *historyArchiverSuite) TestGet_Success_SmallPageSize() {
 	s.Len(response.HistoryBatches, 1)
 	combinedHistory = append(combinedHistory, response.HistoryBatches...)
 
-	s.Equal(s.historyBatchesV100, combinedHistory)
+	s.Equal(append(s.historyBatchesV100[0].Body, s.historyBatchesV100[1].Body...), combinedHistory)
 }
 
 func (s *historyArchiverSuite) TestArchiveAndGet() {
 	mockCtrl := gomock.NewController(s.T())
 	defer mockCtrl.Finish()
 	historyIterator := archiver.NewMockHistoryIterator(mockCtrl)
-	historyBlob := &archiver.HistoryBlob{
-		Header: &archiver.HistoryBlobHeader{
-			IsLast: common.BoolPtr(true),
-		},
-		Body: s.historyBatchesV100,
-	}
 	gomock.InOrder(
 		historyIterator.EXPECT().HasNext().Return(true),
-		historyIterator.EXPECT().Next().Return(historyBlob, nil),
+		historyIterator.EXPECT().Next().Return(s.historyBatchesV100[0], nil),
+		historyIterator.EXPECT().HasNext().Return(true),
+		historyIterator.EXPECT().Next().Return(s.historyBatchesV100[1], nil),
 		historyIterator.EXPECT().HasNext().Return(false),
 	)
 
@@ -569,9 +607,6 @@ func (s *historyArchiverSuite) TestArchiveAndGet() {
 	err = historyArchiver.Archive(context.Background(), URI, archiveRequest)
 	s.NoError(err)
 
-	expectedkey := constructHistoryKey("", testDomainID, testWorkflowID, testRunID, testCloseFailoverVersion)
-	s.assertKeyExists(expectedkey)
-
 	getRequest := &archiver.GetHistoryRequest{
 		DomainID:   testDomainID,
 		WorkflowID: testWorkflowID,
@@ -582,7 +617,7 @@ func (s *historyArchiverSuite) TestArchiveAndGet() {
 	s.NoError(err)
 	s.NotNil(response)
 	s.Nil(response.NextPageToken)
-	s.Equal(s.historyBatchesV100, response.HistoryBatches)
+	s.Equal(append(s.historyBatchesV100[0].Body, s.historyBatchesV100[1].Body...), response.HistoryBatches)
 }
 
 func (s *historyArchiverSuite) newTestHistoryArchiver(historyIterator archiver.HistoryIterator) *historyArchiver {
@@ -597,39 +632,60 @@ func (s *historyArchiverSuite) newTestHistoryArchiver(historyIterator archiver.H
 }
 
 func (s *historyArchiverSuite) setupHistoryDirectory() {
-	s.historyBatchesV1 = []*shared.History{
-		&shared.History{
-			Events: []*shared.HistoryEvent{
-				&shared.HistoryEvent{
-					EventId:   common.Int64Ptr(testNextEventID - 1),
-					Timestamp: common.Int64Ptr(time.Now().UnixNano()),
-					Version:   common.Int64Ptr(1),
+	s.historyBatchesV1 = []*archiver.HistoryBlob{
+		{
+			Header: &archiver.HistoryBlobHeader{
+				IsLast: common.BoolPtr(true),
+			},
+			Body: []*shared.History{
+				{
+					Events: []*shared.HistoryEvent{
+						&shared.HistoryEvent{
+							EventId:   common.Int64Ptr(testNextEventID - 1),
+							Timestamp: common.Int64Ptr(time.Now().UnixNano()),
+							Version:   common.Int64Ptr(1),
+						},
+					},
 				},
 			},
 		},
 	}
 
-	s.historyBatchesV100 = []*shared.History{
-		&shared.History{
-			Events: []*shared.HistoryEvent{
-				&shared.HistoryEvent{
-					EventId:   common.Int64Ptr(common.FirstEventID + 1),
-					Timestamp: common.Int64Ptr(time.Now().UnixNano()),
-					Version:   common.Int64Ptr(testCloseFailoverVersion),
-				},
-				&shared.HistoryEvent{
-					EventId:   common.Int64Ptr(common.FirstEventID + 1),
-					Timestamp: common.Int64Ptr(time.Now().UnixNano()),
-					Version:   common.Int64Ptr(testCloseFailoverVersion),
+	s.historyBatchesV100 = []*archiver.HistoryBlob{
+		{
+			Header: &archiver.HistoryBlobHeader{
+				IsLast: common.BoolPtr(false),
+			},
+			Body: []*shared.History{
+				&shared.History{
+					Events: []*shared.HistoryEvent{
+						&shared.HistoryEvent{
+							EventId:   common.Int64Ptr(common.FirstEventID + 1),
+							Timestamp: common.Int64Ptr(time.Now().UnixNano()),
+							Version:   common.Int64Ptr(testCloseFailoverVersion),
+						},
+						&shared.HistoryEvent{
+							EventId:   common.Int64Ptr(common.FirstEventID + 1),
+							Timestamp: common.Int64Ptr(time.Now().UnixNano()),
+							Version:   common.Int64Ptr(testCloseFailoverVersion),
+						},
+					},
 				},
 			},
 		},
-		&shared.History{
-			Events: []*shared.HistoryEvent{
-				&shared.HistoryEvent{
-					EventId:   common.Int64Ptr(testNextEventID - 1),
-					Timestamp: common.Int64Ptr(time.Now().UnixNano()),
-					Version:   common.Int64Ptr(testCloseFailoverVersion),
+		{
+			Header: &archiver.HistoryBlobHeader{
+				IsLast: common.BoolPtr(true),
+			},
+			Body: []*shared.History{
+				&shared.History{
+					Events: []*shared.HistoryEvent{
+						&shared.HistoryEvent{
+							EventId:   common.Int64Ptr(testNextEventID - 1),
+							Timestamp: common.Int64Ptr(time.Now().UnixNano()),
+							Version:   common.Int64Ptr(testCloseFailoverVersion),
+						},
+					},
 				},
 			},
 		},
@@ -639,16 +695,18 @@ func (s *historyArchiverSuite) setupHistoryDirectory() {
 	s.writeHistoryBatchesForGetTest(s.historyBatchesV100, testCloseFailoverVersion)
 }
 
-func (s *historyArchiverSuite) writeHistoryBatchesForGetTest(historyBatches []*shared.History, version int64) {
-	data, err := encode(historyBatches)
-	s.Require().NoError(err)
-	key := constructHistoryKey("", testDomainID, testWorkflowID, testRunID, version)
-	_, err = s.s3cli.PutObjectWithContext(context.Background(), &s3.PutObjectInput{
-		Bucket: aws.String(testBucket),
-		Key:    aws.String(key),
-		Body:   bytes.NewReader(data),
-	})
-	s.Require().NoError(err)
+func (s *historyArchiverSuite) writeHistoryBatchesForGetTest(historyBatches []*archiver.HistoryBlob, version int64) {
+	for i, batch := range historyBatches {
+		data, err := encode(batch)
+		s.Require().NoError(err)
+		key := constructHistoryKey("", testDomainID, testWorkflowID, testRunID, version, i)
+		_, err = s.s3cli.PutObjectWithContext(context.Background(), &s3.PutObjectInput{
+			Bucket: aws.String(testBucket),
+			Key:    aws.String(key),
+			Body:   bytes.NewReader(data),
+		})
+		s.Require().NoError(err)
+	}
 }
 
 func (s *historyArchiverSuite) assertKeyExists(key string) {