resilien
/
terraform
2
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

80
communicator/ssh/communicator.go
View File

@ -292,8 +292,25 @@ func (c *Communicator) Upload(path string, input io.Reader) error {
// which works for unix and windows // which works for unix and windows
targetDir = filepath.ToSlash(targetDir) 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 { 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) return c.scpSession("scp -vt "+targetDir, scpFunc)
@ -490,45 +507,50 @@ func checkSCPStatus(r *bufio.Reader) error {
return nil 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 {
// Create a temporary file where we can copy the contents of the src if size == 0 {
// so that we can determine the length, since SCP is length-prefixed. // Create a temporary file where we can copy the contents of the src
tf, err := ioutil.TempFile("", "terraform-upload") // so that we can determine the length, since SCP is length-prefixed.
if err != nil { tf, err := ioutil.TempFile("", "terraform-upload")
return fmt.Errorf("Error creating temporary file for upload: %s", err) if err != nil {
} return fmt.Errorf("Error creating temporary file for upload: %s", err)
defer os.Remove(tf.Name()) }
defer tf.Close() defer os.Remove(tf.Name())
defer tf.Close()
log.Println("Copying input data into temporary file so we can read the length") log.Println("Copying input data into temporary file so we can read the length")
if _, err := io.Copy(tf, src); err != nil { if _, err := io.Copy(tf, src); err != nil {
return err return err
} }
// Sync the file so that the contents are definitely on disk, then // Sync the file so that the contents are definitely on disk, then
// read the length of it. // read the length of it.
if err := tf.Sync(); err != nil { if err := tf.Sync(); err != nil {
return fmt.Errorf("Error creating temporary file for upload: %s", err) return fmt.Errorf("Error creating temporary file for upload: %s", err)
} }
// Seek the file to the beginning so we can re-read all of it // Seek the file to the beginning so we can re-read all of it
if _, err := tf.Seek(0, 0); err != nil { if _, err := tf.Seek(0, 0); err != nil {
return fmt.Errorf("Error creating temporary file for upload: %s", err) return fmt.Errorf("Error creating temporary file for upload: %s", err)
} }
fi, err := tf.Stat() fi, err := tf.Stat()
if err != nil { if err != nil {
return fmt.Errorf("Error creating temporary file for upload: %s", err) return fmt.Errorf("Error creating temporary file for upload: %s", err)
}
src = tf
size = fi.Size()
} }
// Start the protocol // Start the protocol
log.Println("Beginning file upload...") log.Println("Beginning file upload...")
fmt.Fprintln(w, "C0644", fi.Size(), dst) fmt.Fprintln(w, "C0644", size, dst)
if err := checkSCPStatus(r); err != nil { if err := checkSCPStatus(r); err != nil {
return err return err
} }
if _, err := io.Copy(w, tf); err != nil { if _, err := io.Copy(w, src); err != nil {
return err return err
} }
@ -592,7 +614,7 @@ func scpUploadDir(root string, fs []os.FileInfo, w io.Writer, r *bufio.Reader) e
err = func() error { err = func() error {
defer f.Close() defer f.Close()
return scpUploadFile(fi.Name(), f, w, r) return scpUploadFile(fi.Name(), f, w, r, fi.Size())
}() }()
if err != nil { if err != nil {