Skip upload copy if we know the length

If the source length is known, we can skip copying the file.
This commit is contained in:
James Bardin 2017-02-09 16:41:07 -05:00
parent 1448cb66fb
commit e1f06e5d0f
1 changed files with 51 additions and 29 deletions

View File

@ -292,8 +292,25 @@ func (c *Communicator) Upload(path string, input io.Reader) error {
// which works for unix and windows
targetDir = filepath.ToSlash(targetDir)
// Skip copying if we can get the file size directly from common io.Readers
size := int64(0)
switch src := input.(type) {
case *os.File:
fi, err := src.Stat()
if err != nil {
size = fi.Size()
}
case *bytes.Buffer:
size = int64(src.Len())
case *bytes.Reader:
size = int64(src.Len())
case *strings.Reader:
size = int64(src.Len())
}
scpFunc := func(w io.Writer, stdoutR *bufio.Reader) error {
return scpUploadFile(targetFile, input, w, stdoutR)
return scpUploadFile(targetFile, input, w, stdoutR, size)
}
return c.scpSession("scp -vt "+targetDir, scpFunc)
@ -490,7 +507,8 @@ func checkSCPStatus(r *bufio.Reader) error {
return nil
}
func scpUploadFile(dst string, src io.Reader, w io.Writer, r *bufio.Reader) error {
func scpUploadFile(dst string, src io.Reader, w io.Writer, r *bufio.Reader, size int64) error {
if size == 0 {
// Create a temporary file where we can copy the contents of the src
// so that we can determine the length, since SCP is length-prefixed.
tf, err := ioutil.TempFile("", "terraform-upload")
@ -521,14 +539,18 @@ func scpUploadFile(dst string, src io.Reader, w io.Writer, r *bufio.Reader) erro
return fmt.Errorf("Error creating temporary file for upload: %s", err)
}
src = tf
size = fi.Size()
}
// Start the protocol
log.Println("Beginning file upload...")
fmt.Fprintln(w, "C0644", fi.Size(), dst)
fmt.Fprintln(w, "C0644", size, dst)
if err := checkSCPStatus(r); err != nil {
return err
}
if _, err := io.Copy(w, tf); err != nil {
if _, err := io.Copy(w, src); err != nil {
return err
}
@ -592,7 +614,7 @@ func scpUploadDir(root string, fs []os.FileInfo, w io.Writer, r *bufio.Reader) e
err = func() error {
defer f.Close()
return scpUploadFile(fi.Name(), f, w, r)
return scpUploadFile(fi.Name(), f, w, r, fi.Size())
}()
if err != nil {