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-core.git
The following commit(s) were added to refs/heads/master by this push:
new 42edf052 [YUNIKORN-3323] Replacing google btree with tidwall btree
(#1117)
42edf052 is described below
commit 42edf052aa81bf71137179991721d1372f3704c0
Author: Aditya Maheshwari <[email protected]>
AuthorDate: Mon Aug 10 13:56:12 2026 +0530
[YUNIKORN-3323] Replacing google btree with tidwall btree (#1117)
Initiated tidwall btree with below options
1. With degree as 7
2. No locks
Used Walk() method as it is performant than Scan()
Closes: #1117
Signed-off-by: mani <[email protected]>
---
go.mod | 2 +-
go.sum | 4 ++--
pkg/scheduler/objects/application.go | 7 +++++-
pkg/scheduler/objects/node_collection.go | 34 +++++++++++++----------------
pkg/scheduler/objects/node_iterator.go | 15 +++++++------
pkg/scheduler/objects/node_iterator_test.go | 12 +++++-----
pkg/scheduler/objects/utilities_test.go | 8 +++----
7 files changed, 42 insertions(+), 40 deletions(-)
diff --git a/go.mod b/go.mod
index 76556ee2..2d3ee6f5 100644
--- a/go.mod
+++ b/go.mod
@@ -24,7 +24,6 @@ go 1.25.0
require (
github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260727092410-674338955bdf
github.com/go-ldap/ldap/v3 v3.4.13
- github.com/google/btree v1.1.3
github.com/google/go-cmp v0.7.0
github.com/google/uuid v1.6.0
github.com/julienschmidt/httprouter v1.3.0
@@ -33,6 +32,7 @@ require (
github.com/prometheus/client_model v0.6.2
github.com/prometheus/common v0.67.5
github.com/sasha-s/go-deadlock v0.3.9
+ github.com/tidwall/btree v1.8.1
go.uber.org/zap v1.27.1
go.yaml.in/yaml/v3 v3.0.4
golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90
diff --git a/go.sum b/go.sum
index 10e3c159..82d63445 100644
--- a/go.sum
+++ b/go.sum
@@ -20,8 +20,6 @@ github.com/go-logr/stdr v1.2.2
h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod
h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/golang/protobuf v1.5.4
h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
github.com/golang/protobuf v1.5.4/go.mod
h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
-github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg=
-github.com/google/btree v1.1.3/go.mod
h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod
h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
@@ -72,6 +70,8 @@ github.com/sasha-s/go-deadlock v0.3.9
h1:fiaT9rB7g5sr5ddNZvlwheclN9IP86eFW9WgqlE
github.com/sasha-s/go-deadlock v0.3.9/go.mod
h1:KuZj51ZFmx42q/mPaYbRk0P1xcwe697zsJKE03vD4/Y=
github.com/stretchr/testify v1.11.1
h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
github.com/stretchr/testify v1.11.1/go.mod
h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
+github.com/tidwall/btree v1.8.1 h1:27ehoXvm5AG/g+1VxLS1SD3vRhp/H7LuEfwNvddEdmA=
+github.com/tidwall/btree v1.8.1/go.mod
h1:jBbTdUWhSZClZWoDg54VnvV7/54modSOzDN7VXftj1A=
go.opentelemetry.io/auto/sdk v1.2.1
h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
go.opentelemetry.io/auto/sdk v1.2.1/go.mod
h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
go.opentelemetry.io/otel v1.43.0
h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I=
diff --git a/pkg/scheduler/objects/application.go
b/pkg/scheduler/objects/application.go
index a125e0ac..3d4d8130 100644
--- a/pkg/scheduler/objects/application.go
+++ b/pkg/scheduler/objects/application.go
@@ -2114,9 +2114,14 @@ func (sa *Application) RemoveAllAllocations()
[]*Allocation {
sa.allocatedResource = resources.NewResource()
sa.allocatedPlaceholder = resources.NewResource()
sa.allocations = make(map[string]*Allocation)
+
// When the resource trackers are zero we should not expect anything to
come in later.
if resources.IsZero(sa.pending) {
- if err := sa.HandleApplicationEvent(CompleteApplication); err
!= nil {
+ event := CompleteApplication
+ if sa.IsFailing() {
+ event = FailApplication
+ }
+ if err := sa.HandleApplicationEvent(event); err != nil {
log.Log(log.SchedApplication).Warn("Application state
not changed to Completing while removing all allocations",
zap.String("currentState", sa.CurrentState()),
zap.Error(err))
diff --git a/pkg/scheduler/objects/node_collection.go
b/pkg/scheduler/objects/node_collection.go
index cea1b2a6..90dfa464 100644
--- a/pkg/scheduler/objects/node_collection.go
+++ b/pkg/scheduler/objects/node_collection.go
@@ -21,7 +21,7 @@ package objects
import (
"fmt"
- "github.com/google/btree"
+ "github.com/tidwall/btree"
"go.uber.org/zap"
"github.com/apache/yunikorn-core/pkg/locking"
@@ -56,27 +56,23 @@ type nodeRef struct {
nodeScore float64 // node score
}
-func (nr nodeRef) Less(than btree.Item) bool {
- other, ok := than.(nodeRef)
- if !ok {
- return false
- }
- if nr.nodeScore < other.nodeScore {
+func nodeRefLess(a, b nodeRef) bool {
+ if a.nodeScore < b.nodeScore {
return true
}
- if other.nodeScore < nr.nodeScore {
+ if b.nodeScore < a.nodeScore {
return false
}
- return nr.node.NodeID < other.node.NodeID
+ return a.node.NodeID < b.node.NodeID
}
type baseNodeCollection struct {
Partition string // partition used with this collection
// Private fields need protection
- nsp NodeSortingPolicy // node sorting policy
- nodes map[string]*nodeRef // nodes assigned to this collection
- sortedNodes *btree.BTree // nodes sorted by score
+ nsp NodeSortingPolicy // node sorting policy
+ nodes map[string]*nodeRef // nodes assigned to this collection
+ sortedNodes *btree.BTreeG[nodeRef] // nodes sorted by score
unreservedIterator *treeIterator
fullIterator *treeIterator
@@ -109,7 +105,7 @@ func (nc *baseNodeCollection) AddNode(node *Node) error {
nodeScore: nc.scoreNode(node),
}
nc.nodes[node.NodeID] = &nref
- nc.sortedNodes.ReplaceOrInsert(nref)
+ nc.sortedNodes.Set(nref)
return nil
}
@@ -174,11 +170,11 @@ func (nc *baseNodeCollection) GetFullNodeIterator()
NodeIterator {
return nc.fullIterator
}
-func (nc *baseNodeCollection) cloneSortedNodes() *btree.BTree {
+func (nc *baseNodeCollection) cloneSortedNodes() *btree.BTreeG[nodeRef] {
nc.Lock()
defer nc.Unlock()
- return nc.sortedNodes.Clone()
+ return nc.sortedNodes.Copy()
}
// Sets the node sorting policy.
@@ -188,11 +184,11 @@ func (nc *baseNodeCollection) SetNodeSortingPolicy(policy
NodeSortingPolicy) {
nc.nsp = policy
// sortedNodes must be rebuilt since sort ordering is different
- nc.sortedNodes.Clear(false)
+ nc.sortedNodes.Clear()
for _, nref := range nc.nodes {
node := nref.node
nref.nodeScore = nc.scoreNode(node)
- nc.sortedNodes.ReplaceOrInsert(*nref)
+ nc.sortedNodes.Set(*nref)
}
}
@@ -217,7 +213,7 @@ func (nc *baseNodeCollection) NodeUpdated(node *Node) {
if nref.nodeScore != updatedScore {
nc.sortedNodes.Delete(*nref)
nref.nodeScore = nc.scoreNode(node)
- nc.sortedNodes.ReplaceOrInsert(*nref)
+ nc.sortedNodes.Set(*nref)
}
}
@@ -227,7 +223,7 @@ func NewNodeCollection(partition string) NodeCollection {
Partition: partition,
nsp:
NewNodeSortingPolicy(policies.FairSortPolicy.String(), nil),
nodes: make(map[string]*nodeRef),
- sortedNodes: btree.New(7), // Degree=7 here is experimentally
the most efficient for up to around 5k nodes
+ sortedNodes: btree.NewBTreeGOptions(nodeRefLess,
btree.Options{Degree: 7, NoLocks: true}), // Degree=7 here is experimentally
the most efficient for up to around 5k nodes
}
unreservedIterator := NewTreeIterator(acceptUnreserved,
bsc.cloneSortedNodes)
diff --git a/pkg/scheduler/objects/node_iterator.go
b/pkg/scheduler/objects/node_iterator.go
index 1bb7b361..9e7f321c 100644
--- a/pkg/scheduler/objects/node_iterator.go
+++ b/pkg/scheduler/objects/node_iterator.go
@@ -19,7 +19,7 @@
package objects
import (
- "github.com/google/btree"
+ "github.com/tidwall/btree"
)
// NodeIterator iterates over a list of nodes based on the defined policy
@@ -30,25 +30,26 @@ type NodeIterator interface {
type treeIterator struct {
accept func(*Node) bool
- getTree func() *btree.BTree
+ getTree func() *btree.BTreeG[nodeRef]
}
// ForEachNode Calls the provided "f" function on the sorted Node object until
it returns false.
// The accept() function checks if the node should be a candidate or not.
func (ti *treeIterator) ForEachNode(f func(*Node) bool) {
- ti.getTree().Ascend(func(item btree.Item) bool {
- if ref, ok := item.(nodeRef); ok {
+ ti.getTree().Walk(func(refs []nodeRef) bool {
+ for _, ref := range refs {
node := ref.node
if ti.accept(node) {
- return f(node)
+ if !f(node) {
+ return false
+ }
}
}
-
return true
})
}
-func NewTreeIterator(accept func(*Node) bool, getTree func() *btree.BTree)
*treeIterator {
+func NewTreeIterator(accept func(*Node) bool, getTree func()
*btree.BTreeG[nodeRef]) *treeIterator {
ti := &treeIterator{
getTree: getTree,
accept: accept,
diff --git a/pkg/scheduler/objects/node_iterator_test.go
b/pkg/scheduler/objects/node_iterator_test.go
index 56594c6d..b983b52b 100644
--- a/pkg/scheduler/objects/node_iterator_test.go
+++ b/pkg/scheduler/objects/node_iterator_test.go
@@ -22,7 +22,7 @@ import (
"strconv"
"testing"
- "github.com/google/btree"
+ "github.com/tidwall/btree"
"gotest.tools/v3/assert"
"github.com/apache/yunikorn-core/pkg/common/resources"
@@ -30,7 +30,7 @@ import (
func TestTreeIterator_AcceptAll(t *testing.T) {
tree := getTree()
- treeItr := NewTreeIterator(acceptAll, func() *btree.BTree {
+ treeItr := NewTreeIterator(acceptAll, func() *btree.BTreeG[nodeRef] {
return tree
})
@@ -45,7 +45,7 @@ func TestTreeIterator_AcceptAll(t *testing.T) {
func TestTreeIterator_AcceptUnreserved(t *testing.T) {
tree := getTree()
- treeItr := NewTreeIterator(acceptUnreserved, func() *btree.BTree {
+ treeItr := NewTreeIterator(acceptUnreserved, func()
*btree.BTreeG[nodeRef] {
return tree
})
@@ -65,14 +65,14 @@ func TestTreeIterator_AcceptUnreserved(t *testing.T) {
}
}
-func getTree() *btree.BTree {
+func getTree() *btree.BTreeG[nodeRef] {
nodesReserved := newSchedNodeList(0, 5, true)
nodes := newSchedNodeList(5, 10, false)
nodes = append(nodes, nodesReserved...)
- tree := btree.New(7)
+ tree := btree.NewBTreeG(nodeRefLess)
for _, n := range nodes {
- tree.ReplaceOrInsert(nodeRef{
+ tree.Set(nodeRef{
node: n,
nodeScore: 1,
})
diff --git a/pkg/scheduler/objects/utilities_test.go
b/pkg/scheduler/objects/utilities_test.go
index 255f6bc5..7c2cee80 100644
--- a/pkg/scheduler/objects/utilities_test.go
+++ b/pkg/scheduler/objects/utilities_test.go
@@ -23,7 +23,7 @@ import (
"testing"
"time"
- "github.com/google/btree"
+ "github.com/tidwall/btree"
"gotest.tools/v3/assert"
"github.com/apache/yunikorn-core/pkg/common/configs"
@@ -331,15 +331,15 @@ func assertAllocationLog(t *testing.T, ask *Allocation,
message []string) {
}
func getNodeIteratorFn(nodes ...*Node) func() NodeIterator {
- tree := btree.New(7)
+ tree := btree.NewBTreeG(nodeRefLess)
for _, node := range nodes {
- tree.ReplaceOrInsert(nodeRef{
+ tree.Set(nodeRef{
node, 1,
})
}
return func() NodeIterator {
- return NewTreeIterator(acceptAll, func() *btree.BTree {
+ return NewTreeIterator(acceptAll, func() *btree.BTreeG[nodeRef]
{
return tree
})
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
