[s3 cache] Faster download with parallelism

[bpaquet]

No description provided.

[bpaquet]

I will deploy this PR internally to check how it behaves at scale, but a review is always interesting

[tonistiigi]

So how much memory does this use?

I wonder if instead of trying to get this functionality behind ReadAt call (where it really does not belong because it is only for a specific range), this return value could instead implement https://pkg.go.dev/io#WriterTo that could be detected on the caller side where it is writing the blob to disk (maybe even via WriteAt()).

[tonistiigi]

These don't need to be public if only used by private functions.

[bpaquet]

Done

[tonistiigi]

You can test this, but for registry layers, our default is 4.

[bpaquet]

These default are not really used, they are overriden in s3.go, where the default value is 4. I changed them anyway.

[bpaquet]

So how much memory does this use?

I wonder if instead of trying to get this functionality behind ReadAt call (where it really does not belong because it is only for a specific range), this return value could instead implement https://pkg.go.dev/io#WriterTo that could be detected on the caller side where it is writing the blob to disk (maybe even via WriteAt()).

If we do the hypothesis the solver will read faster than we can download from s3, we will use parallelism * part size, so 20MB by default.

We can also build a wrapper to "convert" the WriterAt from the SDK to a ReaderAt, but it seems more complicated to me than this solution, especially from a memory allocation point of view. We can anyway consider is if you think it's better.

[mlevieux]

Suggested change

	downloadParallism = downloadPartSizeInt
	downloadPartSize = downloadPartSizeInt

Right?

[mlevieux]

Nit: since go1.22, you can now use the quite shorter range-over-integer variant:

Suggested change

	for i := 0; i < r.totalChunk; i++ {
	r.inChan <- i
	}
	for k := 0; k < r.input.Parallelism; k++ {
	for i := range r.totalChunk {
	r.inChan <- i
	}
	for range r.input.Parallelism {

See this playground link

[mlevieux]

I would most probably replace done with an unbuffered chan struct{} and go with (untested):

Suggested change

	for {
	if r.chunks[r.currentChunk].done {
	break
	}
	err := <-r.outChan
	if err != nil {
	r.Close()
	return 0, err
	}
	}
	readyOrConsume:
	for {
	select {
	case <-r.chunks[r.currentChunk].done:
	// this chunk is okay, can process
	break readyOrConsome
	case err := <-r.outChan:
	// process error in case there is any
	}
	}

And close the done channel from the downloadChunk function before returning.

Which IMO looks more idiomatic and removes the race on done (even if I assume this particular race wouldn't have any side effect).

[mlevieux]

Nit: for the sake of performance / brevity, I would suggest going with a one-time copy of the result of io.ReadAll(resp.Body) into r.chunks[chunk]. This would avoid the need to reference r.chunks[chunk] to get a valid io.Writer implementation, along with the need for chunkStatus to even implement io.Writer. Also that would remove the need for a writeOffset.

No strong opinion however if you think this is clearer like that.

[mlevieux]

Quick note: this actually declares an anonymous struct field with type io.Writer, which I believe is not what you're looking for, since this field's Write method is shadowed by that of chunkStatus below anyways.

If what you're trying to do is to make explicit and enforce the fact that chunkStatus implements io.Writer, the more idiomatic way to go is to add this line somewhere in the file (most generally this is below the struct declaration itself):

var _ io.Writer = &chunkStatus{}

But Go's interface implementation being implicit, this is not even necessary.

[mlevieux]

Same as for io.Writer above.

[mlevieux]

Also ultra nitpick:

Suggested change

	downloadParallism := 4
	downloadParallelism := 4