webdav: make memFile.Write allocate less often.

benchmark                 old ns/op     new ns/op     delta
BenchmarkMemFileWrite     8498028       625563        -92.64%

Change-Id: Iec7dd3931c891c9d6d9d5c6ccd05300b031a5f86

2 files changed, 84 insertions, 10 deletions

diff --git a/webdav/file.go b/webdav/file.go
index 34c8df8..4a0c7ff 100644
--- a/webdav/file.go
+++ b/webdav/file.go
@@ -404,16 +404,36 @@ func (f *memFile) Stat() (os.FileInfo, error) {
 }
 
 func (f *memFile) Write(p []byte) (int, error) {
+	lenp := len(p)
 	f.n.mu.Lock()
 	defer f.n.mu.Unlock()
-	epos := len(p) + f.pos
-	if epos > len(f.n.data) {
-		d := make([]byte, epos)
-		copy(d, f.n.data)
-		f.n.data = d
+
+	if f.pos < len(f.n.data) {
+		n := copy(f.n.data[f.pos:], p)
+		f.pos += n
+		p = p[n:]
+	} else if f.pos > len(f.n.data) {
+		// Write permits the creation of holes, if we've seek'ed past the
+		// existing end of file.
+		if f.pos <= cap(f.n.data) {
+			oldLen := len(f.n.data)
+			f.n.data = f.n.data[:f.pos]
+			hole := f.n.data[oldLen:]
+			for i := range hole {
+				hole[i] = 0
+			}
+		} else {
+			d := make([]byte, f.pos, f.pos+len(p))
+			copy(d, f.n.data)
+			f.n.data = d
+		}
+	}
+
+	if len(p) > 0 {
+		// We should only get here if f.pos == len(f.n.data).
+		f.n.data = append(f.n.data, p...)
+		f.pos = len(f.n.data)
 	}
-	copy(f.n.data[f.pos:], p)
-	f.pos = epos
 	f.n.modTime = time.Now()
-	return len(p), nil
+	return lenp, nil
 }
diff --git a/webdav/file_test.go b/webdav/file_test.go
index ff68c4a..d95d5a8 100644
--- a/webdav/file_test.go
+++ b/webdav/file_test.go
@@ -252,9 +252,12 @@ func TestMemFile(t *testing.T) {
 		"write C",
 		"wantData abcdEFghijkyyyzzstsZZAAAABBBBB..........C",
 		"wantSize 41",
-		"seek set 43 want 43",
 		"write D",
-		"wantData abcdEFghijkyyyzzstsZZAAAABBBBB..........C..D",
+		"wantData abcdEFghijkyyyzzstsZZAAAABBBBB..........CD",
+		"wantSize 42",
+		"seek set 43 want 43",
+		"write E",
+		"wantData abcdEFghijkyyyzzstsZZAAAABBBBB..........CD.E",
 		"wantSize 44",
 		"seek set 0 want 0",
 		"write 5*123456789_",
@@ -396,3 +399,54 @@ func TestMemFile(t *testing.T) {
 		}
 	}
 }
+
+// TestMemFileWriteAllocs tests that writing N consecutive 1KiB chunks to a
+// memFile doesn't allocate a new buffer for each of those N times. Otherwise,
+// calling io.Copy(aMemFile, src) is likely to have quadratic complexity.
+func TestMemFileWriteAllocs(t *testing.T) {
+	fs := NewMemFS()
+	f, err := fs.OpenFile("/xxx", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
+	if err != nil {
+		t.Fatalf("OpenFile: %v", err)
+	}
+	defer f.Close()
+
+	xxx := make([]byte, 1024)
+	for i := range xxx {
+		xxx[i] = 'x'
+	}
+
+	a := testing.AllocsPerRun(100, func() {
+		f.Write(xxx)
+	})
+	// AllocsPerRun returns an integral value, so we compare the rounded-down
+	// number to zero.
+	if a > 0 {
+		t.Fatalf("%v allocs per run, want 0", a)
+	}
+}
+
+func BenchmarkMemFileWrite(b *testing.B) {
+	fs := NewMemFS()
+	xxx := make([]byte, 1024)
+	for i := range xxx {
+		xxx[i] = 'x'
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		f, err := fs.OpenFile("/xxx", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
+		if err != nil {
+			b.Fatalf("OpenFile: %v", err)
+		}
+		for j := 0; j < 100; j++ {
+			f.Write(xxx)
+		}
+		if err := f.Close(); err != nil {
+			b.Fatalf("Close: %v", err)
+		}
+		if err := fs.RemoveAll("/xxx"); err != nil {
+			b.Fatalf("RemoveAll: %v", err)
+		}
+	}
+}