pkg/ioutil: configure pageOffset in PageWriter

etcd-io · Sep 28, 2016 · 2ccaa04 · 2ccaa04
1 parent 1e1dd24
commit 2ccaa04
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/pkg/ioutil/pagewriter.go b/pkg/ioutil/pagewriter.go
@@ -14,9 +14,7 @@
 
 package ioutil
 
-import (
-	"io"
-)
+import "io"
 
 var defaultBufferBytes = 128 * 1024
 
@@ -38,9 +36,10 @@ type PageWriter struct {
 	bufWatermarkBytes int
 }
 
-func NewPageWriter(w io.Writer, pageBytes int) *PageWriter {
+func NewPageWriter(w io.Writer, pageOffset, pageBytes int) *PageWriter {
 	return &PageWriter{
 		w:                 w,
+		pageOffset:        pageOffset,
 		pageBytes:         pageBytes,
 		buf:               make([]byte, defaultBufferBytes+pageBytes),
 		bufWatermarkBytes: defaultBufferBytes,

diff --git a/pkg/ioutil/pagewriter_test.go b/pkg/ioutil/pagewriter_test.go
@@ -25,7 +25,7 @@ func TestPageWriterRandom(t *testing.T) {
 	pageBytes := 128
 	buf := make([]byte, 4*defaultBufferBytes)
 	cw := &checkPageWriter{pageBytes: pageBytes, t: t}
-	w := NewPageWriter(cw, pageBytes)
+	w := NewPageWriter(cw, 0, pageBytes)
 	n := 0
 	for i := 0; i < 4096; i++ {
 		c, err := w.Write(buf[:rand.Intn(len(buf))])
@@ -51,7 +51,7 @@ func TestPageWriterPartialSlack(t *testing.T) {
 	pageBytes := 128
 	buf := make([]byte, defaultBufferBytes)
 	cw := &checkPageWriter{pageBytes: 64, t: t}
-	w := NewPageWriter(cw, pageBytes)
+	w := NewPageWriter(cw, 0, pageBytes)
 	// put writer in non-zero page offset
 	if _, err := w.Write(buf[:64]); err != nil {
 		t.Fatal(err)
@@ -82,6 +82,35 @@ func TestPageWriterPartialSlack(t *testing.T) {
 	}
 }
 
+// TestPageWriterOffset tests if page writer correctly repositions when offset is given.
+func TestPageWriterOffset(t *testing.T) {
+	defaultBufferBytes = 1024
+	pageBytes := 128
+	buf := make([]byte, defaultBufferBytes)
+	cw := &checkPageWriter{pageBytes: 64, t: t}
+	w := NewPageWriter(cw, 0, pageBytes)
+	if _, err := w.Write(buf[:64]); err != nil {
+		t.Fatal(err)
+	}
+	if err := w.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if w.pageOffset != 64 {
+		t.Fatalf("w.pageOffset expected 64, got %d", w.pageOffset)
+	}
+
+	w = NewPageWriter(cw, w.pageOffset, pageBytes)
+	if _, err := w.Write(buf[:64]); err != nil {
+		t.Fatal(err)
+	}
+	if err := w.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if w.pageOffset != 0 {
+		t.Fatalf("w.pageOffset expected 0, got %d", w.pageOffset)
+	}
+}
+
 // checkPageWriter implements an io.Writer that fails a test on unaligned writes.
 type checkPageWriter struct {
 	pageBytes  int