This is an automated email from the ASF dual-hosted git repository.
manirajv06 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/yunikorn-k8shim.git
The following commit(s) were added to refs/heads/master by this push:
new 102e24bb [YUNIKORN-3369] Abort AssumePod retry loop on shim shutdown
(#1078)
102e24bb is described below
commit 102e24bb5b8892c39e41f5ffa36e537ac1c890c2
Author: jimmycmlin <[email protected]>
AuthorDate: Mon Aug 31 14:03:26 2026 +0530
[YUNIKORN-3369] Abort AssumePod retry loop on shim shutdown (#1078)
Closes: #1078
Signed-off-by: Manikandan R <[email protected]>
---
pkg/cache/scheduler_callback.go | 36 ++++++++++++----
pkg/cache/scheduler_callback_test.go | 58 ++++++++++++++++++++++++-
pkg/common/test/volumebinder_mock.go | 17 +++++---
pkg/common/test/volumebinder_mock_test.go | 39 +++++++++++++++++
pkg/shim/scheduler.go | 23 ++++++++--
pkg/shim/scheduler_mock_test.go | 6 ++-
pkg/shim/scheduler_test.go | 71 +++++++++++++++++++++++++++++--
7 files changed, 226 insertions(+), 24 deletions(-)
diff --git a/pkg/cache/scheduler_callback.go b/pkg/cache/scheduler_callback.go
index a2cdaf0a..dd53828f 100644
--- a/pkg/cache/scheduler_callback.go
+++ b/pkg/cache/scheduler_callback.go
@@ -19,12 +19,13 @@
package cache
import (
+ "context"
+ "errors"
"fmt"
"time"
"go.uber.org/zap"
"k8s.io/apimachinery/pkg/util/wait"
- "k8s.io/client-go/util/retry"
"github.com/apache/yunikorn-k8shim/pkg/common/utils"
"github.com/apache/yunikorn-k8shim/pkg/dispatcher"
@@ -37,13 +38,14 @@ import (
// asynchronously to avoid blocking the scheduler.
type AsyncRMCallback struct {
context *Context
+ stopCtx context.Context
}
var _ api.ResourceManagerCallback = &AsyncRMCallback{}
var _ api.StateDumpPlugin = &AsyncRMCallback{}
-func NewAsyncRMCallback(ctx *Context) *AsyncRMCallback {
- return &AsyncRMCallback{context: ctx}
+func NewAsyncRMCallback(ctx *Context, stopCtx context.Context)
*AsyncRMCallback {
+ return &AsyncRMCallback{context: ctx, stopCtx: stopCtx}
}
func (callback *AsyncRMCallback) UpdateAllocation(response
*si.AllocationResponse) error {
@@ -70,13 +72,31 @@ func (callback *AsyncRMCallback) UpdateAllocation(response
*si.AllocationRespons
Duration: time.Second,
Cap: 30 * time.Second,
}
- err := retry.OnError(backOff, func(err error) bool {
- log.Log(log.ShimRMCallback).Error("AssumePod failed,
retrying", zap.Error(err))
- return true
- }, func() error {
- return callback.context.AssumePod(alloc.AllocationKey,
alloc.NodeID)
+ var lastErr error
+ // ConditionWithContextFunc receives a context derived from
callback.stopCtx; closure accesses callback.context directly.
+ err := wait.ExponentialBackoffWithContext(callback.stopCtx,
backOff, func(_ context.Context) (bool, error) {
+ if assumeErr :=
callback.context.AssumePod(alloc.AllocationKey, alloc.NodeID); assumeErr != nil
{
+ log.Log(log.ShimRMCallback).Error("AssumePod
failed, retrying", zap.Error(assumeErr))
+ lastErr = assumeErr
+ return false, nil
+ }
+ return true, nil
})
if err != nil {
+ if errors.Is(err, context.Canceled) || errors.Is(err,
context.DeadlineExceeded) {
+ // Shim is shutting down
(context.DeadlineExceeded is checked defensively for future-proofing).
+ // Do not fall through to
rollbackOnAssumePodFailure:
+ // rollback re-queues the ask to the core
scheduler, which is still running
+ // independently of stopCtx/dispatcher.Stop(),
so the core immediately
+ // re-allocates this task and re-triggers
UpdateAllocation, which hits this
+ // same canceled context and rolls back again —
an unbounded busy loop on the
+ // RM-proxy callback goroutine that Stop() can
never interrupt. See YUNIKORN-3369.
+ log.Log(log.ShimRMCallback).Info("AssumePod
retry aborted due to context cancellation", zap.Error(err))
+ continue
+ }
+ if wait.Interrupted(err) && lastErr != nil {
+ err = lastErr
+ }
if task.IsPlaceholder() {
// Placeholder tasks do not have volume
bindings, so AssumePod failure
// is unexpected and unrecoverable; wrap the
error with context.
diff --git a/pkg/cache/scheduler_callback_test.go
b/pkg/cache/scheduler_callback_test.go
index 90619a77..bf97c437 100644
--- a/pkg/cache/scheduler_callback_test.go
+++ b/pkg/cache/scheduler_callback_test.go
@@ -19,6 +19,7 @@
package cache
import (
+ ctx "context"
"encoding/json"
"strings"
"sync/atomic"
@@ -166,6 +167,61 @@ func TestUpdateAllocation_PlaceholderTask_AssumePodFails(t
*testing.T) {
assert.Assert(t, assumePodErrorFound, "no event with reason
'AssumePodError' was recorded")
}
+// TestUpdateAllocation_ContextCanceled_ProcessesRemainingBatch verifies that a
+// canceled stopCtx aborts AssumePod retries without dropping the rest of the
+// response.New batch or rolling back the already-canceled allocation.
+func TestUpdateAllocation_ContextCanceled_ProcessesRemainingBatch(t
*testing.T) {
+ _, context := initCallbackTest(t, false, false)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ // register a second task under the same app so a single
UpdateAllocation
+ // batch can carry two `New` allocations; mirrors initCallbackTest's own
+ // pod-construction pattern for taskUID1 (Annotations-based app ID, no
+ // Spec.NodeName so it is unassigned and must go through AssumePod).
+ pod2 := &v1.Pod{
+ TypeMeta: apis.TypeMeta{
+ Kind: "Pod",
+ APIVersion: "v1",
+ },
+ ObjectMeta: apis.ObjectMeta{
+ Name: "yunikorn-test-00002",
+ UID: taskUID2,
+ Annotations: map[string]string{
+ constants.AnnotationApplicationID: appID,
+ },
+ },
+ Spec: v1.PodSpec{SchedulerName: "yunikorn"},
+ }
+ context.AddPod(pod2)
+ task2 := context.getTask(appID, taskUID2)
+ assert.Assert(t, task2 != nil)
+ task2.sm.SetState(TaskStates().Scheduling)
+
+ // swap in a callback whose stopCtx is already canceled, simulating
+ // KubernetesShim.Stop() having already run before this batch arrives.
+ cancelCtx, cancel := ctx.WithCancel(ctx.Background())
+ cancel()
+ callback := NewAsyncRMCallback(context, cancelCtx)
+
+ task1 := context.getTask(appID, taskUID1)
+ err := callback.UpdateAllocation(&si.AllocationResponse{
+ New: []*si.Allocation{
+ {ApplicationID: appID, AllocationKey: taskUID1, NodeID:
fakeNodeName},
+ {ApplicationID: appID, AllocationKey: taskUID2, NodeID:
fakeNodeName},
+ },
+ })
+ assert.NilError(t, err, "a canceled-context batch must not surface an
error to the RM proxy")
+ assert.Assert(t, !context.schedulerCache.IsAssumedPod(taskUID1))
+ assert.Assert(t, !context.schedulerCache.IsAssumedPod(taskUID2))
+ // setAllocationKey runs unconditionally as the first step of each loop
+ // iteration, before the retry/cancellation logic; task2 only has it set
+ // if the loop actually reached the second allocation.
+ assert.Equal(t, taskUID2, task2.GetAllocationKey(), "second allocation
in the batch must still be reached after the first hits a canceled context")
+ assert.Equal(t, TaskStates().Scheduling, task1.GetTaskState(),
"canceled allocation must not be rolled back")
+ assert.Equal(t, TaskStates().Scheduling, task2.GetTaskState(),
"canceled allocation must not be rolled back")
+}
+
func TestUpdateAllocation_NewTask_PodAlreadyAssigned(t *testing.T) {
callback, context := initCallbackTest(t, true, false)
defer dispatcher.UnregisterAllEventHandlers()
@@ -592,7 +648,7 @@ func initCallbackTest(t *testing.T, podAssigned,
placeholder bool) (*AsyncRMCall
dispatcher.Start()
dispatcher.RegisterEventHandler("TestAppHandler",
dispatcher.EventTypeApp, context.ApplicationEventHandler())
dispatcher.RegisterEventHandler("TestTaskHandler",
dispatcher.EventTypeTask, context.TaskEventHandler())
- callback := NewAsyncRMCallback(context)
+ callback := NewAsyncRMCallback(context, t.Context())
apiProvider.MockSchedulerAPIUpdateNodeFn(func(request *si.NodeRequest)
error {
for _, node := range request.Nodes {
dispatcher.Dispatch(CachedSchedulerNodeEvent{
diff --git a/pkg/common/test/volumebinder_mock.go
b/pkg/common/test/volumebinder_mock.go
index 607d24e4..a5c0cbe0 100644
--- a/pkg/common/test/volumebinder_mock.go
+++ b/pkg/common/test/volumebinder_mock.go
@@ -38,11 +38,12 @@ type VolumeBinderMock struct {
bindError error
conflictReasons volumebinding.ConflictReasons
- podVolumeClaim *volumebinding.PodVolumeClaims
- podVolumes *volumebinding.PodVolumes
- allBound bool
- revertCalledCount int
- bindCount atomic.Int32
+ podVolumeClaim *volumebinding.PodVolumeClaims
+ podVolumes *volumebinding.PodVolumes
+ allBound bool
+ revertCalledCount int
+ bindCount atomic.Int32
+ volumeClaimsCallCount atomic.Int32
}
func NewVolumeBinderMock() *VolumeBinderMock {
@@ -52,10 +53,10 @@ func NewVolumeBinderMock() *VolumeBinderMock {
}
func (v *VolumeBinderMock) GetPodVolumeClaims(_ klog.Logger, _ *v1.Pod)
(*volumebinding.PodVolumeClaims, error) {
+ v.volumeClaimsCallCount.Add(1)
if v.volumeClaimError != nil {
return nil, v.volumeClaimError
}
-
return v.podVolumeClaim, nil
}
@@ -104,6 +105,10 @@ func (v *VolumeBinderMock) GetBindCount() int32 {
return v.bindCount.Load()
}
+func (v *VolumeBinderMock) GetVolumeClaimsCallCount() int32 {
+ return v.volumeClaimsCallCount.Load()
+}
+
func (v *VolumeBinderMock) EnableVolumeClaimsError(message string) {
v.volumeClaimError = errors.New(message)
}
diff --git a/pkg/common/test/volumebinder_mock_test.go
b/pkg/common/test/volumebinder_mock_test.go
new file mode 100644
index 00000000..db9bce01
--- /dev/null
+++ b/pkg/common/test/volumebinder_mock_test.go
@@ -0,0 +1,39 @@
+/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you 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 test
+
+import (
+ "testing"
+
+ "gotest.tools/v3/assert"
+ "k8s.io/klog/v2"
+)
+
+func TestVolumeBinderMock_GetVolumeClaimsCallCount(t *testing.T) {
+ binder := NewVolumeBinderMock()
+ assert.Equal(t, int32(0), binder.GetVolumeClaimsCallCount())
+
+ _, err := binder.GetPodVolumeClaims(klog.NewKlogr(), nil)
+ assert.NilError(t, err)
+ assert.Equal(t, int32(1), binder.GetVolumeClaimsCallCount())
+
+ _, err = binder.GetPodVolumeClaims(klog.NewKlogr(), nil)
+ assert.NilError(t, err)
+ assert.Equal(t, int32(2), binder.GetVolumeClaimsCallCount(), "call
count must increment on every invocation")
+}
diff --git a/pkg/shim/scheduler.go b/pkg/shim/scheduler.go
index 9b8b5c9b..3d55070a 100644
--- a/pkg/shim/scheduler.go
+++ b/pkg/shim/scheduler.go
@@ -51,6 +51,7 @@ type KubernetesShim struct {
callback api.ResourceManagerCallback
stopChan chan struct{}
stopOnce sync.Once
+ cancelCallbacks ctx.CancelFunc
lock *locking.RWMutex
outstandingAppsFound bool
}
@@ -81,7 +82,7 @@ func NewShimScheduler(scheduler api.SchedulerAPI, configs
*conf.SchedulerConf, b
log.Log(log.Shim).Fatal("problem in creating the context")
return nil
}
- rmCallback := cache.NewAsyncRMCallback(context)
+ rmCallback, cancelCallbacks := newCallbackWithCancel(context)
eventBroadcaster := k8events.NewBroadcaster(&k8events.EventSinkImpl{
Interface: kubeClient.GetClientSet().EventsV1()})
@@ -94,17 +95,29 @@ func NewShimScheduler(scheduler api.SchedulerAPI, configs
*conf.SchedulerConf, b
events.SetRecorder(eventRecorder)
}
- return newShimSchedulerInternal(context, apiFactory, rmCallback)
+ return newShimSchedulerInternal(context, apiFactory, rmCallback,
cancelCallbacks)
}
-// this is visible for testing
-func newShimSchedulerInternal(ctx *cache.Context, apiFactory
client.APIProvider, cb api.ResourceManagerCallback) *KubernetesShim {
+// newCallbackWithCancel builds the async RM callback used by the shim together
+// with the CancelFunc that aborts any in-flight AssumePod retry when the shim
+// shuts down. Split out of NewShimScheduler so it can be tested without a live
+// Kubernetes client.
+func newCallbackWithCancel(cacheCtx *cache.Context)
(api.ResourceManagerCallback, ctx.CancelFunc) {
+ callbackCtx, cancelCallbacks := ctx.WithCancel(ctx.Background())
+ return cache.NewAsyncRMCallback(cacheCtx, callbackCtx), cancelCallbacks
+}
+
+func newShimSchedulerInternal(ctx *cache.Context, apiFactory
client.APIProvider, cb api.ResourceManagerCallback, cancelCallbacks
ctx.CancelFunc) *KubernetesShim {
+ if cancelCallbacks == nil {
+ cancelCallbacks = func() {}
+ }
ss := &KubernetesShim{
apiFactory: apiFactory,
context: ctx,
phManager:
cache.NewPlaceholderManager(apiFactory.GetAPIs()),
callback: cb,
stopChan: make(chan struct{}),
+ cancelCallbacks: cancelCallbacks,
lock: &locking.RWMutex{},
outstandingAppsFound: false,
}
@@ -231,6 +244,8 @@ func (ss *KubernetesShim) Stop() {
stopped = true
log.Log(log.ShimScheduler).Info("stopping scheduler")
close(ss.stopChan)
+ // stop the AssumePod retry loop running on the RM proxy
callback goroutine
+ ss.cancelCallbacks()
// stop the client library code that communicates with
Kubernetes
ss.apiFactory.Stop()
// stop the placeholder manager
diff --git a/pkg/shim/scheduler_mock_test.go b/pkg/shim/scheduler_mock_test.go
index 96f4024f..2aa0353f 100644
--- a/pkg/shim/scheduler_mock_test.go
+++ b/pkg/shim/scheduler_mock_test.go
@@ -19,6 +19,7 @@
package shim
import (
+ ctx "context"
"fmt"
"sync/atomic"
"testing"
@@ -66,9 +67,10 @@ func (fc *MockScheduler) init() {
mockedAPIProvider.GetAPIs().SchedulerAPI = fc.rmProxy
events.SetRecorder(events.NewMockedRecorder())
+ callbackCtx, cancelCallbacks := ctx.WithCancel(ctx.Background())
context := cache.NewContext(mockedAPIProvider)
- rmCallback := cache.NewAsyncRMCallback(context)
- ss := newShimSchedulerInternal(context, mockedAPIProvider, rmCallback)
+ rmCallback := cache.NewAsyncRMCallback(context, callbackCtx)
+ ss := newShimSchedulerInternal(context, mockedAPIProvider, rmCallback,
cancelCallbacks)
fc.context = context
fc.scheduler = ss
diff --git a/pkg/shim/scheduler_test.go b/pkg/shim/scheduler_test.go
index c3e482bf..d7a8cae5 100644
--- a/pkg/shim/scheduler_test.go
+++ b/pkg/shim/scheduler_test.go
@@ -173,7 +173,7 @@ func TestSchedulerRegistrationFailed(t *testing.T) {
})
ctx := cache.NewContext(mockedAPIProvider)
- shim := newShimSchedulerInternal(ctx, mockedAPIProvider, callback)
+ shim := newShimSchedulerInternal(ctx, mockedAPIProvider, callback,
func() {})
assert.Error(t, shim.Run(), "some error")
assertStopChannelClosed(t, shim)
@@ -183,7 +183,7 @@ func TestSchedulerRegistrationFailed(t *testing.T) {
func TestSchedulerStopClosesStopChannel(t *testing.T) {
mockedAPIProvider := client.NewMockedAPIProvider(false)
- shim := newShimSchedulerInternal(cache.NewContext(mockedAPIProvider),
mockedAPIProvider, nil)
+ shim := newShimSchedulerInternal(cache.NewContext(mockedAPIProvider),
mockedAPIProvider, nil, func() {})
shim.Stop()
assertStopChannelClosed(t, shim)
@@ -195,7 +195,7 @@ func TestSchedulerStopClosesStopChannel(t *testing.T) {
func TestSchedulerStopStopsAPIFactory(t *testing.T) {
mockedAPIProvider := client.NewMockedAPIProvider(false)
apiProvider := &trackingAPIProvider{APIProvider: mockedAPIProvider}
- shim := newShimSchedulerInternal(cache.NewContext(apiProvider),
apiProvider, nil)
+ shim := newShimSchedulerInternal(cache.NewContext(apiProvider),
apiProvider, nil, func() {})
shim.Stop()
assert.Check(t, apiProvider.stopped.Load(), "API provider should be
stopped with the scheduler")
@@ -203,6 +203,24 @@ func TestSchedulerStopStopsAPIFactory(t *testing.T) {
shim.Stop()
}
+func TestNewCallbackWithCancel(t *testing.T) {
+ mockedAPIProvider := client.NewMockedAPIProvider(false)
+ shimCtx := cache.NewContext(mockedAPIProvider)
+
+ callback, cancel := newCallbackWithCancel(shimCtx)
+ assert.Assert(t, callback != nil, "expected a non-nil callback")
+ assert.Assert(t, cancel != nil, "expected a non-nil cancel func")
+ cancel() // must be safe to invoke
+}
+
+func TestSchedulerNilCancelCallbacksFallsBackToNoop(t *testing.T) {
+ mockedAPIProvider := client.NewMockedAPIProvider(false)
+ shim := newShimSchedulerInternal(cache.NewContext(mockedAPIProvider),
mockedAPIProvider, nil, nil)
+
+ assert.Assert(t, shim.cancelCallbacks != nil, "nil cancelCallbacks must
fall back to a no-op")
+ shim.cancelCallbacks() // must not panic
+ shim.Stop() // Stop() invokes cancelCallbacks(); must not
panic either
+}
func assertStopChannelClosed(t *testing.T, shim *KubernetesShim) {
t.Helper()
select {
@@ -324,6 +342,53 @@ func TestAssumePodError(t *testing.T) {
"no pods should be bound when AssumePod always fails")
}
+// TestAssumePodError_AbortsOnShutdown verifies that Stop() interrupts the
+// AssumePod retry loop instead of letting it run out its full 30-step,
+// ~29s backoff after shutdown.
+func TestAssumePodError_AbortsOnShutdown(t *testing.T) {
+ cluster := MockScheduler{}
+ cluster.init()
+ binder := test.NewVolumeBinderMock()
+ binder.EnableVolumeClaimsError("unable to get volume claims")
+ cluster.apiProvider.SetVolumeBinder(binder)
+ assert.NilError(t, cluster.start(), "failed to start cluster")
+
+ err := cluster.updateConfig(configData, nil)
+ assert.NilError(t, err, "update config failed")
+ addNode(&cluster, "node-1")
+
+ taskResource := common.NewResourceBuilder().
+ AddResource(siCommon.Memory, 1000).
+ AddResource(siCommon.CPU, 1).
+ Build()
+ pod1 := createTestPod("root.a", "app0001", "task0001", taskResource)
+ cluster.AddPod(pod1)
+
+ // wait until the retry loop has made repeated attempts, proving it is
+ // actively blocking the RM callback goroutine on the 1s-per-step
backoff
+ err = utils.WaitForCondition(func() bool {
+ return binder.GetVolumeClaimsCallCount() >= 2
+ }, 100*time.Millisecond, 10*time.Second)
+ assert.NilError(t, err, "AssumePod retry never made repeated attempts")
+
+ countAtStop := binder.GetVolumeClaimsCallCount()
+ cluster.stop()
+
+ // at most the single attempt already in flight when Stop() ran may
still land
+ err = utils.WaitForCondition(func() bool {
+ return binder.GetVolumeClaimsCallCount() <= countAtStop+1
+ }, 50*time.Millisecond, time.Second)
+ assert.NilError(t, err, "AssumePod made more than one additional
attempt after shim.Stop()")
+
+ settledCount := binder.GetVolumeClaimsCallCount()
+ // the pre-fix backoff retries roughly once per second; wait more than a
+ // full step to prove the loop actually exited rather than merely being
+ // mid-sleep
+ time.Sleep(2 * time.Second)
+ assert.Equal(t, settledCount, binder.GetVolumeClaimsCallCount(),
+ "AssumePod kept retrying after shim.Stop(); the retry loop did
not abort on shutdown")
+}
+
func TestForeignPodTracking(t *testing.T) {
cluster := MockScheduler{}
cluster.init()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
