From cba1ebcbe4870a2184e60ec8649d45f35d0375cb Mon Sep 17 00:00:00 2001
From: Filippo Valsorda <hi@filippo.io>
Date: Sat, 2 Jan 2021 17:19:25 +0100
Subject: [PATCH] internal/format: buffer newlineWriter writes

Most writes in the cmd/age Writer stack are chunk-sized, so
approximately 64KiB. However, the newlineWriter, which splits lines at
64 columns, was doing a Write on the underlying Writer for each line,
making chunks effectively 48 bytes (before base64). There is no
buffering underneath it, so it was resulting in a lot of write syscalls.

Add a reusable bytes.Buffer to buffer the output of each
(*newlineWriter).Write call, and Write it all at once on the
destination.

This makes --armor just 50% slower than plain, instead of 10x.

Fixes #167
---
 cmd/age/age.go            |  3 ++-
 internal/format/format.go | 28 ++++++++++++++++------------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/cmd/age/age.go b/cmd/age/age.go
index 039cd21c..5b2ae816 100644
--- a/cmd/age/age.go
+++ b/cmd/age/age.go
@@ -154,7 +154,8 @@ func main() {
 		}
 	}
 
-	var in, out io.ReadWriter = os.Stdin, os.Stdout
+	var in io.Reader = os.Stdin
+	var out io.Writer = os.Stdout
 	if name := flag.Arg(0); name != "" && name != "-" {
 		f, err := os.Open(name)
 		if err != nil {
diff --git a/internal/format/format.go b/internal/format/format.go
index 07eba83d..d84ba2b1 100644
--- a/internal/format/format.go
+++ b/internal/format/format.go
@@ -55,29 +55,33 @@ func NewlineWriter(dst io.Writer) io.Writer {
 type newlineWriter struct {
 	dst     io.Writer
 	written int
+	buf     bytes.Buffer
 }
 
-func (w *newlineWriter) Write(p []byte) (n int, err error) {
+func (w *newlineWriter) Write(p []byte) (int, error) {
+	if w.buf.Len() != 0 {
+		panic("age: internal error: non-empty newlineWriter.buf")
+	}
 	for len(p) > 0 {
 		remainingInLine := ColumnsPerLine - (w.written % ColumnsPerLine)
 		if remainingInLine == ColumnsPerLine && w.written != 0 {
-			if _, err := w.dst.Write([]byte("\n")); err != nil {
-				return n, err
-			}
+			w.buf.Write([]byte("\n"))
 		}
 		toWrite := remainingInLine
 		if toWrite > len(p) {
 			toWrite = len(p)
 		}
-		nn, err := w.dst.Write(p[:toWrite])
-		n += nn
-		w.written += nn
-		p = p[nn:]
-		if err != nil {
-			return n, err
-		}
+		n, _ := w.buf.Write(p[:toWrite])
+		w.written += n
+		p = p[n:]
+	}
+	if _, err := w.buf.WriteTo(w.dst); err != nil {
+		// We always return n = 0 on error because it's hard to work back to the
+		// input length that ended up written out. Not ideal, but Write errors
+		// are not recoverable anyway.
+		return 0, err
 	}
-	return n, nil
+	return len(p), nil
 }
 
 const intro = "age-encryption.org/v1\n"