diff options
Diffstat (limited to 'pkg/linewriter')
-rw-r--r-- | pkg/linewriter/BUILD | 16 | ||||
-rw-r--r-- | pkg/linewriter/linewriter.go | 78 | ||||
-rw-r--r-- | pkg/linewriter/linewriter_test.go | 81 |
3 files changed, 175 insertions, 0 deletions
diff --git a/pkg/linewriter/BUILD b/pkg/linewriter/BUILD new file mode 100644 index 000000000..4a96c6f1d --- /dev/null +++ b/pkg/linewriter/BUILD @@ -0,0 +1,16 @@ +package(licenses = ["notice"]) # Apache 2.0 + +load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "linewriter", + srcs = ["linewriter.go"], + importpath = "gvisor.googlesource.com/gvisor/pkg/linewriter", + visibility = ["//visibility:public"], +) + +go_test( + name = "linewriter_test", + srcs = ["linewriter_test.go"], + embed = [":linewriter"], +) diff --git a/pkg/linewriter/linewriter.go b/pkg/linewriter/linewriter.go new file mode 100644 index 000000000..98f974410 --- /dev/null +++ b/pkg/linewriter/linewriter.go @@ -0,0 +1,78 @@ +// Copyright 2018 Google Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package linewriter provides an io.Writer which calls an emitter on each line. +package linewriter + +import ( + "bytes" + "sync" +) + +// Writer is an io.Writer which buffers input, flushing +// individual lines through an emitter function. +type Writer struct { + // the mutex locks buf. + sync.Mutex + + // buf holds the data we haven't emitted yet. + buf bytes.Buffer + + // emit is used to flush individual lines. + emit func(p []byte) +} + +// NewWriter creates a Writer which emits using emitter. +// The emitter must not retain p. It may change after emitter returns. +func NewWriter(emitter func(p []byte)) *Writer { + return &Writer{emit: emitter} +} + +// Write implements io.Writer.Write. +// It calls emit on each line of input, not including the newline. +// Write may be called concurrently. +func (w *Writer) Write(p []byte) (int, error) { + w.Lock() + defer w.Unlock() + + total := 0 + for len(p) > 0 { + emit := true + i := bytes.IndexByte(p, '\n') + if i < 0 { + // No newline, we will buffer everything. + i = len(p) + emit = false + } + + n, err := w.buf.Write(p[:i]) + if err != nil { + return total, err + } + total += n + + p = p[i:] + + if emit { + // Skip the newline, but still count it. + p = p[1:] + total++ + + w.emit(w.buf.Bytes()) + w.buf.Reset() + } + } + + return total, nil +} diff --git a/pkg/linewriter/linewriter_test.go b/pkg/linewriter/linewriter_test.go new file mode 100644 index 000000000..ce97cca05 --- /dev/null +++ b/pkg/linewriter/linewriter_test.go @@ -0,0 +1,81 @@ +// Copyright 2018 Google Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package linewriter + +import ( + "bytes" + "testing" +) + +func TestWriter(t *testing.T) { + testCases := []struct { + input []string + want []string + }{ + { + input: []string{"1\n", "2\n"}, + want: []string{"1", "2"}, + }, + { + input: []string{"1\n", "\n", "2\n"}, + want: []string{"1", "", "2"}, + }, + { + input: []string{"1\n2\n", "3\n"}, + want: []string{"1", "2", "3"}, + }, + { + input: []string{"1", "2\n"}, + want: []string{"12"}, + }, + { + // Data with no newline yet is omitted. + input: []string{"1\n", "2\n", "3"}, + want: []string{"1", "2"}, + }, + } + + for _, c := range testCases { + var lines [][]byte + + w := NewWriter(func(p []byte) { + // We must not retain p, so we must make a copy. + b := make([]byte, len(p)) + copy(b, p) + + lines = append(lines, b) + }) + + for _, in := range c.input { + n, err := w.Write([]byte(in)) + if err != nil { + t.Errorf("Write(%q) err got %v want nil (case %+v)", in, err, c) + } + if n != len(in) { + t.Errorf("Write(%q) b got %d want %d (case %+v)", in, n, len(in), c) + } + } + + if len(lines) != len(c.want) { + t.Errorf("len(lines) got %d want %d (case %+v)", len(lines), len(c.want), c) + } + + for i := range lines { + if !bytes.Equal(lines[i], []byte(c.want[i])) { + t.Errorf("item %d got %q want %q (case %+v)", i, lines[i], c.want[i], c) + } + } + } +} |