PR Review

Author: lukasfrank

claimutils/claim/claimer.go

@@ -7,7 +7,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"sync/atomic"
+	"sync"
 
 	"github.com/go-logr/logr"
 	"github.com/ironcore-dev/ironcore/api/core/v1alpha1"
@@ -26,15 +26,19 @@ type Claimer interface {
 	Claim(ctx context.Context, resources v1alpha1.ResourceList) (Claims, error)
 	Release(ctx context.Context, claims Claims) error
 	Start(ctx context.Context) error
+	WaitUntilStarted(ctx context.Context) error
 }
 
-func NewResourceClaimer(log logr.Logger, started chan struct{}, plugins ...Plugin) (*claimer, error) {
+func NewResourceClaimer(log logr.Logger, plugins ...Plugin) (*claimer, error) {
 	c := claimer{
-		log:       log,
-		started:   started,
-		plugins:   map[string]Plugin{},
+		log:     log,
+		plugins: map[string]Plugin{},
+
 		toClaim:   make(chan claimReq, 1),
 		toRelease: make(chan releaseReq, 1),
+
+		started:  make(chan struct{}),
+		shutdown: make(chan struct{}),
 	}
 
 	for _, plugin := range plugins {
@@ -74,8 +78,9 @@ type claimer struct {
 	toClaim   chan claimReq
 	toRelease chan releaseReq
 
-	running atomic.Bool
-	started chan struct{}
+	startOnce sync.Once
+	started   chan struct{}
+	shutdown  chan struct{}
 }
 
 func (c *claimer) checkPluginsForResources(resources v1alpha1.ResourceList) error {
@@ -107,13 +112,14 @@ func (c *claimer) checkPluginsForClaims(claims Claims) error {
 }
 
 func (c *claimer) Start(ctx context.Context) error {
-	if !c.running.CompareAndSwap(false, true) {
-		return ErrAlreadyStarted
-	}
+	var called bool
+	c.startOnce.Do(func() {
+		called = true
+		go c.run(ctx)
+	})
 
-	go c.run(ctx)
-	if c.started != nil {
-		close(c.started)
+	if !called {
+		return ErrAlreadyStarted
 	}
 
 	<-ctx.Done()
@@ -123,21 +129,23 @@ func (c *claimer) Start(ctx context.Context) error {
 
 func (c *claimer) run(ctx context.Context) {
 	defer func() {
-		c.running.Store(false)
-		close(c.toClaim)
-		close(c.toRelease)
-
 		for req := range c.toClaim {
 			req.resultChan <- claimRes{err: ctx.Err()}
 		}
 		for req := range c.toRelease {
 			req.resultChan <- ctx.Err()
 		}
+
+		close(c.toClaim)
+		close(c.toRelease)
 	}()
 
+	close(c.started)
+
 	for {
 		select {
 		case <-ctx.Done():
+			close(c.shutdown)
 			return
 		case req := <-c.toClaim:
 			res := claimRes{}
@@ -154,6 +162,22 @@ func (c *claimer) run(ctx context.Context) {
 	}
 }
 
+func (c *claimer) ensureRunning() error {
+	select {
+	case <-c.started:
+	default:
+		return ErrNotStarted
+	}
+
+	select {
+	case <-c.shutdown:
+		return ErrNotStarted
+	default:
+	}
+
+	return nil
+}
+
 func (c *claimer) claim(resources v1alpha1.ResourceList) (Claims, error) {
 	var insufficientResourceErrors []error
 	for resourceName := range resources {
@@ -192,8 +216,8 @@ func (c *claimer) Claim(ctx context.Context, resources v1alpha1.ResourceList) (C
 		return nil, errors.Join(ErrMissingPlugins, err)
 	}
 
-	if !c.running.Load() {
-		return nil, ErrNotStarted
+	if err := c.ensureRunning(); err != nil {
+		return nil, err
 	}
 
 	req := claimReq{
@@ -202,6 +226,8 @@ func (c *claimer) Claim(ctx context.Context, resources v1alpha1.ResourceList) (C
 	}
 	select {
 	case c.toClaim <- req:
+	case <-c.shutdown:
+		return nil, ErrNotStarted
 	case <-ctx.Done():
 		return nil, ctx.Err()
 	}
@@ -235,16 +261,17 @@ func (c *claimer) Release(ctx context.Context, claims Claims) error {
 		return errors.Join(ErrMissingPlugins, err)
 	}
 
-	if !c.running.Load() {
-		return ErrNotStarted
+	if err := c.ensureRunning(); err != nil {
+		return err
 	}
-
 	req := releaseReq{
 		claims:     claims,
 		resultChan: make(chan error, 1),
 	}
 	select {
 	case c.toRelease <- req:
+	case <-c.shutdown:
+		return ErrNotStarted
 	case <-ctx.Done():
 		return ctx.Err()
 	}
@@ -256,3 +283,12 @@ func (c *claimer) Release(ctx context.Context, claims Claims) error {
 		return res
 	}
 }
+
+func (c *claimer) WaitUntilStarted(ctx context.Context) error {
+	select {
+	case <-c.started:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}

claimutils/claim/claimer_test.go

@@ -27,13 +27,9 @@ func (m *mockReader) Read() ([]pci.Address, error) {
 
 var _ = Describe("Resource Claimer", func() {
 	It("should claim composite resources", func(ctx SpecContext) {
-
-		claimerStarted := make(chan struct{})
-
 		By("init plugin")
 		resourceClaimer, err := claim.NewResourceClaimer(
 			log.FromContext(ctx),
-			claimerStarted,
 			gpu.NewGPUClaimPlugin(log.FromContext(ctx), "nvidia.com/gpu", &mockReader{
 				devices: []pci.Address{
 					{},
@@ -59,9 +55,9 @@ var _ = Describe("Resource Claimer", func() {
 		})
 
 		By("waiting until claimer is started")
-		Eventually(claimerStarted).To(BeClosed())
+		Expect(resourceClaimer.WaitUntilStarted(ctx)).To(Succeed())
 
-		By("failing if not not existing resource is claimed")
+		By("failing if nonexistent resource is claimed")
 		resourceClaim, err := resourceClaimer.Claim(ctx, v1alpha1.ResourceList{
 			"not_existing_plugin": resource.MustParse("1"),
 		})

claimutils/gpu/gpu_test.go

@@ -47,7 +47,7 @@ var _ = Describe("GPU Claimer", func() {
 		Expect(plugin.Init()).Should(MatchError(testErr))
 	})
 
-	It("should error if no resource left (not init)", func(ctx SpecContext) {
+	It("should error if no resource left after init", func(ctx SpecContext) {
 		By("init plugin")
 		plugin := gpu.NewGPUClaimPlugin(log.FromContext(ctx), "test-plugin", &MockReader{}, nil)
 		Expect(plugin.Init()).ShouldNot(HaveOccurred())
@@ -130,20 +130,20 @@ var _ = Describe("GPU Claimer", func() {
 		gpuClaim1, err := plugin.Claim(resource.MustParse("1"))
 		Expect(err).ToNot(HaveOccurred())
 
-		ociAddress1, ok := gpuClaim1.(gpu.Claim)
+		pciAddress1, ok := gpuClaim1.(gpu.Claim)
 		Expect(ok).To(BeTrue())
-		Expect(ociAddress1.PCIAddresses()).To(HaveLen(1))
+		Expect(pciAddress1.PCIAddresses()).To(HaveLen(1))
 
 		By("claim resources again")
 		gpuClaim2, err := plugin.Claim(resource.MustParse("1"))
 		Expect(err).ToNot(HaveOccurred())
 
-		ociAddress2, ok := gpuClaim2.(gpu.Claim)
+		pciAddress2, ok := gpuClaim2.(gpu.Claim)
 		Expect(ok).To(BeTrue())
-		Expect(ociAddress2.PCIAddresses()).To(HaveLen(1))
+		Expect(pciAddress2.PCIAddresses()).To(HaveLen(1))
 
 		By("ensure claims are not equal")
-		Expect(ociAddress1.PCIAddresses()[0]).NotTo(Equal(ociAddress2.PCIAddresses()[0]))
+		Expect(pciAddress1.PCIAddresses()[0]).NotTo(Equal(pciAddress2.PCIAddresses()[0]))
 	})
 
 	It("should handle zero-quantity claims", func(ctx SpecContext) {