Add TryLockPath, CreateAndLock, UnlockAndDelete functions

Signed-off-by: Jan Rodák <[email protected]>

Author: Honny1

internal/staging_lockfile/staging_lockfile.go

@@ -2,6 +2,7 @@ package staging_lockfile
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
 	"sync"
 
@@ -19,14 +20,13 @@ type StagingLockFile struct {
 	// They are safe to access without any other locking.
 	file string
 
-	// rwMutex serializes concurrent reader-writer acquisitions in the same process space
-	rwMutex *sync.RWMutex
-	// stateMutex is used to synchronize concurrent accesses to the state below
-	stateMutex *sync.Mutex
-	locked     bool
-	fd         rawfilelock.FileHandle
+	// The following fields are only set when the lock is acquired, and must never be modified afterwards.
+	locked bool
+	fd     rawfilelock.FileHandle
 }
 
+const maxRetries = 1000
+
 var (
 	stagingLockFiles    map[string]*StagingLockFile
 	stagingLockFileLock sync.Mutex
@@ -41,7 +41,7 @@ func (l *StagingLockFile) AssertLocked() {
 	//
 	// Hence, this “AssertLocked” method, which exists only for sanity checks.
 
-	// Don’t even bother with l.stateMutex: The caller is expected to hold the lock, and in that case l.locked is constant true
+	// The caller is expected to hold the lock, and in that case l.locked is constant true
 	// with no possible writers.
 	// If the caller does not hold the lock, we are violating the locking/memory model anyway, and accessing the data
 	// without the lock is more efficient for callers, and potentially more visible to lock analysers for incorrect callers.
@@ -50,75 +50,124 @@ func (l *StagingLockFile) AssertLocked() {
 	}
 }
 
-// getLockfile returns a StagingLockFile object associated with the specified path.
-// It ensures only one StagingLockFile object exists per path within the process.
-// If a StagingLockFile for the path already exists, it returns that instance.
-// Otherwise, it creates a new one.
-func getLockfile(path string) (*StagingLockFile, error) {
+// tryAcquireLockForFile attempts to acquire a lock for the specified file path.
+// It first checks if the lock is already in use by another thread.
+// If the lock is not in use, it creates a new StagingLockFile, opens the file, and tries to lock it.
+// If successful, it adds the StagingLockFile to the global map of staging lock files.
+// If the lock is already in use, it returns an error.
+func tryAcquireLockForFile(cleanPath string) (*StagingLockFile, error) {
 	stagingLockFileLock.Lock()
 	defer stagingLockFileLock.Unlock()
+
 	if stagingLockFiles == nil {
 		stagingLockFiles = make(map[string]*StagingLockFile)
 	}
-	cleanPath, err := filepath.Abs(path)
-	if err != nil {
-		return nil, fmt.Errorf("ensuring that path %q is an absolute path: %w", path, err)
+
+	if _, ok := stagingLockFiles[cleanPath]; ok {
+		return nil, fmt.Errorf("lock is used already with other thread %q", cleanPath)
 	}
-	if lockFile, ok := stagingLockFiles[cleanPath]; ok {
-		return lockFile, nil
+
+	lockFile := &StagingLockFile{
+		file:   cleanPath,
+		locked: false,
 	}
-	lockFile, err := createStagingLockFileForPath(cleanPath) // platform-dependent LockFile
+
+	fd, err := rawfilelock.OpenLock(lockFile.file, false)
 	if err != nil {
 		return nil, err
 	}
+	lockFile.fd = fd
+
+	if err = rawfilelock.TryLockFile(lockFile.fd, rawfilelock.WriteLock); err != nil {
+		rawfilelock.CloseHandle(lockFile.fd) // This is safe because we hold stagingLockFileLock so we are the only possible holder of the lock within this process.
+		return nil, fmt.Errorf("failed to acquire lock on %q: %w", cleanPath, err)
+	}
+
+	lockFile.locked = true
+
 	stagingLockFiles[cleanPath] = lockFile
 	return lockFile, nil
 }
 
-// createStagingLockFileForPath creates a new StagingLockFile instance for the given path.
-// It verifies that the file can be opened before returning the StagingLockFile.
-// This function will be called at most once for each unique path within a process.
-func createStagingLockFileForPath(path string) (*StagingLockFile, error) {
-	// Check if we can open the lock.
-	fd, err := rawfilelock.OpenLock(path, false)
-	if err != nil {
-		return nil, err
+// UnlockAndDelete releases the lock, removes the associated file from the filesystem,
+// and removes this StagingLockFile from the global map of StagingLockFile.
+//
+// WARNING: After this operation, the StagingLockFile becomes invalid for further use as the file field is cleared.
+func (l *StagingLockFile) UnlockAndDelete() error {
+	stagingLockFileLock.Lock()
+	defer stagingLockFileLock.Unlock()
+
+	if !l.locked {
+		// Panic when unlocking an unlocked lock. That's a violation
+		// of the lock semantics and will reveal such.
+		panic("calling Unlock on unlocked lock")
 	}
-	rawfilelock.UnlockAndCloseHandle(fd)
-
-	return &StagingLockFile{
-		file:       path,
-		rwMutex:    &sync.RWMutex{},
-		stateMutex: &sync.Mutex{},
-		locked:     false,
-	}, nil
-}
 
-// tryLock attempts to acquire an exclusive lock on the StagingLockFile without blocking.
-// It first tries to acquire the internal rwMutex, then opens and tries to lock the file.
-// Returns nil on success or an error if any step fails.
-func (l *StagingLockFile) tryLock() error {
-	success := l.rwMutex.TryLock()
-	rwMutexUnlocker := l.rwMutex.Unlock
+	delete(stagingLockFiles, l.file)
+	l.locked = false
 
-	if !success {
-		return fmt.Errorf("resource temporarily unavailable")
-	}
-	l.stateMutex.Lock()
-	defer l.stateMutex.Unlock()
-	fd, err := rawfilelock.OpenLock(l.file, false)
-	if err != nil {
-		rwMutexUnlocker()
+	defer func() {
+		rawfilelock.UnlockAndCloseHandle(l.fd)
+		l.file = ""
+	}()
+	if err := os.Remove(l.file); err != nil && !os.IsNotExist(err) {
 		return err
 	}
-	l.fd = fd
+	return nil
+}
 
-	if err = rawfilelock.TryLockFile(l.fd, rawfilelock.WriteLock); err != nil {
-		rawfilelock.CloseHandle(fd)
-		rwMutexUnlocker()
-		return err
+// CreateAndLock creates a new temporary file in the specified directory with the given pattern,
+// then creates and locks a StagingLockFile for it. The file is created using os.CreateTemp.
+// Caller MUST call UnlockAndDelete() on the returned StagingLockFile to release the lock and delete the file.
+//
+// Returns:
+//   - The locked StagingLockFile
+//   - The absolute path to the created file
+//   - Any error that occurred during the process
+//
+// If the file cannot be locked, this function will retry up to maxRetries times before failing.
+func CreateAndLock(dir string, pattern string) (*StagingLockFile, string, error) {
+	for try := 0; ; try++ {
+		file, err := os.CreateTemp(dir, pattern)
+		if err != nil {
+			return nil, "", err
+		}
+		file.Close()
+
+		cleanPath, err := filepath.Abs(file.Name())
+		if err != nil {
+			return nil, "", err
+		}
+
+		l, err := tryAcquireLockForFile(cleanPath)
+		if err != nil {
+			if try < maxRetries {
+				continue // Retry if the lock cannot be acquired
+			}
+			stagingLockFileLock.Lock()
+			delete(stagingLockFiles, cleanPath)
+			stagingLockFileLock.Unlock()
+			return nil, "", fmt.Errorf(
+				"failed to allocate lock in %q after %d attempts; last failure on %q: %w",
+				dir, try, cleanPath, err,
+			)
+		}
+
+		return l, cleanPath, nil
 	}
+}
 
-	l.locked = true
-	return nil
+// TryLockPath attempts to acquire a lock on an existing file. If the file does not exist,
+// it will be created.
+//
+// Warning: If acquiring a lock is successful, it returns a new StagingLockFile
+// instance for the file. Caller MUST call UnlockAndDelete() on the returned StagingLockFile
+// to release the lock and delete the file.
+func TryLockPath(path string) (*StagingLockFile, error) {
+	cleanPath, err := filepath.Abs(path)
+	if err != nil {
+		return nil, fmt.Errorf("ensuring that path %q is an absolute path: %w", path, err)
+	}
+
+	return tryAcquireLockForFile(cleanPath)
 }