Repository: ollama/ollama
File: server/download.go (Go)

Code before:
}

var downloadMu sync.Mutex // mutex to check to resume a download while monitoring

// monitorDownload monitors the download progress of a blob and resumes it if it is interrupted
func monitorDownload(ctx context.Context, mp ModelPath, regOpts *RegistryOptions, f *FileDownload, fn func(api.ProgressResponse)) error {
	tick := time.NewTicker(time.Second)
	for range tick.C {
		downloadMu.Lock()
		val, downloading := inProgress.Load(f.Digest)
		if !downloading {
			// check once again if the download is complete
			if fi, _ := os.Stat(f.FilePath); fi != nil {
				downloadMu.Unlock()
				// successfull download while monitoring
				fn(api.ProgressResponse{
					Digest:    f.Digest,
					Total:     int(fi.Size()),
					Completed: int(fi.Size()),
				})
				return nil
			}
			// resume the download
			resp, err := requestDownload(ctx, mp, regOpts, f)
			if err != nil {
				return fmt.Errorf("resume: %w", err)
			}
			inProgress.Store(f.Digest, f)
			downloadMu.Unlock()
			return doDownload(ctx, f, resp, fn)
		}
		downloadMu.Unlock()
		f, ok := val.(*FileDownload)
		if !ok {
			return fmt.Errorf("invalid type for in progress download: %T", val)
		}
		fn(api.ProgressResponse{
			Status:    fmt.Sprintf("downloading %s", f.Digest),
			Digest:    f.Digest,
			Total:     int(f.Total),
			Completed: int(f.Completed),
		})
	}
	return nil
}

var chunkSize = 1024 * 1024 // 1 MiB in bytes

// requestDownload requests a blob from the registry and returns the response
func requestDownload(ctx context.Context, mp ModelPath, regOpts *RegistryOptions, f *FileDownload) (*http.Response, error) {
	var size int64

Review comment (suggestion):
`downloadMu` doesn't unlock if `requestDownload` errors

Code after:
// monitorDownload monitors the download progress of a blob and resumes it if it is interrupted
func monitorDownload(ctx context.Context, mp ModelPath, regOpts *RegistryOptions, f *FileDownload, fn func(api.ProgressResponse)) error {
	tick := time.NewTicker(time.Second)
	for range tick.C {
		done, resume, err := func() (bool, bool, error) {
			downloadMu.Lock()
			defer downloadMu.Unlock()
			val, downloading := inProgress.Load(f.Digest)
			if !downloading {
				// check once again if the download is complete
				if fi, _ := os.Stat(f.FilePath); fi != nil {
					// successful download while monitoring
					fn(api.ProgressResponse{
						Digest:    f.Digest,
						Total:     int(fi.Size()),
						Completed: int(fi.Size()),
					})
					return true, false, nil
				}
				// resume the download
				inProgress.Store(f.Digest, f) // store the file download again to claim the resume
				return false, true, nil
			}
			f, ok := val.(*FileDownload)
			if !ok {
				return false, false, fmt.Errorf("invalid type for in progress download: %T", val)
			}
			fn(api.ProgressResponse{
				Status:    fmt.Sprintf("downloading %s", f.Digest),
				Digest:    f.Digest,
				Total:     int(f.Total),
				Completed: int(f.Completed),
			})
			return false, false, nil
		}()
		if err != nil {
			return err
		}
		if done {
			// done downloading
			return nil
		}
		if resume {
			return doDownload(ctx, mp, regOpts, f, fn)
		}
	}
	return nil
}

var chunkSize = 1024 * 1024 // 1 MiB in bytes