Add Writer struct

This lets you Write() directly and get timestamp annotations.

lib/tss.go

@@ -1,28 +1,100 @@
 package tss
 
 import (
-	"bufio"
 	"bytes"
-	"fmt"
 	"io"
-	"strings"
+	"strconv"
 	"time"
 )
 
+type Writer struct {
+	w         io.Writer
+	start     time.Time
+	lastLine  time.Time
+	buf       bytes.Buffer
+	endOfLine bool
+}
+
+func NewWriter(w io.Writer, start time.Time) *Writer {
+	if start.IsZero() {
+		start = time.Now()
+	}
+	return &Writer{w: w, start: start, endOfLine: true}
+}
+
+var padding = bytes.Repeat([]byte{' '}, 9)
+
+// Write writes the contents of p into the buffer. It returns the number of
+// bytes written. If nn < len(p), it also returns an error explaining why the
+// write is short.
+func (w *Writer) Write(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+	wrote := 0
+	now := time.Now()
+	pos := 0
+	var sbuf [9]byte
+	for {
+		// write everything up to the next newline
+		if w.endOfLine {
+			// print timing info
+			sinceStart := now.Sub(w.start).Round(100 * time.Microsecond)
+			s := TimeScaler(sinceStart)
+			for i := 0; i < 8-len(s); i++ {
+				w.buf.WriteByte(' ')
+				sbuf[i] = ' '
+			}
+			w.buf.WriteString(s)
+			w.buf.WriteByte(' ')
+			if w.lastLine.IsZero() {
+				w.buf.Write(padding)
+				w.lastLine = now
+			} else {
+				sinceLastLine := now.Sub(w.lastLine).Round(100 * time.Microsecond)
+				s := TimeScaler(sinceLastLine)
+				for i := 0; i < 8-len(s); i++ {
+					w.buf.WriteByte(' ')
+				}
+				w.buf.WriteString(s)
+				w.buf.WriteByte(' ')
+			}
+			w.endOfLine = false
+		}
+		idx := bytes.IndexByte(p[pos:], '\n')
+		if idx >= 0 {
+			w.buf.Write(p[pos : pos+idx+1])
+			wrote += idx + 1
+			pos += idx + 1
+			w.endOfLine = true
+			if pos >= len(p) {
+				break
+			}
+		} else {
+			w.buf.Write(p[pos:])
+			wrote += len(p) - pos
+			break
+		}
+	}
+	_, err := w.w.Write(w.buf.Bytes())
+	w.buf.Reset()
+	return wrote, err
+}
+
 // TimeScaler returns a format string for the given Duration where all of the
 // decimals will line up in the same column (fourth from the end).
 func TimeScaler(d time.Duration) string {
 	switch {
 	case d == 0:
 		return "0.0ms"
 	case d >= time.Second:
-		return fmt.Sprintf("%.2fs", float64(d.Nanoseconds())/1e9)
+		return strconv.FormatFloat(float64(d.Nanoseconds())/1e9, 'f', 2, 64) + "s"
 	case d >= 50*time.Microsecond:
-		return fmt.Sprintf("%.1fms", float64(d.Nanoseconds())/1e9*1000)
+		return strconv.FormatFloat(float64(d.Nanoseconds())/1e9*1000, 'f', 1, 64) + "ms"
 	case d >= time.Microsecond:
-		return fmt.Sprintf("%.1fµs", float64(d.Nanoseconds())/1e9*1000*1000)
+		return strconv.FormatFloat(float64(d.Nanoseconds())/1e9*1000*1000, 'f', 1, 64) + "µs"
 	default:
-		return fmt.Sprintf("%.1fns", float64(d.Nanoseconds()))
+		return strconv.FormatFloat(float64(d.Nanoseconds()), 'f', 1, 64) + "ns"
 	}
 }
 
