Repository: kubernetes/kubernetes
File: pkg/scheduler/framework/plugins/dynamicresources/dra_manager.go (Go)

Code before:
	obj, err := c.cache.Get(namespace + "/" + claimName)
	if err != nil {
		return nil, err
	}
	claim, ok := obj.(*resourceapi.ResourceClaim)
	if !ok {
		return nil, fmt.Errorf("unexpected object type %T for assumed object %s/%s", obj, namespace, claimName)
	}
	return claim, nil
}

func (c *claimTracker) List() ([]*resourceapi.ResourceClaim, error) {
	var result []*resourceapi.ResourceClaim
	// Probably not worth adding an index for?
	objs := c.cache.List(nil)
	for _, obj := range objs {
		claim, ok := obj.(*resourceapi.ResourceClaim)
		if ok {
			result = append(result, claim)
		}
	}
	return result, nil
}

func (c *claimTracker) ListAllAllocatedDevices() (sets.Set[structured.DeviceID], error) {
	// There is a rare race when a claim is initially in-flight:
	// - allocated is created from cache (claim not there)
	// - someone removes from the in-flight claims and adds to the cache
	// - we start checking in-flight claims (claim not there anymore)
	// => claim ignored
	//
	// A proper fix would be to rewrite the assume cache, allocatedDevices,
	// and the in-flight map so that they are under a single lock. But that's
	// a pretty big change and prevents reusing the assume cache. So instead
	// we check for changes in the set of allocated devices and keep trying
	// until we get an attempt with no concurrent changes.
	//
	// A claim being first in the cache, then only in-flight cannot happen,
	// so we don't need to re-check the in-flight claims.
	for {
		// Start with a fresh set that matches the current known state of the
		// world according to the informers.
		revision, allocated := c.allocatedDevices.Get()

		// Whatever is in flight also has to be checked.
		c.inFlightAllocations.Range(func(key, value any) bool {
			claim := value.(*resourceapi.ResourceClaim)
			foreachAllocatedDevice(claim, func(deviceID structured.DeviceID) {
				c.logger.V(6).Info("Device is in flight for allocation", "device", deviceID, "claim", klog.KObj(claim))
				allocated.Insert(deviceID)
			}, false, func(structured.SharedDeviceID) {}, func(structured.DeviceConsumedCapacity) {})

Review comment (bug):
I tried writing a unit test for this race, but in contrast to the AssumeCache with its event handler invocation I don't see a way to control the timing of the calls.

I tried this stress test, but it always passed also without this fix, so I am not committing it:

```Go
/*
Copyright 2024 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dynamicresources

import (
	"sync"
	"testing"

	resourceapi "k8s.io/api/resource/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes/fake"
	"k8s.io/client-go/tools/cache"
	"k8s.io/dynamic-resource-allocation/resourceslice/tracker"
	"k8s.io/dynamic-resource-allocation/structured"
	"k8s.io/kubernetes/pkg/scheduler/util/assumecache"
	"k8s.io/kubernetes/test/utils/ktesting"
)

func TestListAllAllocatedDevicesRace(t *testing.T) { testListAllAllocatedDevicesRace(ktesting.Init(t)) }
func testListAllAllocatedDevicesRace(tCtx ktesting.TContext) {
	driver := "dra-driver"
	pool := "test-pool"
	device := "device-1"
	claim := &resourceapi.ResourceClaim{
		ObjectMeta: metav1.ObjectMeta{
			Name:            "test-claim",
			Namespace:       "default",
			UID:             "abc",
			ResourceVersion: "1",
		},
		Status: resourceapi.ResourceClaimStatus{
			Allocation: &resourceapi.AllocationResult{
				Devices: resourceapi.DeviceAllocationResult{
					Results: []resourceapi.DeviceRequestAllocationResult{{
						Driver: driver,
						Pool:   pool,
						Device: device,
					}},
				},
			},
		},
	}
	client := fake.NewClientset(claim)
	informerFactory := informers.NewFilteredSharedInformerFactory(client, 0, "", nil)
	ac := assumecache.NewAssumeCache(tCtx.Logger(), informerFactory.Resource().V1().ResourceClaims().Informer(), "", "", nil)
	tracker, err := tracker.StartTracker(tCtx, tracker.Options{
		SliceInformer: informerFactory.Resource().V1().ResourceSlices(),
	})
	tCtx.ExpectNoError(err, "create ResourceSlice tracker")
	defer tracker.Stop()
	manager := NewDRAManager(tCtx, ac, tracker, informerFactory)
	informerFactory.Start(tCtx.Done())
	defer func() {
		tCtx.Cancel("test completed")
		informerFactory.Shutdown()
	}()
	informerFactory.WaitForCacheSync(tCtx.Done())
	handle := ac.AddEventHandler(cache.ResourceEventHandlerFuncs{})
	cache.WaitForCacheSync(tCtx.Done(), handle.HasSynced)

	// The race occurs when a claim gets moved from in-flight to the cache
	// while ListAllAllocatedDevicesRace runs. We cannot enforce the timing
	// to go wrong, but we can repeat multiple times.
	var wg sync.WaitGroup
	expectDeviceID := structured.MakeDeviceID(driver, pool, device)
	expectAllocatedDevices := sets.New(expectDeviceID)
	for range 100000 {
		manager.ResourceClaims().SignalClaimPendingAllocation(claim.UID, claim)
		var allocatedDevices sets.Set[structured.DeviceID]
		var allocatedDevicesErr error
		wg.Go(func() {
			allocatedDevices, allocatedDevicesErr = manager.resourceClaimTracker.ListAllAllocatedDevices()
		})
		var assumeErr error
		var removed bool
		wg.Go(func() {
			assumeErr = ac.Assume(claim)
			removed = manager.ResourceClaims().RemoveClaimPendingAllocation(claim.UID)
		})
		wg.Wait()
		tCtx.ExpectNoError(assumeErr, "Assume failed")
		if !removed {
			tCtx.Fatal("In-flight claim should have been removed and wasn't.")
		}
		tCtx.ExpectNoError(allocatedDevicesErr, "ListAllAllocatedDevices failed")
		if !allocatedDevices.Equal(expectAllocatedDevices) {
			tCtx.Fatalf("Expected exactly one allocated device %s, got: %v", expectDeviceID, allocatedDevices.UnsortedList())
		}
		ac.Restore(claim.Namespace + "/" + claim.Name)
	}
}
```

Code after:
func (c *claimTracker) Get(namespace, claimName string) (*resourceapi.ResourceClaim, error) {
	obj, err := c.cache.Get(namespace + "/" + claimName)
	if err != nil {
		return nil, err
	}
	claim, ok := obj.(*resourceapi.ResourceClaim)
	if !ok {
		return nil, fmt.Errorf("unexpected object type %T for assumed object %s/%s", obj, namespace, claimName)
	}
	return claim, nil
}

func (c *claimTracker) List() ([]*resourceapi.ResourceClaim, error) {
	var result []*resourceapi.ResourceClaim
	// Probably not worth adding an index for?
	objs := c.cache.List(nil)
	for _, obj := range objs {
		claim, ok := obj.(*resourceapi.ResourceClaim)
		if ok {
			result = append(result, claim)
		}
	}
	return result, nil
}

// errClaimTrackerConcurrentModification gets returned if ListAllAllocatedDevices
// or GatherAllocatedState need to be retried.
//
// There is a rare race when a claim is initially in-flight:
// - allocated is created from cache (claim not there)
// - someone removes from the in-flight claims and adds to the cache
// - we start checking in-flight claims (claim not there anymore)
// => claim ignored
//
// A proper fix would be to rewrite the assume cache, allocatedDevices,
// and the in-flight map so that they are under a single lock. But that's
// a pretty big change and prevents reusing the assume cache. So instead
// we check for changes in the set of allocated devices and keep trying
// until we get an attempt with no concurrent changes.
//
// A claim being first in the cache, then only in-flight cannot happen,
// so we don't need to re-check the in-flight claims.
var errClaimTrackerConcurrentModification = errors.New("conflicting concurrent modification")

func (c *claimTracker) ListAllAllocatedDevices() (sets.Set[structured.DeviceID], error) {
	// Start with a fresh set that matches the current known state of the
	// world according to the informers.
	allocated, revision := c.allocatedDevices.Get()

	// Whatever is in flight also has to be checked.
	c.inFlightAllocations.Range(func(key, value any) bool {