github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/storage/s3.go (about) 1 // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0. 2 3 package storage 4 5 import ( 6 "bytes" 7 "context" 8 "fmt" 9 "io" 10 "net/url" 11 "path" 12 "regexp" 13 "strconv" 14 "strings" 15 "time" 16 17 "github.com/aws/aws-sdk-go/aws" 18 "github.com/aws/aws-sdk-go/aws/awserr" 19 "github.com/aws/aws-sdk-go/aws/client" 20 "github.com/aws/aws-sdk-go/aws/credentials" 21 "github.com/aws/aws-sdk-go/aws/request" 22 "github.com/aws/aws-sdk-go/aws/session" 23 "github.com/aws/aws-sdk-go/service/s3" 24 "github.com/aws/aws-sdk-go/service/s3/s3iface" 25 "github.com/pingcap/errors" 26 backuppb "github.com/pingcap/kvproto/pkg/backup" 27 "github.com/pingcap/log" 28 "github.com/spf13/pflag" 29 "go.uber.org/zap" 30 31 berrors "github.com/pingcap/br/pkg/errors" 32 "github.com/pingcap/br/pkg/logutil" 33 ) 34 35 const ( 36 s3EndpointOption = "s3.endpoint" 37 s3RegionOption = "s3.region" 38 s3StorageClassOption = "s3.storage-class" 39 s3SseOption = "s3.sse" 40 s3SseKmsKeyIDOption = "s3.sse-kms-key-id" 41 s3ACLOption = "s3.acl" 42 s3ProviderOption = "s3.provider" 43 notFound = "NotFound" 44 // number of retries to make of operations. 45 maxRetries = 7 46 // max number of retries when meets error 47 maxErrorRetries = 3 48 49 // the maximum number of byte to read for seek. 50 maxSkipOffsetByRead = 1 << 16 // 64KB 51 52 // TODO make this configurable, 5 mb is a good minimum size but on low latency/high bandwidth network you can go a lot bigger 53 hardcodedS3ChunkSize = 5 * 1024 * 1024 54 ) 55 56 var permissionCheckFn = map[Permission]func(*s3.S3, *backuppb.S3) error{ 57 AccessBuckets: checkS3Bucket, 58 ListObjects: listObjects, 59 GetObject: getObject, 60 } 61 62 // S3Storage info for s3 storage. 63 type S3Storage struct { 64 session *session.Session 65 svc s3iface.S3API 66 options *backuppb.S3 67 } 68 69 // S3Uploader does multi-part upload to s3. 70 type S3Uploader struct { 71 svc s3iface.S3API 72 createOutput *s3.CreateMultipartUploadOutput 73 completeParts []*s3.CompletedPart 74 } 75 76 // UploadPart update partial data to s3, we should call CreateMultipartUpload to start it, 77 // and call CompleteMultipartUpload to finish it. 78 func (u *S3Uploader) Write(ctx context.Context, data []byte) (int, error) { 79 partInput := &s3.UploadPartInput{ 80 Body: bytes.NewReader(data), 81 Bucket: u.createOutput.Bucket, 82 Key: u.createOutput.Key, 83 PartNumber: aws.Int64(int64(len(u.completeParts) + 1)), 84 UploadId: u.createOutput.UploadId, 85 ContentLength: aws.Int64(int64(len(data))), 86 } 87 88 uploadResult, err := u.svc.UploadPartWithContext(ctx, partInput) 89 if err != nil { 90 return 0, errors.Trace(err) 91 } 92 u.completeParts = append(u.completeParts, &s3.CompletedPart{ 93 ETag: uploadResult.ETag, 94 PartNumber: partInput.PartNumber, 95 }) 96 return len(data), nil 97 } 98 99 // Close complete multi upload request. 100 func (u *S3Uploader) Close(ctx context.Context) error { 101 completeInput := &s3.CompleteMultipartUploadInput{ 102 Bucket: u.createOutput.Bucket, 103 Key: u.createOutput.Key, 104 UploadId: u.createOutput.UploadId, 105 MultipartUpload: &s3.CompletedMultipartUpload{ 106 Parts: u.completeParts, 107 }, 108 } 109 _, err := u.svc.CompleteMultipartUploadWithContext(ctx, completeInput) 110 return errors.Trace(err) 111 } 112 113 // S3BackendOptions contains options for s3 storage. 114 type S3BackendOptions struct { 115 Endpoint string `json:"endpoint" toml:"endpoint"` 116 Region string `json:"region" toml:"region"` 117 StorageClass string `json:"storage-class" toml:"storage-class"` 118 Sse string `json:"sse" toml:"sse"` 119 SseKmsKeyID string `json:"sse-kms-key-id" toml:"sse-kms-key-id"` 120 ACL string `json:"acl" toml:"acl"` 121 AccessKey string `json:"access-key" toml:"access-key"` 122 SecretAccessKey string `json:"secret-access-key" toml:"secret-access-key"` 123 Provider string `json:"provider" toml:"provider"` 124 ForcePathStyle bool `json:"force-path-style" toml:"force-path-style"` 125 UseAccelerateEndpoint bool `json:"use-accelerate-endpoint" toml:"use-accelerate-endpoint"` 126 } 127 128 // Apply apply s3 options on backuppb.S3. 129 func (options *S3BackendOptions) Apply(s3 *backuppb.S3) error { 130 if options.Region == "" { 131 options.Region = "us-east-1" 132 } 133 if options.Endpoint != "" { 134 u, err := url.Parse(options.Endpoint) 135 if err != nil { 136 return errors.Trace(err) 137 } 138 if u.Scheme == "" { 139 return errors.Annotate(berrors.ErrStorageInvalidConfig, "scheme not found in endpoint") 140 } 141 if u.Host == "" { 142 return errors.Annotate(berrors.ErrStorageInvalidConfig, "host not found in endpoint") 143 } 144 } 145 // In some cases, we need to set ForcePathStyle to false. 146 // Refer to: https://rclone.org/s3/#s3-force-path-style 147 if options.Provider == "alibaba" || options.Provider == "netease" || 148 options.UseAccelerateEndpoint { 149 options.ForcePathStyle = false 150 } 151 if options.AccessKey == "" && options.SecretAccessKey != "" { 152 return errors.Annotate(berrors.ErrStorageInvalidConfig, "access_key not found") 153 } 154 if options.AccessKey != "" && options.SecretAccessKey == "" { 155 return errors.Annotate(berrors.ErrStorageInvalidConfig, "secret_access_key not found") 156 } 157 158 s3.Endpoint = options.Endpoint 159 s3.Region = options.Region 160 // StorageClass, SSE and ACL are acceptable to be empty 161 s3.StorageClass = options.StorageClass 162 s3.Sse = options.Sse 163 s3.SseKmsKeyId = options.SseKmsKeyID 164 s3.Acl = options.ACL 165 s3.AccessKey = options.AccessKey 166 s3.SecretAccessKey = options.SecretAccessKey 167 s3.ForcePathStyle = options.ForcePathStyle 168 return nil 169 } 170 171 // defineS3Flags defines the command line flags for S3BackendOptions. 172 func defineS3Flags(flags *pflag.FlagSet) { 173 // TODO: remove experimental tag if it's stable 174 flags.String(s3EndpointOption, "", 175 "(experimental) Set the S3 endpoint URL, please specify the http or https scheme explicitly") 176 flags.String(s3RegionOption, "", "(experimental) Set the S3 region, e.g. us-east-1") 177 flags.String(s3StorageClassOption, "", "(experimental) Set the S3 storage class, e.g. STANDARD") 178 flags.String(s3SseOption, "", "Set S3 server-side encryption, e.g. aws:kms") 179 flags.String(s3SseKmsKeyIDOption, "", "KMS CMK key id to use with S3 server-side encryption."+ 180 "Leave empty to use S3 owned key.") 181 flags.String(s3ACLOption, "", "(experimental) Set the S3 canned ACLs, e.g. authenticated-read") 182 flags.String(s3ProviderOption, "", "(experimental) Set the S3 provider, e.g. aws, alibaba, ceph") 183 } 184 185 // parseFromFlags parse S3BackendOptions from command line flags. 186 func (options *S3BackendOptions) parseFromFlags(flags *pflag.FlagSet) error { 187 var err error 188 options.Endpoint, err = flags.GetString(s3EndpointOption) 189 if err != nil { 190 return errors.Trace(err) 191 } 192 options.Region, err = flags.GetString(s3RegionOption) 193 if err != nil { 194 return errors.Trace(err) 195 } 196 options.Sse, err = flags.GetString(s3SseOption) 197 if err != nil { 198 return errors.Trace(err) 199 } 200 options.SseKmsKeyID, err = flags.GetString(s3SseKmsKeyIDOption) 201 if err != nil { 202 return errors.Trace(err) 203 } 204 options.ACL, err = flags.GetString(s3ACLOption) 205 if err != nil { 206 return errors.Trace(err) 207 } 208 options.StorageClass, err = flags.GetString(s3StorageClassOption) 209 if err != nil { 210 return errors.Trace(err) 211 } 212 options.ForcePathStyle = true 213 options.Provider, err = flags.GetString(s3ProviderOption) 214 if err != nil { 215 return errors.Trace(err) 216 } 217 return nil 218 } 219 220 // NewS3StorageForTest creates a new S3Storage for testing only. 221 func NewS3StorageForTest(svc s3iface.S3API, options *backuppb.S3) *S3Storage { 222 return &S3Storage{ 223 session: nil, 224 svc: svc, 225 options: options, 226 } 227 } 228 229 // NewS3Storage initialize a new s3 storage for metadata. 230 // 231 // Deprecated: Create the storage via `New()` instead of using this. 232 func NewS3Storage( // revive:disable-line:flag-parameter 233 backend *backuppb.S3, 234 sendCredential bool, 235 ) (*S3Storage, error) { 236 return newS3Storage(backend, &ExternalStorageOptions{ 237 SendCredentials: sendCredential, 238 CheckPermissions: []Permission{AccessBuckets}, 239 }) 240 } 241 242 func newS3Storage(backend *backuppb.S3, opts *ExternalStorageOptions) (*S3Storage, error) { 243 qs := *backend 244 awsConfig := aws.NewConfig(). 245 WithS3ForcePathStyle(qs.ForcePathStyle). 246 WithRegion(qs.Region) 247 request.WithRetryer(awsConfig, defaultS3Retryer()) 248 if qs.Endpoint != "" { 249 awsConfig.WithEndpoint(qs.Endpoint) 250 } 251 if opts.HTTPClient != nil { 252 awsConfig.WithHTTPClient(opts.HTTPClient) 253 } 254 var cred *credentials.Credentials 255 if qs.AccessKey != "" && qs.SecretAccessKey != "" { 256 cred = credentials.NewStaticCredentials(qs.AccessKey, qs.SecretAccessKey, "") 257 } 258 if cred != nil { 259 awsConfig.WithCredentials(cred) 260 } 261 // awsConfig.WithLogLevel(aws.LogDebugWithSigning) 262 awsSessionOpts := session.Options{ 263 Config: *awsConfig, 264 } 265 ses, err := session.NewSessionWithOptions(awsSessionOpts) 266 if err != nil { 267 return nil, errors.Trace(err) 268 } 269 270 if !opts.SendCredentials { 271 // Clear the credentials if exists so that they will not be sent to TiKV 272 backend.AccessKey = "" 273 backend.SecretAccessKey = "" 274 } else if ses.Config.Credentials != nil { 275 if qs.AccessKey == "" || qs.SecretAccessKey == "" { 276 v, cerr := ses.Config.Credentials.Get() 277 if cerr != nil { 278 return nil, errors.Trace(cerr) 279 } 280 backend.AccessKey = v.AccessKeyID 281 backend.SecretAccessKey = v.SecretAccessKey 282 } 283 } 284 285 c := s3.New(ses) 286 // TODO remove it after BR remove cfg skip-check-path 287 if !opts.SkipCheckPath { 288 err = checkS3Bucket(c, &qs) 289 if err != nil { 290 return nil, errors.Annotatef(berrors.ErrStorageInvalidConfig, "Bucket %s is not accessible: %v", qs.Bucket, err) 291 } 292 } 293 294 if len(qs.Prefix) > 0 && !strings.HasSuffix(qs.Prefix, "/") { 295 qs.Prefix += "/" 296 } 297 298 for _, p := range opts.CheckPermissions { 299 err := permissionCheckFn[p](c, &qs) 300 if err != nil { 301 return nil, errors.Annotatef(berrors.ErrStorageInvalidPermission, "check permission %s failed due to %v", p, err) 302 } 303 } 304 305 return &S3Storage{ 306 session: ses, 307 svc: c, 308 options: &qs, 309 }, nil 310 } 311 312 // checkBucket checks if a bucket exists. 313 func checkS3Bucket(svc *s3.S3, qs *backuppb.S3) error { 314 input := &s3.HeadBucketInput{ 315 Bucket: aws.String(qs.Bucket), 316 } 317 _, err := svc.HeadBucket(input) 318 return errors.Trace(err) 319 } 320 321 // listObjects checks the permission of listObjects 322 func listObjects(svc *s3.S3, qs *backuppb.S3) error { 323 input := &s3.ListObjectsInput{ 324 Bucket: aws.String(qs.Bucket), 325 Prefix: aws.String(qs.Prefix), 326 MaxKeys: aws.Int64(1), 327 } 328 _, err := svc.ListObjects(input) 329 if err != nil { 330 return errors.Trace(err) 331 } 332 return nil 333 } 334 335 // getObject checks the permission of getObject 336 func getObject(svc *s3.S3, qs *backuppb.S3) error { 337 input := &s3.GetObjectInput{ 338 Bucket: aws.String(qs.Bucket), 339 Key: aws.String("not-exists"), 340 } 341 _, err := svc.GetObject(input) 342 if aerr, ok := err.(awserr.Error); ok { 343 if aerr.Code() == "NoSuchKey" { 344 // if key not exists and we reach this error, that 345 // means we have the correct permission to GetObject 346 // other we will get another error 347 return nil 348 } 349 return errors.Trace(err) 350 } 351 return nil 352 } 353 354 // WriteFile writes data to a file to storage. 355 func (rs *S3Storage) WriteFile(ctx context.Context, file string, data []byte) error { 356 input := &s3.PutObjectInput{ 357 Body: aws.ReadSeekCloser(bytes.NewReader(data)), 358 Bucket: aws.String(rs.options.Bucket), 359 Key: aws.String(rs.options.Prefix + file), 360 } 361 if rs.options.Acl != "" { 362 input = input.SetACL(rs.options.Acl) 363 } 364 if rs.options.Sse != "" { 365 input = input.SetServerSideEncryption(rs.options.Sse) 366 } 367 if rs.options.SseKmsKeyId != "" { 368 input = input.SetSSEKMSKeyId(rs.options.SseKmsKeyId) 369 } 370 if rs.options.StorageClass != "" { 371 input = input.SetStorageClass(rs.options.StorageClass) 372 } 373 374 _, err := rs.svc.PutObjectWithContext(ctx, input) 375 if err != nil { 376 return errors.Trace(err) 377 } 378 hinput := &s3.HeadObjectInput{ 379 Bucket: aws.String(rs.options.Bucket), 380 Key: aws.String(rs.options.Prefix + file), 381 } 382 err = rs.svc.WaitUntilObjectExistsWithContext(ctx, hinput) 383 return errors.Trace(err) 384 } 385 386 // ReadFile reads the file from the storage and returns the contents. 387 func (rs *S3Storage) ReadFile(ctx context.Context, file string) ([]byte, error) { 388 input := &s3.GetObjectInput{ 389 Bucket: aws.String(rs.options.Bucket), 390 Key: aws.String(rs.options.Prefix + file), 391 } 392 result, err := rs.svc.GetObjectWithContext(ctx, input) 393 if err != nil { 394 return nil, errors.Annotatef(err, 395 "failed to read s3 file, file info: input.bucket='%s', input.key='%s'", 396 *input.Bucket, *input.Key) 397 } 398 defer result.Body.Close() 399 data, err := io.ReadAll(result.Body) 400 if err != nil { 401 return nil, errors.Trace(err) 402 } 403 return data, nil 404 } 405 406 // FileExists check if file exists on s3 storage. 407 func (rs *S3Storage) FileExists(ctx context.Context, file string) (bool, error) { 408 input := &s3.HeadObjectInput{ 409 Bucket: aws.String(rs.options.Bucket), 410 Key: aws.String(rs.options.Prefix + file), 411 } 412 413 _, err := rs.svc.HeadObjectWithContext(ctx, input) 414 if err != nil { 415 if aerr, ok := errors.Cause(err).(awserr.Error); ok { // nolint:errorlint 416 switch aerr.Code() { 417 case s3.ErrCodeNoSuchBucket, s3.ErrCodeNoSuchKey, notFound: 418 return false, nil 419 } 420 } 421 return false, errors.Trace(err) 422 } 423 return true, nil 424 } 425 426 // WalkDir traverse all the files in a dir. 427 // 428 // fn is the function called for each regular file visited by WalkDir. 429 // The first argument is the file path that can be used in `Open` 430 // function; the second argument is the size in byte of the file determined 431 // by path. 432 func (rs *S3Storage) WalkDir(ctx context.Context, opt *WalkOption, fn func(string, int64) error) error { 433 if opt == nil { 434 opt = &WalkOption{} 435 } 436 prefix := path.Join(rs.options.Prefix, opt.SubDir) 437 if len(prefix) > 0 && !strings.HasSuffix(prefix, "/") { 438 prefix += "/" 439 } 440 maxKeys := int64(1000) 441 if opt.ListCount > 0 { 442 maxKeys = opt.ListCount 443 } 444 req := &s3.ListObjectsInput{ 445 Bucket: aws.String(rs.options.Bucket), 446 Prefix: aws.String(prefix), 447 MaxKeys: aws.Int64(maxKeys), 448 } 449 450 for { 451 // FIXME: We can't use ListObjectsV2, it is not universally supported. 452 // (Ceph RGW supported ListObjectsV2 since v15.1.0, released 2020 Jan 30th) 453 // (as of 2020, DigitalOcean Spaces still does not support V2 - https://developers.digitalocean.com/documentation/spaces/#list-bucket-contents) 454 res, err := rs.svc.ListObjectsWithContext(ctx, req) 455 if err != nil { 456 return errors.Trace(err) 457 } 458 for _, r := range res.Contents { 459 // when walk on specify directory, the result include storage.Prefix, 460 // which can not be reuse in other API(Open/Read) directly. 461 // so we use TrimPrefix to filter Prefix for next Open/Read. 462 path := strings.TrimPrefix(*r.Key, rs.options.Prefix) 463 if err = fn(path, *r.Size); err != nil { 464 return errors.Trace(err) 465 } 466 467 // https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjects.html#AmazonS3-ListObjects-response-NextMarker - 468 // 469 // `res.NextMarker` is populated only if we specify req.Delimiter. 470 // Aliyun OSS and minio will populate NextMarker no matter what, 471 // but this documented behavior does apply to AWS S3: 472 // 473 // "If response does not include the NextMarker and it is truncated, 474 // you can use the value of the last Key in the response as the marker 475 // in the subsequent request to get the next set of object keys." 476 req.Marker = r.Key 477 } 478 if !aws.BoolValue(res.IsTruncated) { 479 break 480 } 481 } 482 483 return nil 484 } 485 486 // URI returns s3://<base>/<prefix>. 487 func (rs *S3Storage) URI() string { 488 return "s3://" + rs.options.Bucket + "/" + rs.options.Prefix 489 } 490 491 // Open a Reader by file path. 492 func (rs *S3Storage) Open(ctx context.Context, path string) (ExternalFileReader, error) { 493 reader, r, err := rs.open(ctx, path, 0, 0) 494 if err != nil { 495 return nil, errors.Trace(err) 496 } 497 return &s3ObjectReader{ 498 storage: rs, 499 name: path, 500 reader: reader, 501 ctx: ctx, 502 rangeInfo: r, 503 }, nil 504 } 505 506 // RangeInfo represents the an HTTP Content-Range header value 507 // of the form `bytes [Start]-[End]/[Size]`. 508 type RangeInfo struct { 509 // Start is the absolute position of the first byte of the byte range, 510 // starting from 0. 511 Start int64 512 // End is the absolute position of the last byte of the byte range. This end 513 // offset is inclusive, e.g. if the Size is 1000, the maximum value of End 514 // would be 999. 515 End int64 516 // Size is the total size of the original file. 517 Size int64 518 } 519 520 // if endOffset > startOffset, should return reader for bytes in [startOffset, endOffset). 521 func (rs *S3Storage) open( 522 ctx context.Context, 523 path string, 524 startOffset, endOffset int64, 525 ) (io.ReadCloser, RangeInfo, error) { 526 input := &s3.GetObjectInput{ 527 Bucket: aws.String(rs.options.Bucket), 528 Key: aws.String(rs.options.Prefix + path), 529 } 530 531 // always set rangeOffset to fetch file size info 532 // s3 endOffset is inclusive 533 var rangeOffset *string 534 if endOffset > startOffset { 535 rangeOffset = aws.String(fmt.Sprintf("bytes=%d-%d", startOffset, endOffset-1)) 536 } else { 537 rangeOffset = aws.String(fmt.Sprintf("bytes=%d-", startOffset)) 538 } 539 input.Range = rangeOffset 540 result, err := rs.svc.GetObjectWithContext(ctx, input) 541 if err != nil { 542 return nil, RangeInfo{}, errors.Trace(err) 543 } 544 545 r, err := ParseRangeInfo(result.ContentRange) 546 if err != nil { 547 return nil, RangeInfo{}, errors.Trace(err) 548 } 549 550 if startOffset != r.Start || (endOffset != 0 && endOffset != r.End+1) { 551 return nil, r, errors.Annotatef(berrors.ErrStorageUnknown, "open file '%s' failed, expected range: %s, got: %v", 552 path, *rangeOffset, result.ContentRange) 553 } 554 555 return result.Body, r, nil 556 } 557 558 var contentRangeRegex = regexp.MustCompile(`bytes (\d+)-(\d+)/(\d+)$`) 559 560 // ParseRangeInfo parses the Content-Range header and returns the offsets. 561 func ParseRangeInfo(info *string) (ri RangeInfo, err error) { 562 if info == nil || len(*info) == 0 { 563 err = errors.Annotate(berrors.ErrStorageUnknown, "ContentRange is empty") 564 return 565 } 566 subMatches := contentRangeRegex.FindStringSubmatch(*info) 567 if len(subMatches) != 4 { 568 err = errors.Annotatef(berrors.ErrStorageUnknown, "invalid content range: '%s'", *info) 569 return 570 } 571 572 ri.Start, err = strconv.ParseInt(subMatches[1], 10, 64) 573 if err != nil { 574 err = errors.Annotatef(err, "invalid start offset value '%s' in ContentRange '%s'", subMatches[1], *info) 575 return 576 } 577 ri.End, err = strconv.ParseInt(subMatches[2], 10, 64) 578 if err != nil { 579 err = errors.Annotatef(err, "invalid end offset value '%s' in ContentRange '%s'", subMatches[2], *info) 580 return 581 } 582 ri.Size, err = strconv.ParseInt(subMatches[3], 10, 64) 583 if err != nil { 584 err = errors.Annotatef(err, "invalid size size value '%s' in ContentRange '%s'", subMatches[3], *info) 585 return 586 } 587 return 588 } 589 590 // s3ObjectReader wrap GetObjectOutput.Body and add the `Seek` method. 591 type s3ObjectReader struct { 592 storage *S3Storage 593 name string 594 reader io.ReadCloser 595 pos int64 596 rangeInfo RangeInfo 597 // reader context used for implement `io.Seek` 598 // currently, lightning depends on package `xitongsys/parquet-go` to read parquet file and it needs `io.Seeker` 599 // See: https://github.com/xitongsys/parquet-go/blob/207a3cee75900b2b95213627409b7bac0f190bb3/source/source.go#L9-L10 600 ctx context.Context 601 retryCnt int 602 } 603 604 // Read implement the io.Reader interface. 605 func (r *s3ObjectReader) Read(p []byte) (n int, err error) { 606 maxCnt := r.rangeInfo.End + 1 - r.pos 607 if maxCnt > int64(len(p)) { 608 maxCnt = int64(len(p)) 609 } 610 n, err = r.reader.Read(p[:maxCnt]) 611 // TODO: maybe we should use !errors.Is(err, io.EOF) here to avoid error lint, but currently, pingcap/errors 612 // doesn't implement this method yet. 613 if err != nil && errors.Cause(err) != io.EOF && r.retryCnt < maxErrorRetries { //nolint:errorlint 614 // if can retry, reopen a new reader and try read again 615 end := r.rangeInfo.End + 1 616 if end == r.rangeInfo.Size { 617 end = 0 618 } 619 _ = r.reader.Close() 620 621 newReader, _, err1 := r.storage.open(r.ctx, r.name, r.pos, end) 622 if err1 != nil { 623 log.Warn("open new s3 reader failed", zap.String("file", r.name), zap.Error(err1)) 624 return 625 } 626 r.reader = newReader 627 r.retryCnt++ 628 n, err = r.reader.Read(p[:maxCnt]) 629 } 630 631 r.pos += int64(n) 632 return 633 } 634 635 // Close implement the io.Closer interface. 636 func (r *s3ObjectReader) Close() error { 637 return r.reader.Close() 638 } 639 640 // Seek implement the io.Seeker interface. 641 // 642 // Currently, tidb-lightning depends on this method to read parquet file for s3 storage. 643 func (r *s3ObjectReader) Seek(offset int64, whence int) (int64, error) { 644 var realOffset int64 645 switch whence { 646 case io.SeekStart: 647 realOffset = offset 648 case io.SeekCurrent: 649 realOffset = r.pos + offset 650 case io.SeekEnd: 651 realOffset = r.rangeInfo.Size + offset 652 default: 653 return 0, errors.Annotatef(berrors.ErrStorageUnknown, "Seek: invalid whence '%d'", whence) 654 } 655 656 if realOffset == r.pos { 657 return realOffset, nil 658 } else if realOffset >= r.rangeInfo.Size { 659 // See: https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.35 660 // because s3's GetObject interface doesn't allow get a range that matches zero length data, 661 // so if the position is out of range, we need to always return io.EOF after the seek operation. 662 663 // close current read and open a new one which target offset 664 if err := r.reader.Close(); err != nil { 665 log.L().Warn("close s3 reader failed, will ignore this error", logutil.ShortError(err)) 666 } 667 668 r.reader = io.NopCloser(bytes.NewReader(nil)) 669 r.pos = r.rangeInfo.Size 670 return r.pos, nil 671 } 672 673 // if seek ahead no more than 64k, we discard these data 674 if realOffset > r.pos && realOffset-r.pos <= maxSkipOffsetByRead { 675 _, err := io.CopyN(io.Discard, r, realOffset-r.pos) 676 if err != nil { 677 return r.pos, errors.Trace(err) 678 } 679 return realOffset, nil 680 } 681 682 // close current read and open a new one which target offset 683 err := r.reader.Close() 684 if err != nil { 685 return 0, errors.Trace(err) 686 } 687 688 newReader, info, err := r.storage.open(r.ctx, r.name, realOffset, 0) 689 if err != nil { 690 return 0, errors.Trace(err) 691 } 692 r.reader = newReader 693 r.rangeInfo = info 694 r.pos = realOffset 695 return realOffset, nil 696 } 697 698 // CreateUploader create multi upload request. 699 func (rs *S3Storage) CreateUploader(ctx context.Context, name string) (ExternalFileWriter, error) { 700 input := &s3.CreateMultipartUploadInput{ 701 Bucket: aws.String(rs.options.Bucket), 702 Key: aws.String(rs.options.Prefix + name), 703 } 704 if rs.options.Acl != "" { 705 input = input.SetACL(rs.options.Acl) 706 } 707 if rs.options.Sse != "" { 708 input = input.SetServerSideEncryption(rs.options.Sse) 709 } 710 if rs.options.SseKmsKeyId != "" { 711 input = input.SetSSEKMSKeyId(rs.options.SseKmsKeyId) 712 } 713 if rs.options.StorageClass != "" { 714 input = input.SetStorageClass(rs.options.StorageClass) 715 } 716 717 resp, err := rs.svc.CreateMultipartUploadWithContext(ctx, input) 718 if err != nil { 719 return nil, errors.Trace(err) 720 } 721 return &S3Uploader{ 722 svc: rs.svc, 723 createOutput: resp, 724 completeParts: make([]*s3.CompletedPart, 0, 128), 725 }, nil 726 } 727 728 // Create creates multi upload request. 729 func (rs *S3Storage) Create(ctx context.Context, name string) (ExternalFileWriter, error) { 730 uploader, err := rs.CreateUploader(ctx, name) 731 if err != nil { 732 return nil, err 733 } 734 uploaderWriter := newBufferedWriter(uploader, hardcodedS3ChunkSize, NoCompression) 735 return uploaderWriter, nil 736 } 737 738 // retryerWithLog wrappes the client.DefaultRetryer, and logging when retry triggered. 739 type retryerWithLog struct { 740 client.DefaultRetryer 741 } 742 743 func (rl retryerWithLog) RetryRules(r *request.Request) time.Duration { 744 backoffTime := rl.DefaultRetryer.RetryRules(r) 745 if backoffTime > 0 { 746 log.Warn("failed to request s3, retrying", zap.Error(r.Error), zap.Duration("backoff", backoffTime)) 747 } 748 return backoffTime 749 } 750 751 func defaultS3Retryer() request.Retryer { 752 return retryerWithLog{ 753 DefaultRetryer: client.DefaultRetryer{ 754 NumMaxRetries: maxRetries, 755 MinRetryDelay: 1 * time.Second, 756 MinThrottleDelay: 2 * time.Second, 757 }, 758 } 759 }