Only open output files when necessary

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

6 files changed, 260 insertions, 8 deletions

diff --git a/internal/iohelpers/readers.go b/internal/iohelpers/readers.go
new file mode 100644
index 00000000..277d7ab9
--- /dev/null
+++ b/internal/iohelpers/readers.go
@@ -0,0 +1,48 @@
+package iohelpers
+
+import (
+	"io"
+	"sync"
+)
+
+// LazyReadCloser provides an interface to a ReadCloser that will open on the
+// first access. The wrapped io.ReadCloser must be provided by 'open'.
+func LazyReadCloser(open func() (io.ReadCloser, error)) io.ReadCloser {
+	return &lazyReadCloser{
+		opened: sync.Once{},
+		open:   open,
+	}
+}
+
+type lazyReadCloser struct {
+	opened sync.Once
+	r      io.ReadCloser
+	// caches the error that came from open(), if any
+	openErr error
+	open    func() (io.ReadCloser, error)
+}
+
+var _ io.ReadCloser = (*lazyReadCloser)(nil)
+
+func (l *lazyReadCloser) openReader() (r io.ReadCloser, err error) {
+	l.opened.Do(func() {
+		l.r, l.openErr = l.open()
+	})
+	return l.r, l.openErr
+}
+
+func (l *lazyReadCloser) Close() error {
+	r, err := l.openReader()
+	if err != nil {
+		return err
+	}
+	return r.Close()
+}
+
+func (l *lazyReadCloser) Read(p []byte) (n int, err error) {
+	r, err := l.openReader()
+	if err != nil {
+		return 0, err
+	}
+	return r.Read(p)
+}
diff --git a/internal/iohelpers/readers_test.go b/internal/iohelpers/readers_test.go
new file mode 100644
index 00000000..ef300ef8
--- /dev/null
+++ b/internal/iohelpers/readers_test.go
@@ -0,0 +1,49 @@
+package iohelpers
+
+import (
+	"bytes"
+	"io"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLazyReadCloser(t *testing.T) {
+	r := newBufferCloser(bytes.NewBufferString("hello world"))
+	opened := false
+	l, ok := LazyReadCloser(func() (io.ReadCloser, error) {
+		opened = true
+		return r, nil
+	}).(*lazyReadCloser)
+	assert.True(t, ok)
+
+	assert.False(t, opened)
+	assert.Nil(t, l.r)
+	assert.False(t, r.closed)
+
+	p := make([]byte, 5)
+	n, err := l.Read(p)
+	assert.NoError(t, err)
+	assert.True(t, opened)
+	assert.Equal(t, r, l.r)
+	assert.Equal(t, 5, n)
+
+	err = l.Close()
+	assert.NoError(t, err)
+	assert.True(t, r.closed)
+
+	// test error propagation
+	l = LazyReadCloser(func() (io.ReadCloser, error) {
+		return nil, os.ErrNotExist
+	}).(*lazyReadCloser)
+
+	assert.Nil(t, l.r)
+
+	p = make([]byte, 5)
+	_, err = l.Read(p)
+	assert.Error(t, err)
+
+	err = l.Close()
+	assert.Error(t, err)
+}
diff --git a/internal/writers/writers.go b/internal/iohelpers/writers.go
index b40f4e73..0358bc43 100644
--- a/internal/writers/writers.go
+++ b/internal/iohelpers/writers.go
@@ -1,4 +1,4 @@
-package writers
+package iohelpers
 
 import (
 	"bufio"
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"sync"
 )
 
 type emptySkipper struct {
@@ -169,3 +170,45 @@ func (f *sameSkipper) Close() error {
 	}
 	return nil
 }
+
+// LazyWriteCloser provides an interface to a WriteCloser that will open on the
+// first access. The wrapped io.WriteCloser must be provided by 'open'.
+func LazyWriteCloser(open func() (io.WriteCloser, error)) io.WriteCloser {
+	return &lazyWriteCloser{
+		opened: sync.Once{},
+		open:   open,
+	}
+}
+
+type lazyWriteCloser struct {
+	opened sync.Once
+	w      io.WriteCloser
+	// caches the error that came from open(), if any
+	openErr error
+	open    func() (io.WriteCloser, error)
+}
+
+var _ io.WriteCloser = (*lazyWriteCloser)(nil)
+
+func (l *lazyWriteCloser) openWriter() (r io.WriteCloser, err error) {
+	l.opened.Do(func() {
+		l.w, l.openErr = l.open()
+	})
+	return l.w, l.openErr
+}
+
+func (l *lazyWriteCloser) Close() error {
+	w, err := l.openWriter()
+	if err != nil {
+		return err
+	}
+	return w.Close()
+}
+
+func (l *lazyWriteCloser) Write(p []byte) (n int, err error) {
+	w, err := l.openWriter()
+	if err != nil {
+		return 0, err
+	}
+	return w.Write(p)
+}
diff --git a/internal/writers/writers_test.go b/internal/iohelpers/writers_test.go
index b3452fcd..d226a03e 100644
--- a/internal/writers/writers_test.go
+++ b/internal/iohelpers/writers_test.go
@@ -1,9 +1,10 @@
-package writers
+package iohelpers
 
 import (
 	"bytes"
 	"fmt"
 	"io"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -37,7 +38,7 @@ func TestEmptySkipper(t *testing.T) {
 	}
 
 	for _, d := range testdata {
-		w := &bufferCloser{&bytes.Buffer{}}
+		w := newBufferCloser(&bytes.Buffer{})
 		opened := false
 		f, ok := NewEmptySkipper(func() (io.WriteCloser, error) {
 			opened = true
@@ -61,11 +62,18 @@ func TestEmptySkipper(t *testing.T) {
 	}
 }
 
+func newBufferCloser(b *bytes.Buffer) *bufferCloser {
+	return &bufferCloser{b, false}
+}
+
 type bufferCloser struct {
 	*bytes.Buffer
+
+	closed bool
 }
 
 func (b *bufferCloser) Close() error {
+	b.closed = true
 	return nil
 }
 
@@ -86,7 +94,7 @@ func TestSameSkipper(t *testing.T) {
 	for _, d := range testdata {
 		t.Run(fmt.Sprintf("in:%q/out:%q/same:%v", d.in, d.out, d.same), func(t *testing.T) {
 			r := bytes.NewBuffer(d.out)
-			w := &bufferCloser{&bytes.Buffer{}}
+			w := newBufferCloser(&bytes.Buffer{})
 			opened := false
 			f, ok := SameSkipper(r, func() (io.WriteCloser, error) {
 				opened = true
@@ -111,3 +119,41 @@ func TestSameSkipper(t *testing.T) {
 		})
 	}
 }
+
+func TestLazyWriteCloser(t *testing.T) {
+	w := newBufferCloser(&bytes.Buffer{})
+	opened := false
+	l, ok := LazyWriteCloser(func() (io.WriteCloser, error) {
+		opened = true
+		return w, nil
+	}).(*lazyWriteCloser)
+	assert.True(t, ok)
+
+	assert.False(t, opened)
+	assert.Nil(t, l.w)
+	assert.False(t, w.closed)
+
+	p := []byte("hello world")
+	n, err := l.Write(p)
+	assert.NoError(t, err)
+	assert.True(t, opened)
+	assert.Equal(t, 11, n)
+
+	err = l.Close()
+	assert.NoError(t, err)
+	assert.True(t, w.closed)
+
+	// test error propagation
+	l = LazyWriteCloser(func() (io.WriteCloser, error) {
+		return nil, os.ErrNotExist
+	}).(*lazyWriteCloser)
+
+	assert.Nil(t, l.w)
+
+	p = []byte("hello world")
+	_, err = l.Write(p)
+	assert.Error(t, err)
+
+	err = l.Close()
+	assert.Error(t, err)
+}
diff --git a/internal/tests/integration/inputdir_test.go b/internal/tests/integration/inputdir_test.go
index 3e3329da..ab791b48 100644
--- a/internal/tests/integration/inputdir_test.go
+++ b/internal/tests/integration/inputdir_test.go
@@ -43,10 +43,6 @@ out/{{ .in | strings.ReplaceAll $f (index .filemap $f) }}.out
 	)
 }
 
-func (s *InputDirSuite) TearDownTest(c *C) {
-	s.tmpDir.Remove()
-}
-
 func (s *InputDirSuite) TestInputDir(c *C) {
 	result := icmd.RunCommand(GomplateBin,
 		"--input-dir", s.tmpDir.Join("in"),
diff --git a/internal/tests/integration/inputdir_unix_test.go b/internal/tests/integration/inputdir_unix_test.go
new file mode 100644
index 00000000..5a6cde18
--- /dev/null
+++ b/internal/tests/integration/inputdir_unix_test.go
@@ -0,0 +1,70 @@
+//+build integration
+//+build !windows
+
+package integration
+
+import (
+	"fmt"
+	"io/ioutil"
+	"math"
+	"os"
+
+	. "gopkg.in/check.v1"
+
+	"golang.org/x/sys/unix"
+	"gotest.tools/v3/assert"
+	"gotest.tools/v3/fs"
+	"gotest.tools/v3/icmd"
+)
+
+func setFileUlimit(b uint64) error {
+	ulimit := unix.Rlimit{
+		Cur: b,
+		Max: math.MaxInt64,
+	}
+	err := unix.Setrlimit(unix.RLIMIT_NOFILE, &ulimit)
+	return err
+}
+
+func (s *InputDirSuite) TestInputDirRespectsUlimit(c *C) {
+	numfiles := 32
+	flist := map[string]string{}
+	for i := 0; i < numfiles; i++ {
+		k := fmt.Sprintf("file_%d", i)
+		flist[k] = fmt.Sprintf("hello world %d\n", i)
+	}
+	testdir := fs.NewDir(c, "ulimittestfiles",
+		fs.WithDir("in", fs.WithFiles(flist)),
+	)
+	defer testdir.Remove()
+
+	// we need another ~11 fds for other various things, so we'd be guaranteed
+	// to hit the limit if we try to have all the input files open
+	// simultaneously
+	setFileUlimit(uint64(numfiles))
+	defer setFileUlimit(8192)
+
+	result := icmd.RunCmd(icmd.Command(GomplateBin,
+		"--input-dir", testdir.Join("in"),
+		"--output-dir", testdir.Join("out"),
+	), func(c *icmd.Cmd) {
+		c.Dir = testdir.Path()
+	})
+	setFileUlimit(8192)
+	result.Assert(c, icmd.Success)
+
+	files, err := ioutil.ReadDir(testdir.Join("out"))
+	assert.NilError(c, err)
+	assert.Equal(c, numfiles, len(files))
+
+	for i := 0; i < numfiles; i++ {
+		f := testdir.Join("out", fmt.Sprintf("file_%d", i))
+		_, err := os.Stat(f)
+		assert.NilError(c, err)
+
+		content, err := ioutil.ReadFile(f)
+		assert.NilError(c, err)
+		expected := fmt.Sprintf("hello world %d\n", i)
+		assert.Equal(c, expected, string(content))
+	}
+}