1 files changed, 41 insertions, 21 deletions

diff --git a/src/cmd/go/internal/modfetch/fetch.go b/src/cmd/go/internal/modfetch/fetch.go
index 51a56028c4..d5516dd8d9 100644
--- a/src/cmd/go/internal/modfetch/fetch.go
+++ b/src/cmd/go/internal/modfetch/fetch.go
@@ -7,6 +7,7 @@ package modfetch
 import (
 	"archive/zip"
 	"bytes"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -22,6 +23,7 @@ import (
 	"cmd/go/internal/module"
 	"cmd/go/internal/par"
 	"cmd/go/internal/renameio"
+	"cmd/go/internal/robustio"
 )
 
 var downloadCache par.Cache
@@ -41,24 +43,27 @@ func Download(mod module.Version) (dir string, err error) {
 		err error
 	}
 	c := downloadCache.Do(mod, func() interface{} {
-		dir, err := DownloadDir(mod)
+		dir, err := download(mod)
 		if err != nil {
 			return cached{"", err}
 		}
-		if err := download(mod, dir); err != nil {
-			return cached{"", err}
-		}
 		checkMod(mod)
 		return cached{dir, nil}
 	}).(cached)
 	return c.dir, c.err
 }
 
-func download(mod module.Version, dir string) (err error) {
-	// If the directory exists, the module has already been extracted.
-	fi, err := os.Stat(dir)
-	if err == nil && fi.IsDir() {
-		return nil
+func download(mod module.Version) (dir string, err error) {
+	// If the directory exists, and no .partial file exists,
+	// the module has already been completely extracted.
+	// .partial files may be created when future versions of cmd/go
+	// extract module zip directories in place instead of extracting
+	// to a random temporary directory and renaming.
+	dir, err = DownloadDir(mod)
+	if err == nil {
+		return dir, nil
+	} else if dir == "" || !errors.Is(err, os.ErrNotExist) {
+		return "", err
 	}
 
 	// To avoid cluttering the cache with extraneous files,
@@ -66,7 +71,7 @@ func download(mod module.Version, dir string) (err error) {
 	// Invoke DownloadZip before locking the file.
 	zipfile, err := DownloadZip(mod)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	if cfg.CmdName != "mod download" {
@@ -75,17 +80,19 @@ func download(mod module.Version, dir string) (err error) {
 
 	unlock, err := lockVersion(mod)
 	if err != nil {
-		return err
+		return "", err
 	}
 	defer unlock()
 
 	// Check whether the directory was populated while we were waiting on the lock.
-	fi, err = os.Stat(dir)
-	if err == nil && fi.IsDir() {
-		return nil
+	_, dirErr := DownloadDir(mod)
+	if dirErr == nil {
+		return dir, nil
 	}
+	_, dirExists := dirErr.(*DownloadDirPartialError)
 
-	// Clean up any remaining temporary directories from previous runs.
+	// Clean up any remaining temporary directories from previous runs, as well
+	// as partially extracted diectories created by future versions of cmd/go.
 	// This is only safe to do because the lock file ensures that their writers
 	// are no longer active.
 	parentDir := filepath.Dir(dir)
@@ -95,6 +102,19 @@ func download(mod module.Version, dir string) (err error) {
 			RemoveAll(path) // best effort
 		}
 	}
+	if dirExists {
+		if err := RemoveAll(dir); err != nil {
+			return "", err
+		}
+	}
+
+	partialPath, err := CachePath(mod, "partial")
+	if err != nil {
+		return "", err
+	}
+	if err := os.Remove(partialPath); err != nil && !os.IsNotExist(err) {
+		return "", err
+	}
 
 	// Extract the zip file to a temporary directory, then rename it to the
 	// final path. That way, we can use the existence of the source directory to
@@ -102,11 +122,11 @@ func download(mod module.Version, dir string) (err error) {
 	// the entire directory (e.g. as an attempt to prune out file corruption)
 	// the module cache will still be left in a recoverable state.
 	if err := os.MkdirAll(parentDir, 0777); err != nil {
-		return err
+		return "", err
 	}
 	tmpDir, err := ioutil.TempDir(parentDir, tmpPrefix)
 	if err != nil {
-		return err
+		return "", err
 	}
 	defer func() {
 		if err != nil {
@@ -117,17 +137,17 @@ func download(mod module.Version, dir string) (err error) {
 	modpath := mod.Path + "@" + mod.Version
 	if err := Unzip(tmpDir, zipfile, modpath, 0); err != nil {
 		fmt.Fprintf(os.Stderr, "-> %s\n", err)
-		return err
+		return "", err
 	}
 
-	if err := os.Rename(tmpDir, dir); err != nil {
-		return err
+	if err := robustio.Rename(tmpDir, dir); err != nil {
+		return "", err
 	}
 
 	// Make dir read-only only *after* renaming it.
 	// os.Rename was observed to fail for read-only directories on macOS.
 	makeDirsReadOnly(dir)
-	return nil
+	return dir, nil
 }
 
 var downloadZipCache par.Cache