refactor download run

This commit is contained in:
Michael Yang 2024-01-26 14:18:45 -08:00
parent efe040f8c0
commit 917bd61084

View file

@ -138,16 +138,13 @@ func (b *blobDownload) Prepare(ctx context.Context, requestURL *url.URL, opts *r
} }
func (b *blobDownload) Run(ctx context.Context, requestURL *url.URL, opts *registryOptions) { func (b *blobDownload) Run(ctx context.Context, requestURL *url.URL, opts *registryOptions) {
b.err = b.run(ctx, requestURL, opts)
}
func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *registryOptions) error {
defer blobDownloadManager.Delete(b.Digest) defer blobDownloadManager.Delete(b.Digest)
ctx, b.CancelFunc = context.WithCancel(ctx) ctx, b.CancelFunc = context.WithCancel(ctx)
file, err := os.OpenFile(b.Name+"-partial", os.O_CREATE|os.O_RDWR, 0o644) file, err := os.OpenFile(b.Name+"-partial", os.O_CREATE|os.O_RDWR, 0o644)
if err != nil { if err != nil {
return err b.err = err
return
} }
defer file.Close() defer file.Close()
@ -188,26 +185,30 @@ func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *regis
} }
if err := g.Wait(); err != nil { if err := g.Wait(); err != nil {
return err b.err = err
return
} }
// explicitly close the file so we can rename it // explicitly close the file so we can rename it
if err := file.Close(); err != nil { if err := file.Close(); err != nil {
return err b.err = err
return
} }
for i := range b.Parts { for i := range b.Parts {
if err := os.Remove(file.Name() + "-" + strconv.Itoa(i)); err != nil { if err := os.Remove(file.Name() + "-" + strconv.Itoa(i)); err != nil {
return err b.err = err
return
} }
} }
if err := os.Rename(file.Name(), b.Name); err != nil { if err := os.Rename(file.Name(), b.Name); err != nil {
return err b.err = err
return
} }
b.done = true b.done = true
return nil return
} }
func (b *blobDownload) downloadChunk(ctx context.Context, requestURL *url.URL, w io.Writer, part *blobDownloadPart, opts *registryOptions) error { func (b *blobDownload) downloadChunk(ctx context.Context, requestURL *url.URL, w io.Writer, part *blobDownloadPart, opts *registryOptions) error {