From acdb1ae8806b0e26e3ff5b71a9ae7916ecf60c43 Mon Sep 17 00:00:00 2001
From: Giuseppe Scrivano <gscrivan@redhat.com>
Date: Tue, 30 Apr 2024 14:56:57 +0200
Subject: [PATCH] lockfile: add functions for non blocking lock

extend the public API to allow a non blocking usage.

Signed-off-by: Giuseppe Scrivano <gscrivan@redhat.com>
---
 pkg/lockfile/lockfile.go         | 48 +++++++++++++++++++++++++++++++-
 pkg/lockfile/lockfile_test.go    |  8 ++++--
 pkg/lockfile/lockfile_unix.go    | 16 +++++++++--
 pkg/lockfile/lockfile_windows.go | 13 ++++++++-
 4 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/pkg/lockfile/lockfile.go b/pkg/lockfile/lockfile.go
index 19694be3bc..3ccd8ff57e 100644
--- a/pkg/lockfile/lockfile.go
+++ b/pkg/lockfile/lockfile.go
@@ -138,6 +138,20 @@ func (l *LockFile) RLock() {
 	l.lock(readLock)
 }
 
+// TryLock attempts to lock the lockfile as a writer.  Panic if the lock is a read-only one.
+func (l *LockFile) TryLock() error {
+	if l.ro {
+		panic("can't take write lock on read-only lock file")
+	} else {
+		return l.tryLock(writeLock)
+	}
+}
+
+// TryRLock attempts to lock the lockfile as a reader.
+func (l *LockFile) TryRLock() error {
+	return l.tryLock(readLock)
+}
+
 // Unlock unlocks the lockfile.
 func (l *LockFile) Unlock() {
 	l.stateMutex.Lock()
@@ -401,9 +415,41 @@ func (l *LockFile) lock(lType lockType) {
 		// Optimization: only use the (expensive) syscall when
 		// the counter is 0.  In this case, we're either the first
 		// reader lock or a writer lock.
-		lockHandle(l.fd, lType)
+		lockHandle(l.fd, lType, false)
 	}
 	l.lockType = lType
 	l.locked = true
 	l.counter++
 }
+
+// lock locks the lockfile via syscall based on the specified type and
+// command.
+func (l *LockFile) tryLock(lType lockType) error {
+	if lType == readLock {
+		l.rwMutex.RLock()
+	} else {
+		l.rwMutex.Lock()
+	}
+	l.stateMutex.Lock()
+	defer l.stateMutex.Unlock()
+	if l.counter == 0 {
+		// If we're the first reference on the lock, we need to open the file again.
+		fd, err := openLock(l.file, l.ro)
+		if err != nil {
+			return err
+		}
+		l.fd = fd
+
+		// Optimization: only use the (expensive) syscall when
+		// the counter is 0.  In this case, we're either the first
+		// reader lock or a writer lock.
+		if err = lockHandle(l.fd, lType, true); err != nil {
+			closeHandle(fd)
+			return err
+		}
+	}
+	l.lockType = lType
+	l.locked = true
+	l.counter++
+	return nil
+}
diff --git a/pkg/lockfile/lockfile_test.go b/pkg/lockfile/lockfile_test.go
index 9512fbf74e..c149dd5275 100644
--- a/pkg/lockfile/lockfile_test.go
+++ b/pkg/lockfile/lockfile_test.go
@@ -213,7 +213,9 @@ func TestLockfileName(t *testing.T) {
 
 	assert.NotEmpty(t, l.name, "lockfile name should be recorded correctly")
 
-	l.Lock()
+	for l.TryLock() != nil {
+		time.Sleep(10 * time.Millisecond)
+	}
 	l.AssertLocked()
 	l.AssertLockedForWriting()
 	l.Unlock()
@@ -226,7 +228,9 @@ func TestLockfileRead(t *testing.T) {
 	require.Nil(t, err, "error getting temporary lock file")
 	defer os.Remove(l.name)
 
-	l.RLock()
+	for l.TryRLock() != nil {
+		time.Sleep(10 * time.Millisecond)
+	}
 	l.AssertLocked()
 	assert.Panics(t, l.AssertLockedForWriting)
 	l.Unlock()
diff --git a/pkg/lockfile/lockfile_unix.go b/pkg/lockfile/lockfile_unix.go
index 38e737e265..4eb76c237a 100644
--- a/pkg/lockfile/lockfile_unix.go
+++ b/pkg/lockfile/lockfile_unix.go
@@ -74,7 +74,7 @@ func openHandle(path string, mode int) (fileHandle, error) {
 	return fileHandle(fd), err
 }
 
-func lockHandle(fd fileHandle, lType lockType) {
+func lockHandle(fd fileHandle, lType lockType, nonblocking bool) error {
 	fType := unix.F_RDLCK
 	if lType != readLock {
 		fType = unix.F_WRLCK
@@ -85,7 +85,15 @@ func lockHandle(fd fileHandle, lType lockType) {
 		Start:  0,
 		Len:    0,
 	}
-	for unix.FcntlFlock(uintptr(fd), unix.F_SETLKW, &lk) != nil {
+	for {
+		cmd := unix.F_SETLKW
+		if nonblocking {
+			cmd = unix.F_SETLK
+		}
+		err := unix.FcntlFlock(uintptr(fd), cmd, &lk)
+		if err == nil || nonblocking {
+			return err
+		}
 		time.Sleep(10 * time.Millisecond)
 	}
 }
@@ -93,3 +101,7 @@ func lockHandle(fd fileHandle, lType lockType) {
 func unlockAndCloseHandle(fd fileHandle) {
 	unix.Close(int(fd))
 }
+
+func closeHandle(fd fileHandle) {
+	unix.Close(int(fd))
+}
diff --git a/pkg/lockfile/lockfile_windows.go b/pkg/lockfile/lockfile_windows.go
index 304c92b158..6482529b3e 100644
--- a/pkg/lockfile/lockfile_windows.go
+++ b/pkg/lockfile/lockfile_windows.go
@@ -81,19 +81,30 @@ func openHandle(path string, mode int) (fileHandle, error) {
 	return fileHandle(fd), err
 }
 
-func lockHandle(fd fileHandle, lType lockType) {
+func lockHandle(fd fileHandle, lType lockType, nonblocking bool) error {
 	flags := 0
 	if lType != readLock {
 		flags = windows.LOCKFILE_EXCLUSIVE_LOCK
 	}
+	if nonblocking {
+		flags |= windows.LOCKFILE_FAIL_IMMEDIATELY
+	}
 	ol := new(windows.Overlapped)
 	if err := windows.LockFileEx(windows.Handle(fd), uint32(flags), reserved, allBytes, allBytes, ol); err != nil {
+		if nonblocking {
+			return err
+		}
 		panic(err)
 	}
+	return nil
 }
 
 func unlockAndCloseHandle(fd fileHandle) {
 	ol := new(windows.Overlapped)
 	windows.UnlockFileEx(windows.Handle(fd), reserved, allBytes, allBytes, ol)
+	closeHandle(fd)
+}
+
+func closeHandle(fd fileHandle) {
 	windows.Close(windows.Handle(fd))
 }