From 07f736fa3c03611c3fd8a8a6bced5ffd88867458 Mon Sep 17 00:00:00 2001
From: Andrey Smirnov <andrey.smirnov@siderolabs.com>
Date: Thu, 15 Aug 2024 16:19:52 +0400
Subject: [PATCH] feat: add support for retry-locking the blockdevice

This allows to wait for others to unlock the blockdevice.

Signed-off-by: Andrey Smirnov <andrey.smirnov@siderolabs.com>
---
 block/device_linux.go      | 27 +++++++++++++++++++++++++++
 block/device_linux_test.go | 19 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/block/device_linux.go b/block/device_linux.go
index 08b22f9..41e9c3f 100644
--- a/block/device_linux.go
+++ b/block/device_linux.go
@@ -6,6 +6,7 @@ package block
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -13,6 +14,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"time"
 	"unsafe"
 
 	"golang.org/x/sys/unix"
@@ -280,6 +282,31 @@ func (d *Device) TryLock(exclusive bool) error {
 	return d.lock(exclusive, unix.LOCK_NB)
 }
 
+// RetryLock until the context deadline.
+func (d *Device) RetryLock(ctx context.Context, exclusive bool) error {
+	ticker := time.NewTicker(100 * time.Millisecond)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ticker.C:
+			if err := d.TryLock(exclusive); err == nil {
+				return nil
+			}
+		}
+	}
+}
+
+// RetryLockWithTimeout retries locking until the timeout.
+func (d *Device) RetryLockWithTimeout(ctx context.Context, exclusive bool, timeout time.Duration) error {
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	return d.RetryLock(ctx, exclusive)
+}
+
 // Unlock releases any lock.
 func (d *Device) Unlock() error {
 	for {
diff --git a/block/device_linux_test.go b/block/device_linux_test.go
index 2331eb9..cdcf44d 100644
--- a/block/device_linux_test.go
+++ b/block/device_linux_test.go
@@ -5,11 +5,13 @@
 package block_test
 
 import (
+	"context"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/freddierice/go-losetup/v2"
 	"github.com/stretchr/testify/assert"
@@ -204,6 +206,23 @@ func TestDevice(t *testing.T) {
 		require.NoError(t, devWhole2.Unlock())
 	})
 
+	t.Run("retry lock", func(t *testing.T) {
+		require.NoError(t, devWhole.Lock(true))
+
+		errCh := make(chan error)
+
+		go func() {
+			errCh <- devWhole2.RetryLockWithTimeout(context.Background(), true, 10*time.Second)
+		}()
+
+		time.Sleep(1 * time.Second)
+
+		require.NoError(t, devWhole.Unlock())
+
+		require.NoError(t, <-errCh)
+		require.NoError(t, devWhole2.Unlock())
+	})
+
 	t.Run("read only", func(t *testing.T) {
 		readOnly, err := devWhole.IsReadOnly()
 		require.NoError(t, err)