This is an automated email from the ASF dual-hosted git repository.
HuangTing-Yao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/yunikorn-core.git
The following commit(s) were added to refs/heads/master by this push:
new 403b7406 [YUNIKORN-3365] Drain RM events on shutdown and buffer
notification reply channels to prevent goroutine leak (#1152)
403b7406 is described below
commit 403b7406ae7a2e3f34f9d78a75e65ef259c6cf64
Author: PoiBlackTea <[email protected]>
AuthorDate: Mon Sep 7 21:01:56 2026 +0800
[YUNIKORN-3365] Drain RM events on shutdown and buffer notification reply
channels to prevent goroutine leak (#1152)
Closes: #1152
Signed-off-by: HuangTing-Yao <[email protected]>
---
pkg/rmproxy/rmproxy.go | 29 ++++++++++++
pkg/rmproxy/rmproxy_test.go | 85 ++++++++++++++++++++++++++++++++++++
pkg/scheduler/context.go | 4 +-
pkg/scheduler/objects/application.go | 2 +-
4 files changed, 117 insertions(+), 3 deletions(-)
diff --git a/pkg/rmproxy/rmproxy.go b/pkg/rmproxy/rmproxy.go
index fa29a097..80a2ba09 100644
--- a/pkg/rmproxy/rmproxy.go
+++ b/pkg/rmproxy/rmproxy.go
@@ -203,11 +203,40 @@ func (rmp *RMProxy) handleRMEvents() {
panic(fmt.Sprintf("%s is not an acceptable type
for RM event.", reflect.TypeOf(v).String()))
}
case <-rmp.stop:
+ rmp.drainPendingEvents()
return
}
}
}
+func (rmp *RMProxy) drainPendingEvents() {
+ for {
+ select {
+ case ev := <-rmp.pendingRMEvents:
+ switch v := ev.(type) {
+ case *rmevent.RMNewAllocationsEvent:
+ drainReplyChannel(v.Channel)
+ case *rmevent.RMReleaseAllocationEvent:
+ drainReplyChannel(v.Channel)
+ }
+ default:
+ return
+ }
+ }
+}
+
+func drainReplyChannel(ch chan *rmevent.Result) {
+ if ch != nil {
+ select {
+ case ch <- &rmevent.Result{
+ Succeeded: false,
+ Reason: "RMProxy is stopping",
+ }:
+ default:
+ }
+ }
+}
+
func (rmp *RMProxy) RegisterResourceManager(request
*si.RegisterResourceManagerRequest, callback api.ResourceManagerCallback)
(*si.RegisterResourceManagerResponse, error) {
rmp.Lock()
defer rmp.Unlock()
diff --git a/pkg/rmproxy/rmproxy_test.go b/pkg/rmproxy/rmproxy_test.go
new file mode 100644
index 00000000..9c98526f
--- /dev/null
+++ b/pkg/rmproxy/rmproxy_test.go
@@ -0,0 +1,85 @@
+/*
+ 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 rmproxy
+
+import (
+ "testing"
+ "time"
+
+ "gotest.tools/v3/assert"
+
+ "github.com/apache/yunikorn-core/pkg/rmproxy/rmevent"
+ "github.com/apache/yunikorn-scheduler-interface/lib/go/si"
+)
+
+func TestRMProxy_StopUnblocksWaitingCaller(t *testing.T) {
+ rmp := NewRMProxy(nil)
+ rmp.StartService()
+
+ c := make(chan *rmevent.Result, 1)
+ rmp.HandleEvent(&rmevent.RMReleaseAllocationEvent{
+ ReleasedAllocations: []*si.AllocationRelease{},
+ RmID: "rm-test",
+ Channel: c,
+ })
+
+ rmp.Stop() // exercises the case <-rmp.stop: drainPendingEvents() wiring
+
+ select {
+ case res := <-c:
+ assert.Assert(t, res != nil) // normal reply or drained —
either is fine, point is no leak
+ case <-time.After(time.Second):
+ t.Fatal("caller leaked: no reply after Stop")
+ }
+}
+
+func TestRMProxy_DrainPendingEvents(t *testing.T) {
+ rmp := NewRMProxy(nil)
+
+ allocResultCh := make(chan *rmevent.Result, 1)
+ releaseResultCh := make(chan *rmevent.Result, 1)
+
+ rmp.HandleEvent(&rmevent.RMNewAllocationsEvent{
+ Allocations: []*si.Allocation{},
+ RmID: "rm-test",
+ Channel: allocResultCh,
+ })
+ rmp.HandleEvent(&rmevent.RMReleaseAllocationEvent{
+ ReleasedAllocations: []*si.AllocationRelease{},
+ RmID: "rm-test",
+ Channel: releaseResultCh,
+ })
+
+ rmp.drainPendingEvents()
+
+ assertDrainFailedResult(t, allocResultCh, "allocResultCh")
+ assertDrainFailedResult(t, releaseResultCh, "releaseResultCh")
+}
+
+func assertDrainFailedResult(t *testing.T, ch <-chan *rmevent.Result, name
string) {
+ t.Helper()
+ select {
+ case res := <-ch:
+ assert.Assert(t, res != nil, "expected non-nil response on %s",
name)
+ assert.Assert(t, !res.Succeeded, "expected Succeeded to be
false on %s", name)
+ assert.Equal(t, res.Reason, "RMProxy is stopping")
+ case <-time.After(1 * time.Second):
+ t.Fatalf("timed out waiting for response on %s", name)
+ }
+}
diff --git a/pkg/scheduler/context.go b/pkg/scheduler/context.go
index de36655b..c5ef60e1 100644
--- a/pkg/scheduler/context.go
+++ b/pkg/scheduler/context.go
@@ -795,7 +795,7 @@ func (cc *ClusterContext)
processAllocationReleases(releases []*si.AllocationRel
// Create a RM update event to notify RM of new allocations
// Lock free call, all updates occur via events.
func (cc *ClusterContext) notifyRMNewAllocation(rmID string, alloc
*objects.Allocation) {
- c := make(chan *rmevent.Result)
+ c := make(chan *rmevent.Result, 1)
// communicate the allocation to the RM synchronously
cc.rmEventHandler.HandleEvent(&rmevent.RMNewAllocationsEvent{
Allocations: []*si.Allocation{alloc.NewSIFromAllocation()},
@@ -815,7 +815,7 @@ func (cc *ClusterContext) notifyRMNewAllocation(rmID
string, alloc *objects.Allo
// Create a RM update event to notify RM of released allocations
// Lock free call, all updates occur via events.
func (cc *ClusterContext) notifyRMAllocationReleased(rmID string,
partitionName string, released []*objects.Allocation, terminationType
si.TerminationType, message string) {
- c := make(chan *rmevent.Result)
+ c := make(chan *rmevent.Result, 1)
releaseEvent := &rmevent.RMReleaseAllocationEvent{
ReleasedAllocations: make([]*si.AllocationRelease, 0),
RmID: rmID,
diff --git a/pkg/scheduler/objects/application.go
b/pkg/scheduler/objects/application.go
index 0fce920b..189d1fe2 100644
--- a/pkg/scheduler/objects/application.go
+++ b/pkg/scheduler/objects/application.go
@@ -2342,7 +2342,7 @@ func (sa *Application)
notifyRMAllocationReleased(released []*Allocation, termin
if len(released) == 0 || sa.rmEventHandler == nil {
return
}
- c := make(chan *rmevent.Result)
+ c := make(chan *rmevent.Result, 1)
releaseEvent := &rmevent.RMReleaseAllocationEvent{
ReleasedAllocations: make([]*si.AllocationRelease, 0),
RmID: sa.rmID,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
