Introduce yet another file locker

Author: pchila

internal/pkg/agent/application/upgrade/step_mark.go

@@ -5,6 +5,7 @@
 package upgrade
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -229,7 +230,7 @@ func markUpgrade(log *logger.Logger, dataDirPath string, agent, previousAgent ag
 		return fmt.Errorf("failed locking marker file: %w", err)
 	}
 
-	defer func(markerFileLock *flock.Flock) {
+	defer func(markerFileLock Locker) {
 		err := markerFileLock.Unlock()
 		if err != nil {
 			log.Warnw("Failed to unlock marker file", "file.path", markerPath, "err", err)
@@ -338,15 +339,95 @@ func markerFilePath(dataDirPath string) string {
 	return filepath.Join(dataDirPath, markerFilename)
 }
 
-func lockMarkerFile(markerFileName string) (*flock.Flock, error) {
-	fLock := flock.New(markerFileName + ".lock")
-	locked, err := fLock.TryLock()
+func fileLockName(markerFileName string) string {
+	return markerFileName + ".lock"
+}
+
+func lockMarkerFile(markerFileName string) (Locker, error) {
+	locker, err := NewFileLocker(fileLockName(markerFileName), WithTimeout(10*time.Second))
+	if err != nil {
+		return nil, fmt.Errorf("instantiating file locker: %w", err)
+	}
+
+	err = locker.Lock()
 	if err != nil {
-		return nil, fmt.Errorf("locking %s: %w", fLock, err)
+		return nil, fmt.Errorf("locking: %w", err)
+	}
+	return locker, nil
+}
+
+type Locker interface {
+	Lock() error
+	Unlock() error
+}
+
+type FileLocker struct {
+	fileLock             *flock.Flock
+	blocking             bool
+	timeout              time.Duration
+	customNotLockedError error
+}
+
+type FileLockerOption func(locker *FileLocker) error
+
+func WithCustomNotLockedError(customError error) FileLockerOption {
+	return func(locker *FileLocker) error {
+		locker.customNotLockedError = customError
+		return nil
+	}
+}
+
+var ErrZeroTimeout = errors.New("must specify a non-zero timeout for a blocking file locker")
+var ErrNotLocked = errors.New("file not locked")
+
+func WithTimeout(timeout time.Duration) FileLockerOption {
+	return func(locker *FileLocker) error {
+
+		if timeout == 0 {
+			return ErrZeroTimeout
+		}
+
+		locker.blocking = true
+		locker.timeout = timeout
+
+		return nil
+	}
+}
+
+func NewFileLocker(lockFilePath string, opts ...FileLockerOption) (*FileLocker, error) {
+	flocker := &FileLocker{fileLock: flock.New(lockFilePath)}
+	for _, opt := range opts {
+		if err := opt(flocker); err != nil {
+			return nil, fmt.Errorf("applying options to new file locker: %w", err)
+		}
+	}
+	return flocker, nil
+}
+
+func (fl *FileLocker) Lock() error {
+	var locked bool
+	var err error
+
+	if fl.blocking {
+		timeoutCtx, cancel := context.WithTimeout(context.TODO(), fl.timeout)
+		defer cancel()
+		locked, err = fl.fileLock.TryLockContext(timeoutCtx, time.Second)
+	} else {
+		locked, err = fl.fileLock.TryLock()
+	}
+
+	if err != nil {
+		return fmt.Errorf("locking %s: %w", fl.fileLock.Path(), err)
 	}
 	if !locked {
-		return nil, fmt.Errorf("failed locking %s", fLock)
+		if fl.customNotLockedError != nil {
+			return fmt.Errorf("failed locking %s: %w", fl.fileLock.Path(), fl.customNotLockedError)
+		}
+		return fmt.Errorf("failed locking %s: %w", fl.fileLock.Path(), ErrNotLocked)
 	}
+	return nil
+}
 
-	return fLock, nil
+func (fl *FileLocker) Unlock() error {
+	return fl.fileLock.Unlock()
 }

internal/pkg/agent/application/upgrade/step_mark_test.go

@@ -373,11 +373,11 @@ func Test_markUpgradeLocking(t *testing.T) {
 		},
 		{
 			name: "Update marker creation should fail if marker is already locked by another process", args: args{
-			agent:          newAgent456,
-			previousAgent:  prevAgent123,
-			action:         nil,
-			upgradeDetails: nil,
-		},
+				agent:          newAgent456,
+				previousAgent:  prevAgent123,
+				action:         nil,
+				upgradeDetails: nil,
+			},
 			beforeUpdateMarkerCreation: func(t *testing.T, dataDir string) {
 				// write some fake data in update marker file
 				updateMarkerFilePath := markerFilePath(dataDir)