Replace in-mem (un)prep lock with file-based lock

Signed-off-by: Dr. Jan-Philip Gehrcke <[email protected]>

Author: jgehrcke

cmd/compute-domain-kubelet-plugin/driver.go

@@ -30,6 +30,7 @@ import (
 	"k8s.io/dynamic-resource-allocation/resourceslice"
 	"k8s.io/klog/v2"
 
+	"github.com/NVIDIA/k8s-dra-driver-gpu/pkg/flock"
 	"github.com/NVIDIA/k8s-dra-driver-gpu/pkg/workqueue"
 )
 
@@ -47,10 +48,10 @@ func isPermanentError(err error) bool {
 }
 
 type driver struct {
-	sync.Mutex
 	client       coreclientset.Interface
 	pluginhelper *kubeletplugin.Helper
 	state        *DeviceState
+	pulock       *flock.Flock
 }
 
 func NewDriver(ctx context.Context, config *Config) (*driver, error) {
@@ -62,6 +63,7 @@ func NewDriver(ctx context.Context, config *Config) (*driver, error) {
 	driver := &driver{
 		client: config.clientsets.Core,
 		state:  state,
+		pulock: flock.NewFlock(DriverPluginPath + "/pu.lock"),
 	}
 
 	helper, err := kubeletplugin.Start(
@@ -184,8 +186,14 @@ func (d *driver) UnprepareResourceClaims(ctx context.Context, claimRefs []kubele
 }
 
 func (d *driver) nodePrepareResource(ctx context.Context, claim *resourceapi.ResourceClaim) (bool, kubeletplugin.PrepareResult) {
-	d.Lock()
-	defer d.Unlock()
+	release, err := d.pulock.Acquire(ctx, flock.WithTimeout(10*time.Second))
+	if err != nil {
+		res := kubeletplugin.PrepareResult{
+			Err: fmt.Errorf("error acquiring prep/unprep lock: %w", err),
+		}
+		return false, res
+	}
+	defer release()
 
 	if claim.Status.Allocation == nil {
 		res := kubeletplugin.PrepareResult{
@@ -207,8 +215,11 @@ func (d *driver) nodePrepareResource(ctx context.Context, claim *resourceapi.Res
 }
 
 func (d *driver) nodeUnprepareResource(ctx context.Context, claimRef kubeletplugin.NamespacedObject) (bool, error) {
-	d.Lock()
-	defer d.Unlock()
+	release, err := d.pulock.Acquire(ctx, flock.WithTimeout(10*time.Second))
+	if err != nil {
+		return false, fmt.Errorf("error acquiring prep/unprep lock: %w", err)
+	}
+	defer release()
 
 	if err := d.state.Unprepare(ctx, claimRef); err != nil {
 		return isPermanentError(err), fmt.Errorf("error unpreparing devices for claim '%v': %w", claimRef.String(), err)

cmd/gpu-kubelet-plugin/driver.go

@@ -19,21 +19,23 @@ package main
 import (
 	"context"
 	"fmt"
-	"sync"
+	"time"
 
 	resourceapi "k8s.io/api/resource/v1beta1"
 	"k8s.io/apimachinery/pkg/types"
 	coreclientset "k8s.io/client-go/kubernetes"
 	"k8s.io/dynamic-resource-allocation/kubeletplugin"
 	"k8s.io/dynamic-resource-allocation/resourceslice"
 	"k8s.io/klog/v2"
+
+	"github.com/NVIDIA/k8s-dra-driver-gpu/pkg/flock"
 )
 
 type driver struct {
-	sync.Mutex
 	client       coreclientset.Interface
 	pluginhelper *kubeletplugin.Helper
 	state        *DeviceState
+	pulock       *flock.Flock
 }
 
 func NewDriver(ctx context.Context, config *Config) (*driver, error) {
@@ -44,6 +46,7 @@ func NewDriver(ctx context.Context, config *Config) (*driver, error) {
 	driver := &driver{
 		client: config.clientsets.Core,
 		state:  state,
+		pulock: flock.NewFlock(DriverPluginPath + "/pu.lock"),
 	}
 
 	helper, err := kubeletplugin.Start(
@@ -110,8 +113,13 @@ func (d *driver) UnprepareResourceClaims(ctx context.Context, claimRefs []kubele
 }
 
 func (d *driver) nodePrepareResource(ctx context.Context, claim *resourceapi.ResourceClaim) kubeletplugin.PrepareResult {
-	d.Lock()
-	defer d.Unlock()
+	release, err := d.pulock.Acquire(ctx, flock.WithTimeout(10*time.Second))
+	if err != nil {
+		return kubeletplugin.PrepareResult{
+			Err: fmt.Errorf("error acquiring prep/unprep lock: %w", err),
+		}
+	}
+	defer release()
 
 	devs, err := d.state.Prepare(ctx, claim)
 
@@ -126,8 +134,11 @@ func (d *driver) nodePrepareResource(ctx context.Context, claim *resourceapi.Res
 }
 
 func (d *driver) nodeUnprepareResource(ctx context.Context, claimNs kubeletplugin.NamespacedObject) error {
-	d.Lock()
-	defer d.Unlock()
+	release, err := d.pulock.Acquire(ctx, flock.WithTimeout(10*time.Second))
+	if err != nil {
+		return fmt.Errorf("error acquiring prep/unprep lock: %w", err)
+	}
+	defer release()
 
 	if err := d.state.Unprepare(ctx, string(claimNs.UID)); err != nil {
 		return fmt.Errorf("error unpreparing devices for claim %v: %w", claimNs.UID, err)

pkg/flock/flock.go

@@ -0,0 +1,116 @@
+package flock
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"syscall"
+	"time"
+)
+
+type Flock struct {
+	path string
+}
+
+type AcquireOption func(*acquireConfig)
+
+type acquireConfig struct {
+	timeout    time.Duration
+	pollPeriod time.Duration
+}
+
+func WithTimeout(timeout time.Duration) AcquireOption {
+	return func(cfg *acquireConfig) {
+		cfg.timeout = timeout
+	}
+}
+
+func WithPollPeriod(period time.Duration) AcquireOption {
+	return func(cfg *acquireConfig) {
+		cfg.pollPeriod = period
+	}
+}
+
+func NewFlock(path string) *Flock {
+	return &Flock{
+		path: path,
+	}
+}
+
+// Acquire attempts to acquire an exclusive file lock, polling with the configured
+// PollPeriod until whichever comes first:
+//
+// - lock successfully acquired
+// - timeout (if provided)
+// - external cancellation of context
+//
+// Returns a release function that must be called to unlock the file, typically
+// with defer().
+//
+// Introduced to protect the work in nodePrepareResource() and
+// nodeUnprepareResource() under a file-based lock because more than one driver
+// pod may be running on a node, but at most one such function must execute at
+// any given time.
+func (l *Flock) Acquire(ctx context.Context, opts ...AcquireOption) (func(), error) {
+	cfg := &acquireConfig{
+		// Default: short period to keep lock acquisition rather responsive
+		pollPeriod: 200 * time.Millisecond,
+		// Default: timeout disabled
+		timeout: 0,
+	}
+	for _, opt := range opts {
+		opt(cfg)
+	}
+
+	f, oerr := os.OpenFile(l.path, os.O_RDWR|os.O_CREATE, 0644)
+	if oerr != nil {
+		return nil, fmt.Errorf("error opening lock file (%s): %w", l.path, oerr)
+	}
+
+	t0 := time.Now()
+	ticker := time.NewTicker(cfg.pollPeriod)
+	defer ticker.Stop()
+
+	// Use non-blocking peek with LOCK_NB flag and polling; trade-off:
+	//
+	// - pro: no need for having to reliably cancel a potentially long-blocking
+	//   flock() system call (can only be done with signals).
+	// - con: lock acquisition time after a release is not immediate, but may
+	//   take up to PollPeriod amount of time. Not an issue in this context.
+	for {
+		flerr := syscall.Flock(int(f.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
+		if flerr == nil {
+			// Lock acquired. Return release function. An exclusive flock() lock
+			// gets released when its file descriptor gets closed (also true
+			// when the lock-holding process crashes).
+			return func() {
+				f.Close()
+			}, nil
+		}
+
+		if flerr != syscall.EWOULDBLOCK {
+			// May be EBADF, EINTR, EINVAl, ENOLCK, and
+			// in general we want an outer retry mechanism
+			// to retry in view of any of those.
+			f.Close()
+			return nil, fmt.Errorf("error acquiring lock (%s): %w", l.path, flerr)
+		}
+
+		// Lock is currently held by other entity. Check for exit criteria;
+		// otherwise retry lock acquisition upon next tick.
+
+		if cfg.timeout > 0 && time.Since(t0) > cfg.timeout {
+			f.Close()
+			return nil, fmt.Errorf("timeout acquiring lock (%s)", l.path)
+		}
+
+		select {
+		case <-ctx.Done():
+			f.Close()
+			return nil, fmt.Errorf("context done before lock acquired (%s)", l.path)
+		case <-ticker.C:
+			// Retry flock().
+			continue
+		}
+	}
+}