@@ -31,29 +103,5 @@ func Copy(w io.Writer, r io.Reader) (written int64, err error) {
 }
 
 func CopyTime(w io.Writer, r io.Reader, start time.Time) (written int64, err error) {
-	bs := bufio.NewScanner(r)
-	n := int64(0)
-	var lastLine time.Time
-	var buf bytes.Buffer
-	for bs.Scan() {
-		gotLine := time.Now()
-		sinceLastLine := gotLine.Sub(lastLine).Round(100 * time.Microsecond)
-		sinceStart := gotLine.Sub(start).Round(100 * time.Microsecond)
-		fmt.Fprintf(&buf, "%8s ", TimeScaler(sinceStart))
-		if lastLine.IsZero() {
-			buf.WriteString(strings.Repeat(" ", 9))
-		} else {
-			fmt.Fprintf(&buf, "%8s ", TimeScaler(sinceLastLine))
-		}
-		buf.Write(bs.Bytes())
-		buf.WriteByte('\n')
-		wn, err := w.Write(buf.Bytes())
-		n += int64(wn)
-		if err != nil {
-			return n, err
-		}
-		buf.Reset()
-		lastLine = gotLine
-	}
-	return n, bs.Err()
+	return io.Copy(NewWriter(w, start), r)
 }

lib/tss_test.go

@@ -26,24 +26,57 @@ func (s *sleepReader) Read(p []byte) (int, error) {
 		return 6, nil
 	}
 	time.Sleep(s.sleepFor)
+	if s.count == 2 {
+		copy(p[:3], "hel")
+		return 3, nil
+	}
+	if s.count == 3 {
+		copy(p[:3], "lo\n")
+		return 3, nil
+	}
+	if s.count == 4 {
+		copy(p[:15], "hello\nhello\nhel")
+		return 15, nil
+	}
+	if s.count == 5 {
+		copy(p[:3], "lo\n")
+		return 3, nil
+	}
 	copy(p[:6], "hello\n")
 	return 6, nil
 }
 
+func TestWriter(t *testing.T) {
+	t.Parallel()
+	max := 6
+	s := &sleepReader{max: max, sleepFor: 2 * time.Millisecond}
+	buf := new(bytes.Buffer)
+	w := tss.NewWriter(buf, time.Time{})
+	n, err := io.Copy(w, s)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if int(n) != len("hello\n")*max {
+		t.Errorf("expected n of 36, got %d:\n%s", n, buf)
+	}
+}
+
 func TestCopy(t *testing.T) {
 	t.Parallel()
-	s := &sleepReader{max: 3, sleepFor: 5 * time.Millisecond}
+	max := 6
+	s := &sleepReader{max: max, sleepFor: 2 * time.Millisecond}
 	buf := new(bytes.Buffer)
 	n, err := tss.Copy(buf, s)
-	if n != 72 {
-		t.Errorf("expected n of 72, got %d", n)
-	}
+	want := len("hello\n") * 6
 	if err != nil {
 		t.Errorf("expected nil error, got %v", err)
 	}
+	if int(n) != want {
+		t.Errorf("expected n of %d, got %d", want, n)
+	}
 	parts := strings.Split(buf.String(), "\n")
-	if len(parts) != 4 {
-		t.Errorf("incorrect number of parts: want 4 got %q", parts)
+	if len(parts) != 7 {
+		t.Errorf("incorrect number of parts: want 6 got %d:\n%q", len(parts), parts)
 	}
 	line1 := parts[0]
 	if len(line1) != 23 {
@@ -91,17 +124,38 @@ func TestTimeScaler(t *testing.T) {
 }
 
 func BenchmarkCopy(b *testing.B) {
-	bs := bytes.Repeat([]byte{'a'}, 512+1)
-	for i := 0; i < len(bs); i += 40 {
+	bs := bytes.Repeat([]byte{'a'}, 2<<12+1)
+	for i := 0; i < len(bs); i += 50 {
+		bs[i] = '\n'
+	}
+	b.SetBytes(int64(len(bs)))
+	rd := bytes.NewReader(bs)
+	buf := new(bytes.Buffer)
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		tss.CopyTime(buf, rd, time.Now().Add(-50*time.Millisecond))
+		buf.Reset()
+		rd.Reset(bs)
+	}
+}
+
+func BenchmarkWriter(b *testing.B) {
+	bs := bytes.Repeat([]byte{'a'}, 2<<12)
+	for i := 0; i < len(bs); i += 50 {
 		bs[i] = '\n'
 	}
+	b.SetBytes(int64(len(bs)))
 	rd := bytes.NewReader(bs)
 	buf := new(bytes.Buffer)
+	w := tss.NewWriter(buf, time.Now().Add(-50*time.Millisecond))
 	b.ReportAllocs()
 	b.ResetTimer()
 
 	for i := 0; i < b.N; i++ {
-		tss.Copy(buf, rd)
+		io.Copy(w, rd)
+		buf.Reset()
 		rd.Reset(bs)
 	}
 }