Merge pull request #18605 from lucasrod16/use-single-function-to-create-WAL-files

Add function to create WAL files

Author: ahrtr

server/storage/wal/file_pipeline.go

@@ -75,7 +75,7 @@ func (fp *filePipeline) Close() error {
 func (fp *filePipeline) alloc() (f *fileutil.LockedFile, err error) {
 	// count % 2 so this file isn't the same as the one last published
 	fpath := filepath.Join(fp.dir, fmt.Sprintf("%d.tmp", fp.count%2))
-	if f, err = fileutil.LockFile(fpath, os.O_CREATE|os.O_WRONLY, fileutil.PrivateFileMode); err != nil {
+	if f, err = createNewWALFile[*fileutil.LockedFile](fpath, false); err != nil {
 		return nil, err
 	}
 	if err = fileutil.Preallocate(f.File, fp.size, true); err != nil {

server/storage/wal/repair.go

@@ -67,7 +67,7 @@ func Repair(lg *zap.Logger, dirpath string) bool {
 
 		case errors.Is(err, io.ErrUnexpectedEOF):
 			brokenName := f.Name() + ".broken"
-			bf, bferr := os.OpenFile(brokenName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, fileutil.PrivateFileMode)
+			bf, bferr := createNewWALFile[*os.File](brokenName, true)
 			if bferr != nil {
 				lg.Warn("failed to create backup file", zap.String("path", brokenName), zap.Error(bferr))
 				return false

server/storage/wal/wal.go

@@ -126,7 +126,7 @@ func Create(lg *zap.Logger, dirpath string, metadata []byte) (*WAL, error) {
 	}
 
 	p := filepath.Join(tmpdirpath, walName(0, 0))
-	f, err := fileutil.LockFile(p, os.O_WRONLY|os.O_CREATE, fileutil.PrivateFileMode)
+	f, err := createNewWALFile[*fileutil.LockedFile](p, false)
 	if err != nil {
 		lg.Warn(
 			"failed to flock an initial WAL file",
@@ -233,6 +233,31 @@ func Create(lg *zap.Logger, dirpath string, metadata []byte) (*WAL, error) {
 	return w, nil
 }
 
+// createNewWALFile creates a WAL file.
+// To create a locked file, use *fileutil.LockedFile type parameter.
+// To create a standard file, use *os.File type parameter.
+// If forceNew is true, the file will be truncated if it already exists.
+func createNewWALFile[T *os.File | *fileutil.LockedFile](path string, forceNew bool) (T, error) {
+	flag := os.O_WRONLY | os.O_CREATE
+	if forceNew {
+		flag |= os.O_TRUNC
+	}
+
+	if _, isLockedFile := any(T(nil)).(*fileutil.LockedFile); isLockedFile {
+		lockedFile, err := fileutil.LockFile(path, flag, fileutil.PrivateFileMode)
+		if err != nil {
+			return nil, err
+		}
+		return any(lockedFile).(T), nil
+	}
+
+	file, err := os.OpenFile(path, flag, fileutil.PrivateFileMode)
+	if err != nil {
+		return nil, err
+	}
+	return any(file).(T), nil
+}
+
 func (w *WAL) Reopen(lg *zap.Logger, snap walpb.Snapshot) (*WAL, error) {
 	err := w.Close()
 	if err != nil {

server/storage/wal/wal_test.go

@@ -96,6 +96,75 @@ func TestNew(t *testing.T) {
 	}
 }
 
+func TestCreateNewWALFile(t *testing.T) {
+	tests := []struct {
+		name     string
+		fileType interface{}
+		forceNew bool
+	}{
+		{
+			name:     "creating standard file should succeed and not truncate file",
+			fileType: &os.File{},
+			forceNew: false,
+		},
+		{
+			name:     "creating locked file should succeed and not truncate file",
+			fileType: &fileutil.LockedFile{},
+			forceNew: false,
+		},
+		{
+			name:     "creating standard file with forceNew should truncate file",
+			fileType: &os.File{},
+			forceNew: true,
+		},
+		{
+			name:     "creating locked file with forceNew should truncate file",
+			fileType: &fileutil.LockedFile{},
+			forceNew: true,
+		},
+	}
+
+	for i, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			p := filepath.Join(t.TempDir(), walName(0, uint64(i)))
+
+			// create initial file with some data to verify truncate behavior
+			err := os.WriteFile(p, []byte("test data"), fileutil.PrivateFileMode)
+			require.NoError(t, err)
+
+			var f interface{}
+			switch tt.fileType.(type) {
+			case *os.File:
+				f, err = createNewWALFile[*os.File](p, tt.forceNew)
+				require.IsType(t, &os.File{}, f)
+			case *fileutil.LockedFile:
+				f, err = createNewWALFile[*fileutil.LockedFile](p, tt.forceNew)
+				require.IsType(t, &fileutil.LockedFile{}, f)
+			default:
+				panic("unknown file type")
+			}
+
+			require.NoError(t, err)
+
+			// validate the file permissions
+			fi, err := os.Stat(p)
+			require.NoError(t, err)
+			expectedPerms := fmt.Sprintf("%o", os.FileMode(fileutil.PrivateFileMode))
+			actualPerms := fmt.Sprintf("%o", fi.Mode().Perm())
+			require.Equal(t, expectedPerms, actualPerms, "unexpected file permissions on %q", p)
+
+			content, err := os.ReadFile(p)
+			require.NoError(t, err)
+
+			if tt.forceNew {
+				require.Empty(t, string(content), "file content should be truncated but it wasn't")
+			} else {
+				require.Equal(t, "test data", string(content), "file content should not be truncated but it was")
+			}
+		})
+	}
+}
+
 func TestCreateFailFromPollutedDir(t *testing.T) {
 	p := t.TempDir()
 	os.WriteFile(filepath.Join(p, "test.wal"), []byte("data"), os.ModeTemporary)