safeio: Extend WriteFile to take operations before the final rename

This patch extends WriteFile to allow arbitrary operations to be applied
to the file before it is atomically renamed.

This will be used in upcoming patches to change the mtime of the file
before it is atomically renamed.

internal/safeio/safeio_test.go

@@ -2,9 +2,11 @@ package safeio
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"os"
+	"strings"
 	"testing"
 )
 
@@ -24,8 +26,8 @@ func mustTempDir(t *testing.T) string {
 	return dir
 }
 
-func testWriteFile(fname string, data []byte, perm os.FileMode) error {
-	err := WriteFile("file1", data, perm)
+func testWriteFile(fname string, data []byte, perm os.FileMode, ops ...FileOp) error {
+	err := WriteFile("file1", data, perm, ops...)
 	if err != nil {
 		return fmt.Errorf("error writing new file: %v", err)
 	}
@@ -81,6 +83,65 @@ func TestWriteFile(t *testing.T) {
 	}
 }
 
+func TestWriteFileWithOp(t *testing.T) {
+	dir := mustTempDir(t)
+
+	var opFile string
+	op := func(f string) error {
+		opFile = f
+		return nil
+	}
+
+	content := []byte("content 1")
+	if err := testWriteFile("file1", content, 0660, op); err != nil {
+		t.Error(err)
+	}
+
+	if opFile == "" {
+		t.Error("operation was not called")
+	}
+	if !strings.Contains(opFile, "file1") {
+		t.Errorf("operation called with suspicious file: %s", opFile)
+	}
+
+	// Remove the test directory, but only if we have not failed. We want to
+	// keep the failed structure for debugging.
+	if !t.Failed() {
+		os.RemoveAll(dir)
+	}
+}
+
+func TestWriteFileWithFailingOp(t *testing.T) {
+	dir := mustTempDir(t)
+
+	var opFile string
+	opOK := func(f string) error {
+		opFile = f
+		return nil
+	}
+
+	opError := errors.New("operation failed")
+	opFail := func(f string) error {
+		return opError
+	}
+
+	content := []byte("content 1")
+	err := WriteFile("file1", content, 0660, opOK, opOK, opFail)
+	if err != opError {
+		t.Errorf("different error, got %v, expected %v", err, opError)
+	}
+
+	if _, err := os.Stat(opFile); err == nil {
+		t.Errorf("temporary file was not removed after failure (%v)", opFile)
+	}
+
+	// Remove the test directory, but only if we have not failed. We want to
+	// keep the failed structure for debugging.
+	if !t.Failed() {
+		os.RemoveAll(dir)
+	}
+}
+
 // TODO: We should test the possible failure scenarios for WriteFile, but it
 // gets tricky without being able to do failure injection (or turning the code
 // into a mess).