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

Code before:
		})

		return nil
	}

	fileDownload := &FileDownload{
		Digest:    digest,
		FilePath:  fp,
		Total:     1, // dummy value to indicate that we don't know the total size yet
		Completed: 0,
	}

	_, downloading := inProgress.LoadOrStore(digest, fileDownload)
	if downloading {
		// this is another client requesting the server to download the same blob concurrently
		return monitorDownload(ctx, mp, regOpts, fileDownload, fn)
	}
	resp, err := requestDownload(ctx, mp, regOpts, fileDownload)
	if err != nil {
		return err
	}
	return doDownload(ctx, fileDownload, resp, fn)
}

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 {
				downloadMu.Unlock()
				return fmt.Errorf("resume: %w", err)
			}
			inProgress.Store(f.Digest, f)

Review comment (refactor):
nit: the code in this for loop should be refactored into a function so the unlock can be deferred which will be less error prone

Code after:
		})

		return nil
	}

	fileDownload := &FileDownload{
		Digest:    digest,
		FilePath:  fp,
		Total:     1, // dummy value to indicate that we don't know the total size yet
		Completed: 0,
	}

	_, downloading := inProgress.LoadOrStore(digest, fileDownload)
	if downloading {
		// this is another client requesting the server to download the same blob concurrently
		return monitorDownload(ctx, mp, regOpts, fileDownload, fn)
	}
	return doDownload(ctx, mp, regOpts, fileDownload, fn)
}

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 {
		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),