diff --git a/cmd/kosli/byteSize.go b/cmd/kosli/byteSize.go new file mode 100644 index 000000000..85bcbcbf0 --- /dev/null +++ b/cmd/kosli/byteSize.go @@ -0,0 +1,66 @@ +package main + +import ( + "errors" + "fmt" + "math" + "strconv" + "strings" + "unicode" +) + +// byteSizeUnits maps a lower-cased unit, with any trailing "b" or "ib" already +// stripped, to bytes. Units are binary, as disk figures are. +var byteSizeUnits = map[string]int64{ + "": 1 << 20, // a bare number is megabytes + "b": 1, + "k": 1 << 10, + "m": 1 << 20, + "g": 1 << 30, + "t": 1 << 40, +} + +// parseByteSize turns "512", "512M", "8GB" or "1.5G" into bytes: a bare number +// is megabytes, a K, M, G or T suffix takes an optional B, and "B" alone is bytes. +func parseByteSize(s string) (int64, error) { + s = strings.TrimSpace(s) + if s == "" { + return 0, errors.New("size is empty") + } + + digits := 0 + for digits < len(s) && (s[digits] >= '0' && s[digits] <= '9' || s[digits] == '.') { + digits++ + } + number, unit := s[:digits], strings.TrimSpace(s[digits:]) + if number == "" || strings.ContainsFunc(unit, func(r rune) bool { return !unicode.IsLetter(r) }) { + return 0, fmt.Errorf("%q is not a size: expected a number with an optional K, M, G or T unit, e.g. 512M", s) + } + value, err := strconv.ParseFloat(number, 64) + if err != nil { + return 0, fmt.Errorf("%q is not a size: expected a number with an optional K, M, G or T unit, e.g. 512M", s) + } + + key := strings.ToLower(unit) + if key != "" && key != "b" { + // Only a single unit letter may precede the optional "b" or "ib", so + // "ib" alone and "bb" are unknown rather than a guess at megabytes or bytes. + key = strings.TrimSuffix(strings.TrimSuffix(key, "ib"), "b") + if len(key) != 1 || key == "b" { + return 0, fmt.Errorf("unknown unit %q in size %q: use K, M, G or T, optionally followed by B", unit, s) + } + } + multiplier, ok := byteSizeUnits[key] + if !ok { + return 0, fmt.Errorf("unknown unit %q in size %q: use K, M, G or T, optionally followed by B", unit, s) + } + + bytes := value * float64(multiplier) + if bytes >= math.MaxInt64 { + return 0, fmt.Errorf("size %q is too large", s) + } + if bytes < 1 { + return 0, fmt.Errorf("size %q must be at least 1 byte", s) + } + return int64(bytes), nil +} diff --git a/cmd/kosli/byteSize_test.go b/cmd/kosli/byteSize_test.go new file mode 100644 index 000000000..7dd90ce2a --- /dev/null +++ b/cmd/kosli/byteSize_test.go @@ -0,0 +1,68 @@ +package main + +import ( + "testing" + + "github.com/kosli-dev/cli/internal/aws" + "github.com/stretchr/testify/require" +) + +func TestParseByteSize(t *testing.T) { + const mib = int64(1) << 20 + for _, tc := range []struct { + input string + want int64 + wantErr string + }{ + {input: "512", want: 512 * mib}, + {input: "1", want: mib}, + {input: " 64 ", want: 64 * mib}, + {input: "512M", want: 512 * mib}, + {input: "512MB", want: 512 * mib}, + {input: "512mb", want: 512 * mib}, + {input: "512MiB", want: 512 * mib}, + {input: "8G", want: 8 << 30}, + {input: "8GB", want: 8 << 30}, + {input: "8 GB", want: 8 << 30}, + {input: "2T", want: 2 << 40}, + {input: "1024K", want: 1 << 20}, + {input: "4096KB", want: 4 << 20}, + {input: "1000B", want: 1000}, + {input: "1.5G", want: 3 << 29}, + {input: "0.5M", want: 512 << 10}, + {input: "", wantErr: "empty"}, + {input: "0", wantErr: "must be at least 1 byte"}, + {input: "0B", wantErr: "must be at least 1 byte"}, + {input: "-1", wantErr: "not a size"}, + {input: "-512M", wantErr: "not a size"}, + {input: "abc", wantErr: "not a size"}, + {input: "M", wantErr: "not a size"}, + {input: "512X", wantErr: `unknown unit "X"`}, + {input: "5ib", wantErr: `unknown unit "ib"`}, + {input: "5bb", wantErr: `unknown unit "bb"`}, + {input: "5KiBB", wantErr: `unknown unit "KiBB"`}, + {input: "512 megabytes", wantErr: `unknown unit "megabytes"`}, + {input: "1e3", wantErr: "not a size"}, + {input: "0x10", wantErr: "not a size"}, + {input: "1,024", wantErr: "not a size"}, + {input: "99999999999T", wantErr: "too large"}, + } { + t.Run(tc.input, func(t *testing.T) { + got, err := parseByteSize(tc.input) + if tc.wantErr != "" { + require.Error(t, err) + require.Contains(t, err.Error(), tc.wantErr) + return + } + require.NoError(t, err) + require.Equal(t, tc.want, got) + }) + } +} + +// The flag default is a string, so it can drift from the value it spells. +func TestDefaultDownloadBudgetMatchesTheAwsDefault(t *testing.T) { + got, err := parseByteSize(defaultDownloadBudget) + require.NoError(t, err) + require.Equal(t, aws.DefaultDownloadLimits.BytesInFlight, got) +} diff --git a/cmd/kosli/root.go b/cmd/kosli/root.go index 18b17bd3a..4e7a6ab6a 100644 --- a/cmd/kosli/root.go +++ b/cmd/kosli/root.go @@ -259,6 +259,8 @@ Paths the list already matches stay excluded whatever is later added there, so k awsSecretKeyFlag = "The AWS secret access key." awsRegionFlag = "The AWS region." bucketNameFlag = "The name of the S3 bucket." + downloadConcurrencyFlag = "[optional] The number of S3 objects to download at the same time when fingerprinting the bucket. Each object in flight may hold up to 40 MB of download buffers in memory, on top of the disk the --download-budget allows." + downloadBudgetFlag = "[optional] The maximum total size of the S3 objects downloading at the same time, which caps the temporary disk the snapshot uses. A bare number is megabytes; add K, M, G or T (optionally followed by B) to choose the unit, e.g. 512M or 8G. An object larger than the budget still downloads, on its own. Objects are downloaded to the OS temporary directory." bucketPathsFlag = "[optional] The comma separated list of file and/or directory paths in the S3 bucket to include when fingerprinting. Paths match by literal prefix. Cannot be used together with --exclude or --exclude-regex." bucketPathsRegexFlag = "[optional] The comma separated list of Go regular expressions matched against object keys in the S3 bucket to include when fingerprinting. Cannot be used together with --exclude or --exclude-regex." excludeBucketPathsFlag = "[optional] The comma separated list of file and/or directory paths in the S3 bucket to exclude when fingerprinting. Paths match by literal prefix. Cannot be used together with --include or --include-regex." diff --git a/cmd/kosli/snapshotS3.go b/cmd/kosli/snapshotS3.go index 975c4a2f8..619ce7ac5 100644 --- a/cmd/kosli/snapshotS3.go +++ b/cmd/kosli/snapshotS3.go @@ -1,6 +1,7 @@ package main import ( + "fmt" "io" "net/http" "net/url" @@ -71,12 +72,15 @@ kosli snapshot s3 yourEnvironmentName \ ` type snapshotS3Options struct { - bucket string - includePaths []string - includeRegex []string - excludePaths []string - excludeRegex []string - awsStaticCreds *aws.AWSStaticCreds + bucket string + includePaths []string + includeRegex []string + excludePaths []string + excludeRegex []string + downloadConcurrency int + downloadBudget string + downloadLimits aws.DownloadLimits + awsStaticCreds *aws.AWSStaticCreds } func newSnapshotS3Cmd(out io.Writer) *cobra.Command { @@ -108,7 +112,7 @@ func newSnapshotS3Cmd(out io.Writer) *cobra.Command { } } - return nil + return o.resolveDownloadLimits() }, RunE: func(cmd *cobra.Command, args []string) error { return o.run(args) @@ -120,6 +124,8 @@ func newSnapshotS3Cmd(out io.Writer) *cobra.Command { cmd.Flags().StringSliceVar(&o.includeRegex, "include-regex", []string{}, bucketPathsRegexFlag) cmd.Flags().StringSliceVarP(&o.excludePaths, "exclude", "x", []string{}, excludeBucketPathsFlag) cmd.Flags().StringSliceVar(&o.excludeRegex, "exclude-regex", []string{}, excludeBucketPathsRegexFlag) + cmd.Flags().IntVar(&o.downloadConcurrency, "download-concurrency", aws.DefaultDownloadLimits.Concurrency, downloadConcurrencyFlag) + cmd.Flags().StringVar(&o.downloadBudget, "download-budget", defaultDownloadBudget, downloadBudgetFlag) addAWSAuthFlags(cmd, o.awsStaticCreds) addDryRunFlag(cmd) @@ -143,7 +149,7 @@ func (o *snapshotS3Options) run(args []string) error { return err } - s3Data, err := o.awsStaticCreds.GetS3Data(o.bucket, o.includePaths, o.includeRegex, o.excludePaths, o.excludeRegex, logger) + s3Data, err := o.awsStaticCreds.GetS3Data(o.bucket, o.includePaths, o.includeRegex, o.excludePaths, o.excludeRegex, o.downloadLimits, logger) if err != nil { return err } @@ -164,3 +170,19 @@ func (o *snapshotS3Options) run(args []string) error { } return err } + +// defaultDownloadBudget is aws.DefaultDownloadLimits.BytesInFlight as the flag +// spells it; a test keeps the two equal. +const defaultDownloadBudget = "512M" + +func (o *snapshotS3Options) resolveDownloadLimits() error { + if o.downloadConcurrency < 1 { + return fmt.Errorf("--download-concurrency must be at least 1, got %d", o.downloadConcurrency) + } + budget, err := parseByteSize(o.downloadBudget) + if err != nil { + return fmt.Errorf("invalid --download-budget: %w", err) + } + o.downloadLimits = aws.DownloadLimits{Concurrency: o.downloadConcurrency, BytesInFlight: budget} + return nil +} diff --git a/cmd/kosli/snapshotS3_test.go b/cmd/kosli/snapshotS3_test.go index b83af80dc..4932700f0 100644 --- a/cmd/kosli/snapshotS3_test.go +++ b/cmd/kosli/snapshotS3_test.go @@ -113,6 +113,34 @@ func (suite *SnapshotS3TestSuite) TestSnapshotS3Cmd() { cmd: fmt.Sprintf(`snapshot s3 %s %s --bucket %s --exclude dummy`, suite.envName, suite.defaultKosliArguments, suite.bucketName), golden: "bucket kosli-cli-public was reported to environment snapshot-s3-env\n", }, + { + name: "download limits can be set, with a bare number read as megabytes", + cmd: fmt.Sprintf(`snapshot s3 %s %s --bucket %s --download-concurrency 2 --download-budget 64`, suite.envName, suite.defaultKosliArguments, suite.bucketName), + golden: "bucket kosli-cli-public was reported to environment snapshot-s3-env\n", + }, + { + name: "the download budget takes a unit suffix", + cmd: fmt.Sprintf(`snapshot s3 %s %s --bucket %s --download-budget 2GB`, suite.envName, suite.defaultKosliArguments, suite.bucketName), + golden: "bucket kosli-cli-public was reported to environment snapshot-s3-env\n", + }, + { + wantError: true, + name: "snapshot s3 fails if --download-concurrency is below 1", + cmd: fmt.Sprintf(`snapshot s3 %s %s --bucket %s --download-concurrency 0`, suite.envName, suite.defaultKosliArguments, suite.bucketName), + golden: "Error: --download-concurrency must be at least 1, got 0\n", + }, + { + wantError: true, + name: "snapshot s3 fails if --download-budget is not a size", + cmd: fmt.Sprintf(`snapshot s3 %s %s --bucket %s --download-budget large`, suite.envName, suite.defaultKosliArguments, suite.bucketName), + golden: "Error: invalid --download-budget: \"large\" is not a size: expected a number with an optional K, M, G or T unit, e.g. 512M\n", + }, + { + wantError: true, + name: "snapshot s3 fails if --download-budget is zero", + cmd: fmt.Sprintf(`snapshot s3 %s %s --bucket %s --download-budget 0`, suite.envName, suite.defaultKosliArguments, suite.bucketName), + golden: "Error: invalid --download-budget: size \"0\" must be at least 1 byte\n", + }, } for _, t := range tests { diff --git a/cmd/kosli/testdata/empty-flag-audit-coverage.json b/cmd/kosli/testdata/empty-flag-audit-coverage.json index fcb151ca5..6a40d2f23 100644 --- a/cmd/kosli/testdata/empty-flag-audit-coverage.json +++ b/cmd/kosli/testdata/empty-flag-audit-coverage.json @@ -874,6 +874,8 @@ "aws-region": "string", "aws-secret-key": "string", "bucket": "string", + "download-budget": "string", + "download-concurrency": "int", "dry-run": "bool", "exclude": "stringSlice", "exclude-regex": "stringSlice", diff --git a/docs/adr/20260911-s3-fingerprint-from-virtual-tree.md b/docs/adr/20260911-s3-fingerprint-from-virtual-tree.md index f5ed44e3e..5487d1b0a 100644 --- a/docs/adr/20260911-s3-fingerprint-from-virtual-tree.md +++ b/docs/adr/20260911-s3-fingerprint-from-virtual-tree.md @@ -1,7 +1,7 @@ --- title: "20260911 - Fingerprint S3 buckets from a virtual tree; object keys never become local paths" description: "Download each object to an anonymous temp file, hash it, delete it, and compute the directory fingerprint from (key, sha256) pairs so that no S3 key is ever used as a filename" -status: "Proposed" +status: "Accepted" date: "2026-09-11" --- @@ -37,7 +37,7 @@ The fingerprint format itself is fixed. An S3 snapshot must match the fingerprin 4. **A root `.kosli_ignore` is honoured virtually.** Its rules are parsed by `digest.ParseIgnoreRules`, the same reading `DirSha256` gives the file, and resolved by `digest/virtualglob.go`, which reproduces `filepathx.Glob`, `filepath.Glob` and `filepath.Walk` step for step over the virtual tree rather than reimplementing what the globs appear to mean. That is what keeps their quirks identical: a literal `**/x` finds a root `x` spelled with a double slash, which the walk's cleaned paths never equal, so a root file `x` survives while a root directory `x` keeps its name and loses its contents; `**/*.log` is rebuilt cleaned and matches outright; and excluding `logs/*` leaves an empty directory whose name is still hashed. Exclusion therefore runs inside the tree walk, not by filtering the file list. The ignore file can never exclude itself, as in `DirSha256`. `digest.FilesNeedingContent` shares that walk so excluded objects are not downloaded at all, and `VirtualDirSha256` refuses a tree that needs a digest it was not given, so a skipped download can never leak into a fingerprint. Equivalence with `DirSha256` on a materialised tree is asserted for every rule shape in `TestVirtualIgnoreTestSuite`. -5. **Downloads may run in parallel**, behind a fixed worker pool and a bytes-in-flight budget derived from listing sizes, with results written by index for determinism and a cancellable context so the first transport error stops in-flight multipart downloads. This is a performance property, not a safety one, and is delivered separately from the change this record describes; see #1167. Until it lands, downloads are sequential, exactly as before, and peak temp disk is one object. +5. **Downloads may run in parallel**, behind a fixed worker pool and a bytes-in-flight budget derived from listing sizes, with results written by index for determinism and a cancellable context so the first transport error stops in-flight multipart downloads. This is a performance property, not a safety one, and was delivered separately from the change this record describes; see #1167. The defaults are eight objects in flight within 512 MB of listed bytes, tunable with `--download-concurrency` and `--download-budget`, so peak temp disk is the budget rather than one object. 6. **The switch is pinned, not argued.** `TestPinnedFingerprints` in `internal/aws` holds fingerprints of representative fake buckets recorded against the key-layout implementation before it was replaced: unusual key shapes, a `.` sorting before a `/`, nested prefixes with folder markers, and a single object with its basename as artifact name. It was green before the switch and must stay green after it, alongside `TestGetS3DataFromClientKeepsTodaysLayoutForUnusualKeys` from #1155. @@ -62,6 +62,6 @@ The attest side is unchanged: `kosli attest artifact --artifact-type dir` still - `localPathForS3Key`, `filepath.IsLocal`, the `O_EXCL` open, the `ENOTDIR` branch, `containsSingleFile` and the platform-conditional tests from #1155 are deleted. No code in the S3 path branches on the operating system any more. The codebase does not get smaller, though: the virtual tree, the key rule with its collision reporting, and above all the faithful simulation of `filepathx.Glob` add several hundred lines, most of them owed to reproducing `.kosli_ignore` semantics exactly. That is the price of the compatibility contract, paid once and shared with #1069. - The bucket's ignore file is recognised by the exact key `.kosli_ignore`. On disk, `ignoreFilePathInTree` also accepted a case-folded spelling such as `.KOSLI_IGNORE` where the operator's filesystem folded case, so on macOS or Windows such a bucket had its rules applied; now it does not, its exclusions stop applying and its fingerprint moves to the Linux value. Same correction as the key case above, and release-noted with it. - `...`, `.. `, `CON`, colon and backslash keys snapshot again. `..` segments and colliding keys remain errors and now name every key involved. -- Excluded and ignored objects are not downloaded. Each object's temp file is removed once hashed, so peak temp disk falls from the whole bucket to one object, and to the configured budget once #1167 lands. The other half of that trade is memory: the old loop streamed the listing a page at a time, while a tree cannot be fingerprinted without holding it, so peak memory rises from one listing page to every object that survives the filters, at a few hundred bytes per object on top of its key across the listing, the key-to-path maps, the manifest and the tree. A million objects is a few hundred megabytes. #1167 sizes its budget against this figure. +- Excluded and ignored objects are not downloaded. Each object's temp file is removed once hashed, so peak temp disk falls from the whole bucket to one object, and to the configured budget now that #1167 has landed. The other half of that trade is memory: the old loop streamed the listing a page at a time, while a tree cannot be fingerprinted without holding it, so peak memory rises from one listing page to every object that survives the filters, at a few hundred bytes per object on top of its key across the listing, the key-to-path maps, the manifest and the tree. A million objects is a few hundred megabytes. #1167 sized its budget against this figure. - #1069 rebases onto the shared layer: metadata mode becomes a sha256 source plugged into the same list, normalise, exclude, tree pipeline, its key rule disappears in favour of rule 2, and its rejection of buckets with a root `.kosli_ignore` becomes a download of that one object. -- Delivery is two pull requests. The first carries this decision with sequential downloads, so the security-relevant review is not mixed with performance work. The second, #1167, adds parallel downloads, the byte budget and the flags that tune them. +- Delivery was two pull requests. The first, #1180, carried this decision with sequential downloads, so the security-relevant review was not mixed with performance work. The second, #1167, added parallel downloads, the byte budget and the flags that tune them. diff --git a/go.mod b/go.mod index 71550fa26..37e4bbdab 100644 --- a/go.mod +++ b/go.mod @@ -50,6 +50,7 @@ require ( github.com/zalando/go-keyring v0.2.8 gitlab.com/gitlab-org/api/client-go v1.46.0 golang.org/x/oauth2 v0.37.0 + golang.org/x/sync v0.22.0 golang.org/x/term v0.46.0 google.golang.org/api v0.297.0 google.golang.org/grpc v1.83.2 @@ -249,7 +250,6 @@ require ( golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f // indirect golang.org/x/mod v0.38.0 // indirect golang.org/x/net v0.58.0 // indirect - golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.48.0 // indirect golang.org/x/text v0.41.0 // indirect golang.org/x/time v0.15.0 // indirect diff --git a/internal/aws/aws.go b/internal/aws/aws.go index 8787dffbf..c2f9d7164 100644 --- a/internal/aws/aws.go +++ b/internal/aws/aws.go @@ -27,6 +27,7 @@ import ( "github.com/kosli-dev/cli/internal/digest" "github.com/kosli-dev/cli/internal/filters" "github.com/kosli-dev/cli/internal/logger" + "golang.org/x/sync/semaphore" ) // EcsEnvRequest represents the PUT request body to be sent to kosli from ECS @@ -169,7 +170,11 @@ func defaultNewS3Client(creds *AWSStaticCreds) (S3API, error) { if err != nil { return nil, err } - return &s3Client{S3ListAPI: client, S3DownloadAPI: transfermanager.New(client)}, nil + // Five parts per object is the SDK's default, pinned so the connection count, + // objects in flight times parts, cannot move with an SDK upgrade. + return &s3Client{S3ListAPI: client, S3DownloadAPI: transfermanager.New(client, func(o *transfermanager.Options) { + o.Concurrency = 5 + })}, nil } // NewS3ClientFunc is the factory used by GetS3Data to create an S3API client. @@ -441,16 +446,16 @@ func objectMatchesFilter(key string, paths []string, patterns []*regexp.Regexp) // includePaths / excludePaths match object keys by literal prefix. // includeRegex / excludeRegex match object keys by Go regular expression. // Include and exclude filters are mutually exclusive (callers enforce this). -func (staticCreds *AWSStaticCreds) GetS3Data(bucket string, includePaths, includeRegex, excludePaths, excludeRegex []string, logger *logger.Logger) ([]*S3Data, error) { +func (staticCreds *AWSStaticCreds) GetS3Data(bucket string, includePaths, includeRegex, excludePaths, excludeRegex []string, limits DownloadLimits, logger *logger.Logger) ([]*S3Data, error) { client, err := NewS3ClientFunc(staticCreds) if err != nil { return []*S3Data{}, err } - return getS3DataFromClient(client, bucket, includePaths, includeRegex, excludePaths, excludeRegex, logger) + return getS3DataFromClient(client, bucket, includePaths, includeRegex, excludePaths, excludeRegex, limits, logger) } // getS3DataFromClient harvests bucket content using the provided S3API client. -func getS3DataFromClient(client S3API, bucket string, includePaths, includeRegex, excludePaths, excludeRegex []string, logger *logger.Logger) ([]*S3Data, error) { +func getS3DataFromClient(client S3API, bucket string, includePaths, includeRegex, excludePaths, excludeRegex []string, limits DownloadLimits, logger *logger.Logger) ([]*S3Data, error) { s3Data := []*S3Data{} includeRegexCompiled, err := compilePathRegex(includeRegex) @@ -480,7 +485,7 @@ func getS3DataFromClient(client S3API, bucket string, includePaths, includeRegex return s3Data, fmt.Errorf("bucket [%s] reported no modification time for any matching object", bucket) } - artifactName, sha256, err := fingerprintS3Objects(client, bucket, objects, logger) + artifactName, sha256, err := fingerprintS3Objects(client, bucket, objects, limits, logger) if err != nil { return s3Data, err } @@ -494,8 +499,26 @@ func getS3DataFromClient(client S3API, bucket string, includePaths, includeRegex type s3Object struct { key string lastModified time.Time + size int64 } +// DownloadLimits bounds the object downloads in flight at once when +// fingerprinting a bucket. +type DownloadLimits struct { + // Concurrency is the number of objects downloading at the same time. Each + // one may buffer up to five 8 MiB parts in memory while it writes, so memory + // rises with this figure independently of BytesInFlight. + Concurrency int + // BytesInFlight caps the sum of the listed sizes of the objects downloading + // at the same time, and so the temp disk they occupy. An object larger than + // the whole budget downloads alone. + BytesInFlight int64 +} + +// DefaultDownloadLimits keeps peak temp disk near half a gigabyte, which fits +// Lambda's default /tmp, and part buffers near 320 MiB of memory. +var DefaultDownloadLimits = DownloadLimits{Concurrency: 8, BytesInFlight: 512 << 20} + // listMatchingS3Objects lists the bucket, dropping folder markers and keys the // filters exclude, in the order S3 returns them. func listMatchingS3Objects(client S3ListAPI, bucket string, includePaths []string, includeRegex []*regexp.Regexp, @@ -532,11 +555,14 @@ func listMatchingS3Objects(client S3ListAPI, bucket string, includePaths []strin seen[*object.Key] = true // An object without a timestamp stays in the fingerprint and out of // the snapshot timestamp, as it was before. - var lastModified time.Time + listed := s3Object{key: *object.Key} if object.LastModified != nil { - lastModified = *object.LastModified + listed.lastModified = *object.LastModified + } + if object.Size != nil { + listed.size = *object.Size } - objects = append(objects, s3Object{key: *object.Key, lastModified: lastModified}) + objects = append(objects, listed) } } return objects, nil @@ -550,8 +576,9 @@ func listMatchingS3Objects(client S3ListAPI, bucket string, includePaths []strin // is fingerprinted as that file and named after it, as before. // // A root .kosli_ignore is downloaded first so its rules can be applied, and -// objects the rules exclude are not downloaded at all. -func fingerprintS3Objects(downloader S3DownloadAPI, bucket string, objects []s3Object, logger *logger.Logger) (string, string, error) { +// objects the rules exclude are not downloaded at all. The remaining objects +// download in parallel within limits; the first failure cancels the rest. +func fingerprintS3Objects(downloader S3DownloadAPI, bucket string, objects []s3Object, limits DownloadLimits, logger *logger.Logger) (string, string, error) { keys := make([]string, len(objects)) for i, object := range objects { keys[i] = object.key @@ -571,8 +598,7 @@ func fingerprintS3Objects(downloader S3DownloadAPI, bucket string, objects []s3O } }() - // The manifest starts as paths only; digests are filled in below by index, - // so it stays in listing order. + // The manifest starts as paths only; digests are filled in by index below. files := make([]digest.VirtualFile, len(objects)) for i, object := range objects { files[i].Path = paths[object.key] @@ -581,7 +607,7 @@ func fingerprintS3Objects(downloader S3DownloadAPI, bucket string, objects []s3O // One object is fingerprinted as that file and named after it, as it was // when the objects were laid out on disk. if file, ok := digest.SingleVirtualFile(files); ok { - sha256, err := downloadAndHashS3Object(downloader, tempDir, bucket, objects[0].key, nil, logger) + sha256, err := downloadAndHashS3Object(context.TODO(), downloader, tempDir, bucket, objects[0].key, nil, logger) if err != nil { return "", "", err } @@ -594,7 +620,7 @@ func fingerprintS3Objects(downloader S3DownloadAPI, bucket string, objects []s3O if paths[key] != digest.IgnoreFileName { continue } - sha256, err := downloadAndHashS3Object(downloader, tempDir, bucket, key, func(file *os.File) error { + sha256, err := downloadAndHashS3Object(context.TODO(), downloader, tempDir, bucket, key, func(file *os.File) error { if _, err := file.Seek(0, io.SeekStart); err != nil { return err } @@ -617,21 +643,25 @@ func fingerprintS3Objects(downloader S3DownloadAPI, bucket string, objects []s3O return "", "", ignoreRuleError(err) } + var toDownload []int for i, object := range objects { sha256, downloaded := contentSha256[object.key] switch { case downloaded: // the ignore-file pass already hashed this object case needed[files[i].Path]: - sha256, err = downloadAndHashS3Object(downloader, tempDir, bucket, object.key, nil, logger) - if err != nil { - return "", "", err - } + toDownload = append(toDownload, i) default: logger.Debug("object key [%s] is excluded by %s and is not downloaded", object.key, digest.IgnoreFileName) } files[i].Sha256 = sha256 } + // Each download writes its own slot, so the manifest stays in listing order + // however the downloads interleave. + if err := downloadS3ObjectsInParallel(downloader, tempDir, bucket, objects, toDownload, files, limits, logger); err != nil { + return "", "", err + } + sha256, err := digest.VirtualDirSha256(files, rules, logger) if err != nil { return "", "", ignoreRuleError(err) @@ -648,11 +678,73 @@ func ignoreRuleError(err error) error { return err } +// downloadS3ObjectsInParallel fetches the objects at indexes and writes each +// digest into files at the same index. A fixed worker pool bounds downloads and +// goroutines alike, a weighted semaphore bounds their listed bytes, and the +// first error cancels the context so nothing further starts. +func downloadS3ObjectsInParallel(downloader S3DownloadAPI, tempDir, bucket string, objects []s3Object, indexes []int, + files []digest.VirtualFile, limits DownloadLimits, logger *logger.Logger) error { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + budget := semaphore.NewWeighted(max(limits.BytesInFlight, 1)) + firstErr := make(chan error, 1) + fail := func(err error) { + select { + case firstErr <- err: + default: // an earlier failure is already recorded + } + cancel() + } + + work := make(chan int) + var wg sync.WaitGroup + for range min(max(limits.Concurrency, 1), len(indexes)) { + wg.Add(1) + go func() { + defer wg.Done() + for i := range work { + object := objects[i] + // An object larger than the budget takes all of it and so runs alone. + weight := max(min(object.size, limits.BytesInFlight), 1) + if err := budget.Acquire(ctx, weight); err != nil { + return // cancelled while waiting + } + sha256, err := downloadAndHashS3Object(ctx, downloader, tempDir, bucket, object.key, nil, logger) + budget.Release(weight) + if err != nil { + fail(err) + return + } + files[i].Sha256 = sha256 + } + }() + } + +feed: + for _, i := range indexes { + select { + case work <- i: + case <-ctx.Done(): + break feed + } + } + close(work) + wg.Wait() + + select { + case err := <-firstErr: + return err + default: + return nil + } +} + // downloadAndHashS3Object fetches one object into a fresh temp file, lets // inspect read it when given, returns the sha256 of its content and removes the // file. The file's name comes from the OS, so nothing about the key reaches the // filesystem. -func downloadAndHashS3Object(downloader S3DownloadAPI, tempDir, bucket, key string, inspect func(*os.File) error, logger *logger.Logger) (string, error) { +func downloadAndHashS3Object(ctx context.Context, downloader S3DownloadAPI, tempDir, bucket, key string, inspect func(*os.File) error, logger *logger.Logger) (string, error) { file, err := os.CreateTemp(tempDir, "object-*") if err != nil { return "", fmt.Errorf("object key [%s]: %w", key, err) @@ -667,7 +759,7 @@ func downloadAndHashS3Object(downloader S3DownloadAPI, tempDir, bucket, key stri } }() - result, err := downloader.DownloadObject(context.TODO(), &transfermanager.DownloadObjectInput{ + result, err := downloader.DownloadObject(ctx, &transfermanager.DownloadObjectInput{ Bucket: aws.String(bucket), Key: aws.String(key), WriterAt: file, diff --git a/internal/aws/aws_test.go b/internal/aws/aws_test.go index 6c6412a15..afc3bacaa 100644 --- a/internal/aws/aws_test.go +++ b/internal/aws/aws_test.go @@ -433,7 +433,7 @@ func (suite *AWSTestSuite) TestGetS3Data() { } { suite.Run(t.name, func() { skipIfCredsUnset(suite.T(), t.requireEnvVars, t.creds) - data, err := t.creds.GetS3Data(t.bucketName, t.includePaths, nil, t.excludePaths, nil, logger.NewStandardLogger()) + data, err := t.creds.GetS3Data(t.bucketName, t.includePaths, nil, t.excludePaths, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.False(suite.T(), (err != nil) != t.wantErr, "GetS3Data() error = %v, wantErr %v", err, t.wantErr) if !t.wantErr { @@ -1178,7 +1178,7 @@ func (suite *AWSTestSuite) TestGetS3DataFromClient() { } data, err := getS3DataFromClient(client, fakeS3TestBucketName, t.includePaths, - t.includeRegex, t.excludePaths, t.excludeRegex, logger.NewStandardLogger()) + t.includeRegex, t.excludePaths, t.excludeRegex, DefaultDownloadLimits, logger.NewStandardLogger()) if t.wantErr { require.Error(suite.T(), err) @@ -1221,7 +1221,7 @@ func (suite *AWSTestSuite) TestGetS3DataFromClientFilterEquivalence() { fingerprint := func(objects map[string][]byte, includePaths, includeRegex, excludePaths, excludeRegex []string) string { client := &FakeS3Client{Bucket: fakeS3TestBucketName, Objects: objects} data, err := getS3DataFromClient(client, fakeS3TestBucketName, includePaths, - includeRegex, excludePaths, excludeRegex, logger.NewStandardLogger()) + includeRegex, excludePaths, excludeRegex, DefaultDownloadLimits, logger.NewStandardLogger()) require.NoError(suite.T(), err) require.Len(suite.T(), data, 1) require.Len(suite.T(), data[0].Digests, 1) @@ -1292,7 +1292,7 @@ func (suite *AWSTestSuite) TestGetS3DataFromClientRejectsKeysWithDotDotSegments( }, } - _, err := getS3DataFromClient(poisoned, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + _, err := getS3DataFromClient(poisoned, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.Error(suite.T(), err, "a key containing a \"..\" segment must fail the snapshot instead of silently overwriting another object's download") require.Contains(suite.T(), err.Error(), "uploads/user-a/../../protected/release.bin") } @@ -1307,7 +1307,7 @@ func (suite *AWSTestSuite) TestGetS3DataFromClientCollidingKeysAreAnError() { }, } - _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.Error(suite.T(), err) require.Contains(suite.T(), err.Error(), "[a//b]", "the error must name both colliding keys") require.Contains(suite.T(), err.Error(), "[a/b]", "the error must name both colliding keys") @@ -1325,7 +1325,7 @@ func (suite *AWSTestSuite) TestGetS3DataFromClientObjectAndPrefixCollideAreAnErr }, } - _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.Error(suite.T(), err) require.Contains(suite.T(), err.Error(), "object key [a]", "the error must name the object") require.Contains(suite.T(), err.Error(), "object key [a/b]", "the error must name an object under the prefix") @@ -1358,9 +1358,9 @@ func (suite *AWSTestSuite) TestGetS3DataFromClientKeepsTodaysLayoutForUnusualKey }, } - unusualData, err := getS3DataFromClient(unusual, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + unusualData, err := getS3DataFromClient(unusual, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.NoError(suite.T(), err) - todayData, err := getS3DataFromClient(today, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + todayData, err := getS3DataFromClient(today, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.NoError(suite.T(), err) require.Len(suite.T(), unusualData, 1) require.Len(suite.T(), todayData, 1) diff --git a/internal/aws/s3_fingerprint_test.go b/internal/aws/s3_fingerprint_test.go index 4fa6e8240..d05066cf8 100644 --- a/internal/aws/s3_fingerprint_test.go +++ b/internal/aws/s3_fingerprint_test.go @@ -5,6 +5,7 @@ import ( "os" "path/filepath" "sort" + "strings" "sync" "testing" "time" @@ -59,7 +60,7 @@ func (r *recordingDownloader) downloadedKeys() []string { func snapshotFake(t *testing.T, client S3API) (artifactName, fingerprint string) { t.Helper() - data, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + data, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.NoError(t, err) require.Len(t, data, 1) require.Len(t, data[0].Digests, 1) @@ -191,14 +192,14 @@ func (suite *S3FingerprintTestSuite) TestDownloadsExactlyTheContributingObjects( "scratch.tmp": []byte("tmp"), "filtered/out.txt": []byte("out"), }}} - data, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, []string{"filtered/"}, nil, logger.NewStandardLogger()) + data, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, []string{"filtered/"}, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.NoError(suite.T(), err) require.Len(suite.T(), data, 1) require.Equal(suite.T(), []string{".kosli_ignore", "app.js", "lib/util.js"}, client.downloadedKeys()) } -// Objects are downloaded to files whose names owe nothing to the key, each is -// removed once hashed, and the download directory is gone at the end. +// How many temp files exist at once is the byte budget's concern, tested in +// S3ParallelTestSuite; this test checks only their names and their removal. func (suite *S3FingerprintTestSuite) TestObjectsNeverLandUnderTheirKeyAndDoNotLinger() { keys := []string{"alpha.bin", "beta/gamma.bin", "delta/epsilon/zeta.bin"} objects := map[string][]byte{} @@ -206,21 +207,27 @@ func (suite *S3FingerprintTestSuite) TestObjectsNeverLandUnderTheirKeyAndDoNotLi objects[key] = []byte(key) } client := &recordingDownloader{S3API: &FakeS3Client{Bucket: fakeS3TestBucketName, Objects: objects}} + // Downloads run concurrently, so onDownload fires from worker goroutines; + // require.* must run only on the test goroutine, so record and assert after. + var mu sync.Mutex + var nilFile bool + var keyLikeNames []string client.onDownload = func(key string, file *os.File) { - require.NotNil(suite.T(), file, "the transfer manager must be handed a real file") - require.NotContains(suite.T(), filepath.Base(file.Name()), filepath.Base(key), - "the local file name must owe nothing to the key") - client.mu.Lock() - earlier := append([]string{}, client.files[:len(client.files)-1]...) - client.mu.Unlock() - for _, previous := range earlier { - _, err := os.Stat(previous) - require.ErrorIs(suite.T(), err, os.ErrNotExist, "an earlier object's file must be gone before the next download") + mu.Lock() + defer mu.Unlock() + if file == nil { + nilFile = true + return + } + if strings.Contains(filepath.Base(file.Name()), filepath.Base(key)) { + keyLikeNames = append(keyLikeNames, key) } } - _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.NoError(suite.T(), err) + require.False(suite.T(), nilFile, "the transfer manager must be handed a real file") + require.Empty(suite.T(), keyLikeNames, "the local file name must owe nothing to the key") require.Len(suite.T(), client.files, len(keys)) for _, file := range client.files { _, err := os.Stat(file) @@ -238,7 +245,7 @@ func (suite *S3FingerprintTestSuite) TestAMalformedIgnoreRuleFailsTheSnapshot() client := &FakeS3Client{Bucket: fakeS3TestBucketName, Objects: map[string][]byte{ ".kosli_ignore": []byte(rule + "\n"), "app.js": []byte("app"), }} - _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.Error(suite.T(), err) require.Contains(suite.T(), err.Error(), "the bucket's .kosli_ignore holds a rule that cannot be applied") require.Contains(suite.T(), err.Error(), rule) @@ -250,10 +257,11 @@ func (suite *S3FingerprintTestSuite) TestADownloadErrorNamesTheKey() { client := &FakeS3Client{Bucket: fakeS3TestBucketName, Objects: map[string][]byte{ "README.md": []byte(fakeReadmeBody), "notes.txt": []byte(fakeNotesBody), }, DownloadObjectErr: os.ErrDeadlineExceeded} - _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.Error(suite.T(), err) require.ErrorIs(suite.T(), err, os.ErrDeadlineExceeded) - require.Contains(suite.T(), err.Error(), "object key [README.md]") + // Downloads overlap, so either object may be the first to fail. + require.Regexp(suite.T(), `object key \[(README\.md|notes\.txt)\]`, err.Error()) require.NotContains(suite.T(), err.Error(), "--exclude-regex", "a transport failure must not advise dropping the object") } @@ -263,7 +271,7 @@ func (suite *S3FingerprintTestSuite) TestAHashErrorNamesTheKey() { client := &FakeS3Client{Bucket: fakeS3TestBucketName, Objects: map[string][]byte{"README.md": []byte(fakeReadmeBody)}} // Deleting the file between download and hash is the one way to make the // hash fail without touching permissions. - _, err := downloadAndHashS3Object(client, suite.T().TempDir(), fakeS3TestBucketName, "README.md", func(file *os.File) error { + _, err := downloadAndHashS3Object(context.TODO(), client, suite.T().TempDir(), fakeS3TestBucketName, "README.md", func(file *os.File) error { return os.Remove(file.Name()) }, logger.NewStandardLogger()) require.Error(suite.T(), err) @@ -280,19 +288,19 @@ func (suite *S3FingerprintTestSuite) TestAListingWithoutModificationTimesDoesNot objects := map[string][]byte{"README.md": []byte(fakeReadmeBody), "notes.txt": []byte(fakeNotesBody)} later := fakeS3LastModified.Add(time.Hour) full, err := getS3DataFromClient(&FakeS3Client{Bucket: fakeS3TestBucketName, Objects: objects, - LastModified: map[string]time.Time{"notes.txt": later}}, fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + LastModified: map[string]time.Time{"notes.txt": later}}, fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.NoError(suite.T(), err) partial, err := getS3DataFromClient(&FakeS3Client{Bucket: fakeS3TestBucketName, Objects: objects, LastModified: map[string]time.Time{"notes.txt": later}, NoLastModified: map[string]bool{"README.md": true}}, - fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.NoError(suite.T(), err) require.Equal(suite.T(), full[0].Digests, partial[0].Digests, "the object without a timestamp stays in the fingerprint") require.Equal(suite.T(), later.Unix(), partial[0].LastModifiedTimestamp) _, err = getS3DataFromClient(&FakeS3Client{Bucket: fakeS3TestBucketName, Objects: objects, NoLastModified: map[string]bool{"README.md": true, "notes.txt": true}}, - fakeS3TestBucketName, nil, nil, nil, nil, logger.NewStandardLogger()) + fakeS3TestBucketName, nil, nil, nil, nil, DefaultDownloadLimits, logger.NewStandardLogger()) require.Error(suite.T(), err) require.Contains(suite.T(), err.Error(), "modification time") } diff --git a/internal/aws/s3_parallel_test.go b/internal/aws/s3_parallel_test.go new file mode 100644 index 000000000..dfb630141 --- /dev/null +++ b/internal/aws/s3_parallel_test.go @@ -0,0 +1,286 @@ +package aws + +import ( + "context" + "errors" + "fmt" + "math/rand" + "path/filepath" + "runtime" + "sync" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/feature/s3/transfermanager" + "github.com/kosli-dev/cli/internal/digest" + "github.com/kosli-dev/cli/internal/logger" + "github.com/kosli-dev/cli/internal/utils" + "github.com/stretchr/testify/require" + "github.com/stretchr/testify/suite" +) + +type S3ParallelTestSuite struct { + suite.Suite + lock sync.Mutex +} + +// trackingDownloader records peak downloads and listed bytes in flight, and +// calls per key. delay keeps downloads overlapping so the bounds are exercised. +type trackingDownloader struct { + S3API + sizes map[string]int64 + delay time.Duration + // hook runs before the delegate and can fail the download in its place. + hook func(ctx context.Context, key string) error + + mu sync.Mutex + inFlight int + maxInFlight int + bytesInFlight int64 + maxBytesInFlight int64 + calls map[string]int +} + +func (d *trackingDownloader) DownloadObject(ctx context.Context, params *transfermanager.DownloadObjectInput, optFns ...func(*transfermanager.Options)) (*transfermanager.DownloadObjectOutput, error) { + key := *params.Key + d.mu.Lock() + if d.calls == nil { + d.calls = map[string]int{} + } + d.calls[key]++ + d.inFlight++ + d.bytesInFlight += d.sizes[key] + d.maxInFlight = max(d.maxInFlight, d.inFlight) + d.maxBytesInFlight = max(d.maxBytesInFlight, d.bytesInFlight) + d.mu.Unlock() + defer func() { + d.mu.Lock() + d.inFlight-- + d.bytesInFlight -= d.sizes[key] + d.mu.Unlock() + }() + + time.Sleep(d.delay) + if d.hook != nil { + if err := d.hook(ctx, key); err != nil { + return nil, err + } + } + return d.S3API.DownloadObject(ctx, params, optFns...) +} + +func (d *trackingDownloader) currentInFlight() int { + d.mu.Lock() + defer d.mu.Unlock() + return d.inFlight +} + +// bucketOf fakes n objects of size bytes each, spread over four prefixes. +func bucketOf(n int, size int, delay time.Duration) (*trackingDownloader, []s3Object) { + objects := map[string][]byte{} + listing := make([]s3Object, 0, n) + sizes := map[string]int64{} + for i := 0; i < n; i++ { + key := fmt.Sprintf("dir%d/object-%03d.bin", i%4, i) + body := []byte(fmt.Sprintf("%0*d", size, i)) + objects[key] = body + sizes[key] = int64(len(body)) + listing = append(listing, s3Object{key: key, size: int64(len(body))}) + } + fake := &FakeS3Client{Bucket: fakeS3TestBucketName, Objects: objects} + return &trackingDownloader{S3API: fake, sizes: sizes, delay: delay}, listing +} + +func (suite *S3ParallelTestSuite) TestMakesExactlyOneCallPerObject() { + client, listing := bucketOf(60, 8, 0) + _, parallel, err := fingerprintS3Objects(client, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: 8, BytesInFlight: 1 << 30}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.Len(suite.T(), client.calls, 60) + for key, n := range client.calls { + require.Equal(suite.T(), 1, n, "key %s", key) + } + + sequential, _ := bucketOf(60, 8, 0) + _, want, err := fingerprintS3Objects(sequential, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: 1, BytesInFlight: 1 << 30}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.Equal(suite.T(), want, parallel) +} + +func (suite *S3ParallelTestSuite) TestRespectsTheConcurrencyBound() { + client, listing := bucketOf(40, 8, 5*time.Millisecond) + _, _, err := fingerprintS3Objects(client, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: 4, BytesInFlight: 1 << 30}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.LessOrEqual(suite.T(), client.maxInFlight, 4) + require.GreaterOrEqual(suite.T(), client.maxInFlight, 2, "downloads must actually overlap for the bound to be tested") +} + +// The byte budget binds before the count bound here: eight slots would allow +// eight 100-byte objects, the budget allows two. +func (suite *S3ParallelTestSuite) TestRespectsTheByteBudget() { + client, listing := bucketOf(20, 100, 5*time.Millisecond) + _, _, err := fingerprintS3Objects(client, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: 8, BytesInFlight: 250}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.LessOrEqual(suite.T(), client.maxBytesInFlight, int64(250)) + require.LessOrEqual(suite.T(), client.maxInFlight, 2) + require.GreaterOrEqual(suite.T(), client.maxInFlight, 2, "downloads must actually overlap for the budget to be tested") +} + +func (suite *S3ParallelTestSuite) TestAnObjectLargerThanTheBudgetRunsAlone() { + client, listing := bucketOf(12, 100, 5*time.Millisecond) + big := "big/huge.bin" + body := make([]byte, 1000) + client.S3API.(*FakeS3Client).Objects[big] = body + client.sizes[big] = 1000 + listing = append([]s3Object{{key: big, size: 1000}}, listing...) + + var aloneChecks, inFlightDuringBig int + client.hook = func(_ context.Context, key string) error { + if key == big { + seen := client.currentInFlight() + suite.mu().Lock() + aloneChecks++ + inFlightDuringBig = max(inFlightDuringBig, seen) + suite.mu().Unlock() + } + return nil + } + + _, _, err := fingerprintS3Objects(client, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: 8, BytesInFlight: 250}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.Equal(suite.T(), 1, aloneChecks) + require.Equal(suite.T(), 1, client.calls[big]) + require.Equal(suite.T(), 1, inFlightDuringBig, "the oversized object must be the only download in flight") +} + +func (suite *S3ParallelTestSuite) mu() *sync.Mutex { return &suite.lock } + +func (suite *S3ParallelTestSuite) TestFingerprintIsIndependentOfCompletionOrder() { + tree := map[string]string{} + for i := 0; i < 30; i++ { + tree[fmt.Sprintf("d%d/sub%d/f%02d.txt", i%3, i%5, i)] = fmt.Sprintf("content %d", i) + } + root := suite.T().TempDir() + objects := map[string][]byte{} + sizes := map[string]int64{} + listing := []s3Object{} + for p, content := range tree { + require.NoError(suite.T(), utils.CreateFileWithContent(filepath.Join(root, filepath.FromSlash(p)), content)) + objects[p] = []byte(content) + sizes[p] = int64(len(content)) + listing = append(listing, s3Object{key: p, size: int64(len(content))}) + } + want, err := digest.DirSha256(root, nil, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + + random := rand.New(rand.NewSource(5)) + for round := 0; round < 3; round++ { + client := &trackingDownloader{S3API: &FakeS3Client{Bucket: fakeS3TestBucketName, Objects: objects}, sizes: sizes} + client.hook = func(_ context.Context, _ string) error { + suite.mu().Lock() + d := time.Duration(random.Intn(4)) * time.Millisecond + suite.mu().Unlock() + time.Sleep(d) + return nil + } + name, got, err := fingerprintS3Objects(client, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: 8, BytesInFlight: 1 << 30}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.Equal(suite.T(), fakeS3TestBucketName, name) + require.Equal(suite.T(), want, got, "round %d", round) + } +} + +// Workers race for the slots, so the hook fails the first arrival rather than a +// fixed key, and blocks every other one until it is cancelled. +func (suite *S3ParallelTestSuite) TestATransportErrorStopsRemainingWork() { + client, listing := bucketOf(10, 8, 0) + errBoom := errors.New("boom") + var arrivals int + var failingKey string + client.hook = func(ctx context.Context, key string) error { + suite.mu().Lock() + arrivals++ + first := arrivals == 1 + if first { + failingKey = key + } + suite.mu().Unlock() + if first { + time.Sleep(5 * time.Millisecond) // let the second slot fill first + return errBoom + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(5 * time.Second): + return errors.New("the context was never cancelled") + } + } + + _, _, err := fingerprintS3Objects(client, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: 2, BytesInFlight: 1 << 30}, logger.NewStandardLogger()) + require.Error(suite.T(), err) + require.ErrorIs(suite.T(), err, errBoom) + require.Contains(suite.T(), err.Error(), fmt.Sprintf("object key [%s]", failingKey)) + require.Len(suite.T(), client.calls, 2, "only the two downloads in flight at the failure may have started: %v", client.calls) +} + +func (suite *S3ParallelTestSuite) TestLimitsReachTheDownloader() { + for _, limit := range []int{1, 4} { + client, _ := bucketOf(12, 8, 3*time.Millisecond) + _, err := getS3DataFromClient(client, fakeS3TestBucketName, nil, nil, nil, nil, + DownloadLimits{Concurrency: limit, BytesInFlight: 1 << 30}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.Len(suite.T(), client.calls, 12) + require.LessOrEqual(suite.T(), client.maxInFlight, limit) + if limit > 1 { + require.GreaterOrEqual(suite.T(), client.maxInFlight, 2) + } + } +} + +// One goroutine per object would be hundreds of megabytes of idle stacks on a +// large bucket. +func (suite *S3ParallelTestSuite) TestGoroutinesDoNotScaleWithTheBucket() { + const objects, concurrency = 2000, 4 + client, listing := bucketOf(objects, 8, 0) + before := runtime.NumGoroutine() + var peak int + client.hook = func(_ context.Context, _ string) error { + suite.mu().Lock() + peak = max(peak, runtime.NumGoroutine()) + suite.mu().Unlock() + return nil + } + + _, _, err := fingerprintS3Objects(client, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: concurrency, BytesInFlight: 1 << 30}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.Len(suite.T(), client.calls, objects) + require.LessOrEqual(suite.T(), peak, before+concurrency+8, + "goroutines during the run must be bounded by the concurrency, not the object count") +} + +// A concurrency far above the object count must not start idle workers. +func (suite *S3ParallelTestSuite) TestWorkersAreClampedToTheWork() { + client, listing := bucketOf(3, 8, 0) + before := runtime.NumGoroutine() + var peak int + client.hook = func(_ context.Context, _ string) error { + suite.mu().Lock() + peak = max(peak, runtime.NumGoroutine()) + suite.mu().Unlock() + return nil + } + + _, _, err := fingerprintS3Objects(client, fakeS3TestBucketName, listing, DownloadLimits{Concurrency: 50000, BytesInFlight: 1 << 30}, logger.NewStandardLogger()) + require.NoError(suite.T(), err) + require.Len(suite.T(), client.calls, 3) + require.LessOrEqual(suite.T(), peak, before+3+8, "workers must be bounded by the objects to download") +} + +func (suite *S3ParallelTestSuite) TestDefaultLimitsAreSane() { + require.GreaterOrEqual(suite.T(), DefaultDownloadLimits.Concurrency, 2) + require.GreaterOrEqual(suite.T(), DefaultDownloadLimits.BytesInFlight, int64(64<<20)) +} + +func TestS3ParallelTestSuite(t *testing.T) { + suite.Run(t, new(S3ParallelTestSuite)) +}