Merge pull request #8193 from ethereum-optimism/aj/atomic-writer

cannon: Use a unique temp file when writing JSON content.

cannon/cmd/json.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"path"
 
 	"github.com/ethereum-optimism/optimism/op-service/ioutil"
 )
@@ -35,18 +34,16 @@ func writeJSON[X any](outputPath string, value X) error {
 	var out io.Writer
 	finish := func() error { return nil }
 	if outputPath != "-" {
-		// Write to a tmp file but reserve the file extension if present
-		tmpPath := outputPath + "-tmp" + path.Ext(outputPath)
-		f, err := ioutil.OpenCompressed(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
+		f, err := ioutil.NewAtomicWriterCompressed(outputPath, 0755)
 		if err != nil {
 			return fmt.Errorf("failed to open output file: %w", err)
 		}
+		// Ensure we close the stream even if failures occur.
 		defer f.Close()
 		out = f
-		finish = func() error {
-			// Rename the file into place as atomically as the OS will allow
-			return os.Rename(tmpPath, outputPath)
-		}
+		// Closing the file causes it to be renamed to the final destination
+		// so make sure we handle any errors it returns
+		finish = f.Close
 	} else {
 		out = os.Stdout
 	}

op-service/ioutil/atomic.go

+package ioutil
+
+import (
+	"io"
+	"os"
+	"path/filepath"
+)
+
+type atomicWriter struct {
+	dest string
+	temp string
+	out  io.WriteCloser
+}
+
+// NewAtomicWriterCompressed creates a io.WriteCloser that performs an atomic write.
+// The contents are initially written to a temporary file and only renamed into place when the writer is closed.
+// NOTE: It's vital to check if an error is returned from Close() as it may indicate the file could not be renamed
+// If path ends in .gz the contents written will be gzipped.
+func NewAtomicWriterCompressed(path string, perm os.FileMode) (io.WriteCloser, error) {
+	f, err := os.CreateTemp(filepath.Dir(path), filepath.Base(path))
+	if err != nil {
+		return nil, err
+	}
+	if err := f.Chmod(perm); err != nil {
+		_ = f.Close()
+		return nil, err
+	}
+	return &atomicWriter{
+		dest: path,
+		temp: f.Name(),
+		out:  CompressByFileType(path, f),
+	}, nil
+}
+
+func (a *atomicWriter) Write(p []byte) (n int, err error) {
+	return a.out.Write(p)
+}
+
+func (a *atomicWriter) Close() error {
+	// Attempt to clean up the temp file even if it can't be renamed into place.
+	defer os.Remove(a.temp)
+	if err := a.out.Close(); err != nil {
+		return err
+	}
+	return os.Rename(a.temp, a.dest)
+}

op-service/ioutil/atomic_test.go

+package ioutil
+
+import (
+	"io"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestAtomicWriter_RenameOnClose(t *testing.T) {
+	dir := t.TempDir()
+	target := filepath.Join(dir, "target.txt")
+	f, err := NewAtomicWriterCompressed(target, 0755)
+	require.NoError(t, err)
+	defer f.Close()
+	_, err = os.Stat(target)
+	require.ErrorIs(t, err, os.ErrNotExist, "should not create target file when created")
+
+	content := ([]byte)("Hello world")
+	n, err := f.Write(content)
+	require.NoError(t, err)
+	require.Equal(t, len(content), n)
+	_, err = os.Stat(target)
+	require.ErrorIs(t, err, os.ErrNotExist, "should not create target file when writing")
+
+	require.NoError(t, f.Close())
+	stat, err := os.Stat(target)
+	require.NoError(t, err, "should create target file when closed")
+	require.EqualValues(t, fs.FileMode(0755), stat.Mode())
+
+	files, err := os.ReadDir(dir)
+	require.NoError(t, err)
+	require.Len(t, files, 1, "should not leave temporary files behind")
+}
+
+func TestAtomicWriter_MultipleClose(t *testing.T) {
+	dir := t.TempDir()
+	target := filepath.Join(dir, "target.txt")
+	f, err := NewAtomicWriterCompressed(target, 0755)
+	require.NoError(t, err)
+
+	require.NoError(t, f.Close())
+	require.ErrorIs(t, f.Close(), os.ErrClosed)
+}
+
+func TestAtomicWriter_ApplyGzip(t *testing.T) {
+	tests := []struct {
+		name       string
+		filename   string
+		compressed bool
+	}{
+		{"Uncompressed", "test.notgz", false},
+		{"Gzipped", "test.gz", true},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			data := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 0, 0, 0, 0, 0, 0}
+			dir := t.TempDir()
+			path := filepath.Join(dir, test.filename)
+			out, err := NewAtomicWriterCompressed(path, 0o644)
+			require.NoError(t, err)
+			defer out.Close()
+			_, err = out.Write(data)
+			require.NoError(t, err)
+			require.NoError(t, out.Close())
+
+			writtenData, err := os.ReadFile(path)
+			require.NoError(t, err)
+			if test.compressed {
+				require.NotEqual(t, data, writtenData, "should have compressed data on disk")
+			} else {
+				require.Equal(t, data, writtenData, "should not have compressed data on disk")
+			}
+
+			in, err := OpenDecompressed(path)
+			require.NoError(t, err)
+			readData, err := io.ReadAll(in)
+			require.NoError(t, err)
+			require.Equal(t, data, readData)
+		})
+	}
+}

op-service/ioutil/gzip.go

@@ -33,10 +33,7 @@ func OpenCompressed(file string, flag int, perm os.FileMode) (io.WriteCloser, er
 	if err != nil {
 		return nil, err
 	}
-	if IsGzip(file) {
-		out = gzip.NewWriter(out)
-	}
-	return out, nil
+	return CompressByFileType(file, out), nil
 }
 
 // WriteCompressedJson writes the object to the specified file as a compressed json object
@@ -58,3 +55,10 @@ func WriteCompressedJson(file string, obj any) error {
 func IsGzip(path string) bool {
 	return strings.HasSuffix(path, ".gz")
 }
+
+func CompressByFileType(file string, out io.WriteCloser) io.WriteCloser {
+	if IsGzip(file) {
+		return gzip.NewWriter(out)
+	}
+	return out
+}