diff --git a/pkg/ioutil/pagewriter.go b/pkg/ioutil/pagewriter.go index ed22d942011a..5f0b762566da 100644 --- 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 index 62cbd649235a..88bbf43b338e 100644 --- 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