sd_daemon/log: cheapen WriteString/WriteBytes.

2 files changed, 151 insertions, 11 deletions

diff --git a/sd_daemon/log.go b/sd_daemon/log.go
index df82d4b..806f248 100644
--- a/sd_daemon/log.go
+++ b/sd_daemon/log.go
@@ -17,7 +17,7 @@
 package sd_daemon
 
 import (
-	"fmt"
+	"bytes"
 	"io"
 	"log/syslog"
 	"os"
@@ -36,6 +36,7 @@ import (
 type Logger struct {
 	mu  sync.Mutex
 	out io.Writer
+	buf []byte
 }
 
 func NewLogger(w io.Writer) Logger {
@@ -50,25 +51,93 @@ func NewLogger(w io.Writer) Logger {
 // talk to syslog or journald directly.
 var Log = Logger{out: os.Stderr}
 
+// Cheap version of
+//
+//     *buf = append(*buf, fmt.Sprintf("<%d>", n)...)
+func appendPrefix(buf []byte, n syslog.Priority) []byte {
+	var b [22]byte // 21 = ceil(log_10(2^64))+len("<>")
+	b[len(b)-1] = '>'
+	i := len(b) - 2
+	for n >= 10 {
+		b[i] = byte('0' + n%10)
+		n = n / 10
+		i--
+	}
+	b[i] = byte('0' + n)
+	i--
+	b[i] = '<'
+	return append(buf, b[i:]...)
+}
+
 // WriteString writes a message with the specified priority to the
 // log.
-func (l Logger) WriteBytes(level syslog.Priority, msg []byte) (n int, err error) {
-	return l.WriteString(level, string(msg))
+func (l Logger) WriteString(level syslog.Priority, msg string) (n int, err error) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	msg = strings.TrimSuffix(msg, "\n")
+
+	// The following is a cheap version of:
+	//
+	//    prefix := fmt.Sprintf("<%d>", level)
+	//    buf := prefix + strings.Replace(msg, "\n", "\n"+prefix, -1)
+	//    return io.WriteString(l.out, buf)
+
+	l.buf = l.buf[:0]
+	l.buf = appendPrefix(l.buf, level) // possible allocation
+	prefix := l.buf
+	nlines := strings.Count(msg, "\n") + 1
+	n = len(msg) + len(prefix)*nlines + 1
+	if cap(l.buf) < n {
+		l.buf = make([]byte, len(l.buf), n) // allocation
+		copy(l.buf, prefix)
+	}
+
+	for nlines > 1 {
+		nl := strings.IndexByte(msg, '\n')
+		l.buf = append(l.buf, msg[:nl+1]...)
+		l.buf = append(l.buf, prefix...)
+		msg = msg[nl+1:]
+		nlines--
+	}
+	l.buf = append(l.buf, msg...)
+	l.buf = append(l.buf, '\n')
+
+	return l.out.Write(l.buf)
 }
 
 // WriteString writes a message with the specified priority to the
 // log.
-func (l Logger) WriteString(level syslog.Priority, msg string) (n int, err error) {
+func (l Logger) WriteBytes(level syslog.Priority, msg []byte) (n int, err error) {
+	// Copy/pasted from WriteString and
+	//  * `strings.` -> `bytes.`
+	//  * `"\n"` -> `[]byte{'\n'}`
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
-	// BUG(lukeshu): Logger uses high-level string functions that
-	// do many small allocations, making it insuitable for in
-	// tight loops; it should use a buffer property to be
-	// essentially zero-allocation.
-	prefix := fmt.Sprintf("<%d>", level)
-	buf := prefix + strings.Replace(strings.TrimSuffix(msg, "\n"), "\n", "\n"+prefix, -1)
-	return io.WriteString(l.out, buf)
+	msg = bytes.TrimSuffix(msg, []byte{'\n'})
+
+	l.buf = l.buf[:0]
+	l.buf = appendPrefix(l.buf, level) // possible allocation
+	prefix := l.buf
+	nlines := bytes.Count(msg, []byte{'\n'}) + 1
+	n = len(msg) + len(prefix)*nlines + 1
+	if cap(l.buf) < n {
+		l.buf = make([]byte, len(l.buf), n) // allocation
+		copy(l.buf, prefix)
+	}
+
+	for nlines > 1 {
+		nl := bytes.IndexByte(msg, '\n')
+		l.buf = append(l.buf, msg[:nl+1]...)
+		l.buf = append(l.buf, prefix...)
+		msg = msg[nl+1:]
+		nlines--
+	}
+	l.buf = append(l.buf, msg...)
+	l.buf = append(l.buf, '\n')
+
+	return l.out.Write(l.buf)
 }
 
 type loggerWriter struct {
diff --git a/sd_daemon/log_test.go b/sd_daemon/log_test.go
new file mode 100644
index 0000000..e2fdef9
--- /dev/null
+++ b/sd_daemon/log_test.go
@@ -0,0 +1,71 @@
+// Copyright 2016 Luke Shumaker
+//
+// 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 sd_daemon
+
+import (
+	"log/syslog"
+	"testing"
+)
+
+type cbWriter func(p []byte) (n int, err error)
+
+func (cb cbWriter) Write(p []byte) (n int, err error) {
+	return cb(p)
+}
+
+func TestLog(t *testing.T) {
+	var written []byte
+	log := NewLogger(cbWriter(func(p []byte) (n int, err error) {
+		written = p
+		return len(p), nil
+	}))
+
+	type testcase struct {
+		level syslog.Priority
+		msg   string
+	}
+
+	testcases := map[testcase]string{
+		testcase{6, "foo"}:             "<6>foo\n",
+		testcase{6, "foo\n"}:           "<6>foo\n",
+		testcase{6, "foo\nbar"}:        "<6>foo\n<6>bar\n",
+		testcase{6, "foo\nbar\n"}:      "<6>foo\n<6>bar\n",
+		testcase{6, "foo\nbar\nbaz"}:   "<6>foo\n<6>bar\n<6>baz\n",
+		testcase{6, "foo\nbar\nbaz\n"}: "<6>foo\n<6>bar\n<6>baz\n",
+
+		testcase{0, "foo"}:  "<0>foo\n",
+		testcase{1, "foo"}:  "<1>foo\n",
+		testcase{10, "foo"}: "<10>foo\n",
+	}
+
+	for in, out := range testcases {
+		written = nil
+		n, err := log.WriteString(in.level, in.msg)
+		if n != len(out) || string(written) != out || err != nil {
+			t.Errorf("WriteString(%#v, %#v)\n -> expected:{%#v, %#v, %#v}\n -> got:{%#v, %#v, %#v}\n",
+				in.level, in.msg,
+				len(out), nil, out,
+				n, err, string(written))
+		}
+		written = nil
+		n, err = log.WriteBytes(in.level, []byte(in.msg))
+		if n != len(out) || string(written) != out || err != nil {
+			t.Errorf("WriteBytes(%#v, %#v)\n -> expected:{%#v, %#v, %#v}\n -> got:{%#v, %#v, %#v}\n",
+				in.level, in.msg,
+				len(out), nil, out,
+				n, err, string(written))
+		}
+	}
+}