Add helper.UnbufferedStartWriter

This type is intended to be used with PostUploadPack, which starts
slow with progress messages but then speeds up to where it can use
buffering.

2 files changed, 63 insertions, 0 deletions

diff --git a/internal/helper/unbuffered_start.go b/internal/helper/unbuffered_start.go
new file mode 100644
index 000000000..b82b1f8ef
--- /dev/null
+++ b/internal/helper/unbuffered_start.go
@@ -0,0 +1,32 @@
+package helper
+
+import (
+	"bufio"
+)
+
+// UnbufferedStartWriter wraps a *bufio.Writer so that early writes flush after each write.
+type UnbufferedStartWriter struct {
+	w         *bufio.Writer
+	slowStart int64
+	n         int64
+}
+
+// NewUnbufferedStartWriter returns a new UnbufferedStartWriter. Early writes
+// automatically call Flush on the underlying writer. Once the total
+// number of bytes written exceeds slowStart, the automatic flushing
+// stops. Callers should always call Flush when they are done.
+func NewUnbufferedStartWriter(w *bufio.Writer, slowStart int64) *UnbufferedStartWriter {
+	return &UnbufferedStartWriter{w: w, slowStart: slowStart}
+}
+
+func (ssw *UnbufferedStartWriter) Write(p []byte) (int, error) {
+	n, err := ssw.w.Write(p)
+	ssw.n += int64(n)
+	if err == nil && ssw.n <= ssw.slowStart {
+		err = ssw.Flush()
+	}
+	return n, err
+}
+
+// Flush flushes the underlying bufio.Writer of ssw.
+func (ssw *UnbufferedStartWriter) Flush() error { return ssw.w.Flush() }
diff --git a/internal/helper/unbuffered_start_test.go b/internal/helper/unbuffered_start_test.go
new file mode 100644
index 000000000..70194c7e2
--- /dev/null
+++ b/internal/helper/unbuffered_start_test.go
@@ -0,0 +1,31 @@
+package helper
+
+import (
+	"bufio"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type writeLogger struct {
+	writes []string
+}
+
+func (wl *writeLogger) Write(p []byte) (int, error) {
+	wl.writes = append(wl.writes, string(p))
+	return len(p), nil
+}
+
+func TestUnbufferedStartWriter(t *testing.T) {
+	wl := &writeLogger{}
+	ssw := NewUnbufferedStartWriter(bufio.NewWriter(wl), 5)
+
+	for _, c := range []byte("hello world") {
+		n, err := ssw.Write([]byte{c})
+		require.NoError(t, err)
+		require.Equal(t, 1, n)
+	}
+	require.NoError(t, ssw.Flush())
+
+	require.Equal(t, []string{"h", "e", "l", "l", "o", " world"}, wl.writes)
+}