Merge branch 'jc-lock-file' into 'master'

Add safe package for safe file writes

See merge request gitlab-org/gitaly!1298

4 files changed, 246 insertions, 0 deletions

diff --git a/changelogs/unreleased/jc-lock-file.yml b/changelogs/unreleased/jc-lock-file.yml
new file mode 100644
index 000000000..a6cdc2877
--- /dev/null
+++ b/changelogs/unreleased/jc-lock-file.yml
@@ -0,0 +1,5 @@
+---
+title: Add lock file package for atomic file writes
+merge_request: 1298
+author:
+type: other
diff --git a/internal/safe/file_writer.go b/internal/safe/file_writer.go
new file mode 100644
index 000000000..62f6860ee
--- /dev/null
+++ b/internal/safe/file_writer.go
@@ -0,0 +1,118 @@
+package safe
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"sync"
+
+	"github.com/sirupsen/logrus"
+)
+
+// FileWriter is a thread safe writer that does an atomic write to the target file. It allows one
+// writer at a time to acquire a lock, write the file, and atomically replace the contents of the target file.
+type FileWriter struct {
+	tmpFile   *os.File
+	path      string
+	closeErr  error
+	closeOnce sync.Once
+}
+
+// CreateFileWriter takes path as an absolute path of the target file and creates a new FileWriter by attempting to create a tempfile
+func CreateFileWriter(ctx context.Context, path string) (*FileWriter, error) {
+	if ctx.Done() == nil {
+		return nil, errors.New("context cannot be cancelled")
+	}
+	var err error
+	writer := &FileWriter{path: path}
+
+	directory := filepath.Dir(path)
+
+	tmpFile, err := ioutil.TempFile(directory, filepath.Base(path))
+	if err != nil {
+		return nil, err
+	}
+
+	writer.tmpFile = tmpFile
+
+	go writer.cleanupOnContextDone(ctx)
+
+	return writer, nil
+}
+
+func (fw *FileWriter) cleanupOnContextDone(ctx context.Context) {
+	<-ctx.Done()
+	if err := fw.cleanup(); err != nil {
+		logrus.WithField("path", fw.path).WithError(err).Error("error when closing FileWriter")
+	}
+}
+
+// Write wraps the temporary file's Write.
+func (fw *FileWriter) Write(p []byte) (n int, err error) {
+	return fw.tmpFile.Write(p)
+}
+
+// Commit will close the temporary file and rename it to the target file name the first call to Commit() will close and
+// delete the temporary file, so subsequenty calls to Commit() are gauaranteed to return an error.
+func (fw *FileWriter) Commit() error {
+	if err := fw.tmpFile.Sync(); err != nil {
+		return fmt.Errorf("syncing temp file: %v", err)
+	}
+
+	if err := fw.close(); err != nil {
+		return fmt.Errorf("closing temp file: %v", err)
+	}
+
+	if err := fw.rename(); err != nil {
+		return fmt.Errorf("renaming temp file: %v", err)
+	}
+
+	if err := fw.syncDir(); err != nil {
+		return fmt.Errorf("syncing dir: %v", err)
+	}
+
+	return nil
+}
+
+// rename renames the temporary file to the target file
+func (fw *FileWriter) rename() error {
+	return os.Rename(fw.tmpFile.Name(), fw.path)
+}
+
+// syncDir will sync the directory
+func (fw *FileWriter) syncDir() error {
+	f, err := os.Open(filepath.Dir(fw.path))
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	return f.Sync()
+}
+
+// cleanup will close the temporary file and remove it.
+func (fw *FileWriter) cleanup() error {
+	var err error
+
+	if err = fw.close(); err != nil {
+		return err
+	}
+
+	if err = os.Remove(fw.tmpFile.Name()); err != nil && !os.IsNotExist(err) {
+		return err
+	}
+
+	return nil
+}
+
+// close uses sync.Once to guarantee that the file gets closed only once
+func (fw *FileWriter) close() error {
+	fw.closeOnce.Do(func() {
+		fw.closeErr = fw.tmpFile.Close()
+	})
+
+	return fw.closeErr
+}
diff --git a/internal/safe/file_writer_test.go b/internal/safe/file_writer_test.go
new file mode 100644
index 000000000..a7e281be7
--- /dev/null
+++ b/internal/safe/file_writer_test.go
@@ -0,0 +1,114 @@
+package safe_test
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitaly/internal/safe"
+	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
+)
+
+func TestFile(t *testing.T) {
+	dir, cleanup := testhelper.TempDir(t, "", t.Name())
+	defer cleanup()
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	filePath := filepath.Join(dir, "test_file_contents")
+	fileContents := "very important contents"
+	file, err := safe.CreateFileWriter(ctx, filePath)
+	require.NoError(t, err)
+
+	_, err = io.Copy(file, bytes.NewBufferString(fileContents))
+	require.NoError(t, err)
+
+	testhelper.AssertFileNotExists(t, filePath)
+
+	require.NoError(t, file.Commit())
+
+	writtenContents, err := ioutil.ReadFile(filePath)
+	require.NoError(t, err)
+	require.Equal(t, fileContents, string(writtenContents))
+
+	filesInTempDir, err := ioutil.ReadDir(dir)
+	require.NoError(t, err)
+	require.Len(t, filesInTempDir, 1)
+	require.Equal(t, filepath.Base(filePath), filesInTempDir[0].Name())
+}
+
+func TestFileContextCancelled(t *testing.T) {
+	dir, cleanup := testhelper.TempDir(t, "", t.Name())
+	defer cleanup()
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	filePath := filepath.Join(dir, "test_file_contents")
+	fileContents := "very important contents"
+	file, err := safe.CreateFileWriter(ctx, filePath)
+	require.NoError(t, err)
+
+	_, err = io.Copy(file, bytes.NewBufferString(fileContents))
+	require.NoError(t, err)
+
+	testhelper.AssertFileNotExists(t, filePath)
+
+	files, err := ioutil.ReadDir(dir)
+	require.NoError(t, err)
+	require.Len(t, files, 1, "expect only the temp file to exist")
+	require.Contains(t, files[0].Name(), "test_file_contents", "the one file in the directory should be the temp file")
+
+	cancel()
+
+	testhelper.AssertFileNotExists(t, filePath)
+
+	// wait for the cleanup functions to run
+	for i := 0; i < 10; i++ {
+		time.Sleep(10 * time.Millisecond)
+		filesInTempDir, err := ioutil.ReadDir(dir)
+		require.NoError(t, err)
+		if len(filesInTempDir) > 0 {
+			continue
+		}
+		return
+	}
+	t.Error("directory should not have any files left")
+}
+
+func TestFileRace(t *testing.T) {
+	dir, cleanup := testhelper.TempDir(t, "", t.Name())
+	defer cleanup()
+
+	filePath := filepath.Join(dir, "test_file_contents")
+
+	var wg sync.WaitGroup
+
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func(i int) {
+			ctx, cancel := testhelper.Context()
+			defer cancel()
+			w, err := safe.CreateFileWriter(ctx, filePath)
+			require.NoError(t, err)
+			_, err = w.Write([]byte(fmt.Sprintf("message # %d", i)))
+			require.NoError(t, err)
+			require.NoError(t, w.Commit())
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+
+	require.FileExists(t, filePath)
+	filesInTempDir, err := ioutil.ReadDir(dir)
+	require.NoError(t, err)
+	require.Len(t, filesInTempDir, 1, "make sure no other files were written")
+}
diff --git a/internal/testhelper/testhelper.go b/internal/testhelper/testhelper.go
index 7342069ee..2ff9c0a3f 100644
--- a/internal/testhelper/testhelper.go
+++ b/internal/testhelper/testhelper.go
@@ -495,3 +495,12 @@ func CreateLooseRef(t *testing.T, repoPath, refName string) {
 	MustRunCommand(t, nil, "git", "-C", repoPath, "update-ref", relRefPath, "master")
 	require.FileExists(t, filepath.Join(repoPath, relRefPath), "ref must be in loose file")
 }
+
+// TempDir is a wrapper around ioutil.TempDir that provides a cleanup function
+func TempDir(t *testing.T, dir, prefix string) (string, func() error) {
+	dirPath, err := ioutil.TempDir(dir, prefix)
+	require.NoError(t, err)
+	return dirPath, func() error {
+		return os.RemoveAll(dirPath)
+	}
+}