From 680eadf0d6bb969266e53d7402b4b2a17adeb33f Mon Sep 17 00:00:00 2001 From: Lucas Rodriguez Date: Tue, 24 Sep 2024 12:55:24 -0500 Subject: [PATCH] Add function to create WAL files Signed-off-by: Lucas Rodriguez --- server/storage/wal/file_pipeline.go | 2 +- server/storage/wal/repair.go | 2 +- server/storage/wal/wal.go | 27 ++++++++++- server/storage/wal/wal_test.go | 69 +++++++++++++++++++++++++++++ 4 files changed, 97 insertions(+), 3 deletions(-) diff --git a/server/storage/wal/file_pipeline.go b/server/storage/wal/file_pipeline.go index bdfa31b5e2a..c8ee4cce429 100644 --- a/server/storage/wal/file_pipeline.go +++ b/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 { diff --git a/server/storage/wal/repair.go b/server/storage/wal/repair.go index d1a887835da..16277540f34 100644 --- a/server/storage/wal/repair.go +++ b/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 diff --git a/server/storage/wal/wal.go b/server/storage/wal/wal.go index 3a313876083..98327d52bcb 100644 --- a/server/storage/wal/wal.go +++ b/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 { diff --git a/server/storage/wal/wal_test.go b/server/storage/wal/wal_test.go index ed3a8893df5..c7679973e36 100644 --- a/server/storage/wal/wal_test.go +++ b/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)