koordlet: revise core sched implementation

Signed-off-by: saintube <[email protected]>

go.mod

@@ -36,7 +36,6 @@ require (
 	go.uber.org/atomic v1.10.0
 	go.uber.org/multierr v1.6.0
 	golang.org/x/crypto v0.14.0
-	golang.org/x/exp v0.0.0-20231006140011-7918f672742d
 	golang.org/x/net v0.16.0
 	golang.org/x/sys v0.13.0
 	golang.org/x/time v0.0.0-20220920022843-2ce7c2934d45

go.sum

@@ -1302,8 +1302,6 @@ golang.org/x/exp v0.0.0-20200119233911-0405dc783f0a/go.mod h1:2RIsYlXP63K8oxa1u0
 golang.org/x/exp v0.0.0-20200207192155-f17229e696bd/go.mod h1:J/WKrq2StrnmMY6+EHIKF9dgMWnmCNThgcyBT1FY9mM=
 golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU=
 golang.org/x/exp v0.0.0-20210220032938-85be41e4509f/go.mod h1:I6l2HNBLBZEcrOoCpyKLdY2lHoRZ8lI4x60KMCQDft4=
-golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
-golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
 golang.org/x/image v0.0.0-20180708004352-c73c2afc3b81/go.mod h1:ux5Hcp/YLpHSI86hEcLt0YII63i6oz57MZXIpbrjZUs=
 golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js=
 golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0=

pkg/koordlet/runtimehooks/hooks/coresched/cookie_cache.go

@@ -20,74 +20,76 @@ import (
 	"sort"
 	"sync"
 
-	"golang.org/x/exp/constraints"
-
 	sysutil "github.com/koordinator-sh/koordinator/pkg/koordlet/util/system"
 )
 
 // CookieCacheEntry is an entry which stores the cookie ID and its belonging PIDs.
 type CookieCacheEntry struct {
 	rwMutex  sync.RWMutex
-	CookieID uint64
-	pidCache *OrderedMap[uint32]
+	cookieID uint64
+	pidCache *PIDCache
 }
 
 func newCookieCacheEntry(cookieID uint64, pids ...uint32) *CookieCacheEntry {
+	m := &PIDCache{}
+	m.AddAny(pids...)
 	return &CookieCacheEntry{
-		CookieID: cookieID,
-		pidCache: newUint32OrderedMap(pids...),
+		cookieID: cookieID,
+		pidCache: m,
 	}
 }
 
 func (c *CookieCacheEntry) DeepCopy() *CookieCacheEntry {
+	c.rwMutex.RLock()
+	defer c.rwMutex.RUnlock()
+	copiedM := c.pidCache.DeepCopy()
 	return &CookieCacheEntry{
-		CookieID: c.CookieID,
-		pidCache: c.pidCache.DeepCopy(),
+		cookieID: c.cookieID,
+		pidCache: copiedM,
 	}
 }
 
 func (c *CookieCacheEntry) GetCookieID() uint64 {
 	c.rwMutex.RLock()
 	defer c.rwMutex.RUnlock()
-	return c.CookieID
+	return c.cookieID
 }
 
 func (c *CookieCacheEntry) SetCookieID(cookieID uint64) {
 	c.rwMutex.Lock()
 	defer c.rwMutex.Unlock()
-	c.CookieID = cookieID
+	c.cookieID = cookieID
 }
 
 func (c *CookieCacheEntry) IsEntryInvalid() bool {
 	c.rwMutex.RLock()
 	defer c.rwMutex.RUnlock()
-	return c.CookieID <= sysutil.DefaultCoreSchedCookieID || c.pidCache.Len() <= 0
+	return c.cookieID <= sysutil.DefaultCoreSchedCookieID || c.pidCache.Len() <= 0
 }
 
 func (c *CookieCacheEntry) HasPID(pid uint32) bool {
 	c.rwMutex.RLock()
 	defer c.rwMutex.RUnlock()
-	_, ok := c.pidCache.Get(pid)
-	return ok
+	return c.pidCache.Has(pid)
 }
 
 func (c *CookieCacheEntry) ContainsPIDs(pids ...uint32) []uint32 {
 	c.rwMutex.RLock()
 	defer c.rwMutex.RUnlock()
-	var notFoundpids []uint32
+	var notFoundPIDs []uint32
 	for _, pid := range pids {
-		_, ok := c.pidCache.Get(pid)
-		if !ok {
-			notFoundpids = append(notFoundpids, pid)
+		if !c.pidCache.Has(pid) {
+			notFoundPIDs = append(notFoundPIDs, pid)
 		}
 	}
-	return notFoundpids
+	return notFoundPIDs
 }
 
+// GetAllPIDs gets all PIDs sorted in ascending order.
 func (c *CookieCacheEntry) GetAllPIDs() []uint32 {
 	c.rwMutex.RLock()
 	defer c.rwMutex.RUnlock()
-	return c.pidCache.GetAll()
+	return c.pidCache.GetAllSorted()
 }
 
 func (c *CookieCacheEntry) AddPIDs(pids ...uint32) {
@@ -108,166 +110,55 @@ func (c *CookieCacheEntry) DeletePIDs(pids ...uint32) {
 	c.pidCache.DeleteAny(pids...)
 }
 
-type IndexedValue[T constraints.Ordered] struct {
-	Value T
-	Index int
-}
-
-const MaxSimpleRecreateSize = 4
+type PIDCache map[uint32]struct{}
 
-// OrderedMap is an ordered map with type T.
-type OrderedMap[T constraints.Ordered] struct {
-	Set  map[T]*IndexedValue[T]
-	List []T // ascending order
+func NewPIDCache(pids ...uint32) *PIDCache {
+	p := &PIDCache{}
+	p.AddAny(pids...)
+	return p
 }
 
-func newUint32OrderedMap(vs ...uint32) *OrderedMap[uint32] {
-	return newOrderedMap[uint32](vs...)
-}
-
-func newOrderedMap[T constraints.Ordered](vs ...T) *OrderedMap[T] {
-	s := &OrderedMap[T]{
-		Set:  map[T]*IndexedValue[T]{},
-		List: make([]T, 0),
-	}
-	if len(vs) > 0 {
-		s.AddAny(vs...)
+func (p PIDCache) DeepCopy() *PIDCache {
+	copiedM := map[uint32]struct{}{}
+	for pid := range p {
+		copiedM[pid] = struct{}{}
 	}
-	return s
+	return (*PIDCache)(&copiedM)
 }
 
-func (o *OrderedMap[T]) DeepCopy() *OrderedMap[T] {
-	u := &OrderedMap[T]{
-		Set:  map[T]*IndexedValue[T]{},
-		List: make([]T, len(o.List)),
-	}
-	for i := range o.List {
-		v := o.List[i]
-		u.List[i] = v
-		u.Set[v] = &IndexedValue[T]{
-			Value: v,
-			Index: i,
-		}
-	}
-	return u
+func (p PIDCache) Len() int {
+	return len(p)
 }
 
-func (o *OrderedMap[T]) Len() int {
-	return len(o.List)
-}
-
-func (o *OrderedMap[T]) Get(v T) (T, bool) {
-	e, ok := o.Set[v]
-	if ok {
-		return e.Value, true
-	}
-	var oo T
-	return oo, false
+func (p PIDCache) Has(pid uint32) bool {
+	_, ok := p[pid]
+	return ok
 }
 
-func (o *OrderedMap[T]) GetAll() []T {
-	if o.Len() <= 0 {
+func (p PIDCache) GetAllSorted() []uint32 {
+	if len(p) <= 0 {
 		return nil
 	}
-	l := make([]T, len(o.List))
-	copy(l, o.List)
-	return l
-}
-
-func (o *OrderedMap[T]) Add(v T) {
-	if _, ok := o.Set[v]; ok {
-		return
-	}
-
-	o.Set[v] = &IndexedValue[T]{
-		Value: v,
+	pids := make([]uint32, len(p))
+	i := 0
+	for pid := range p {
+		pids[i] = pid
+		i++
 	}
-	n := len(o.List)
-	idx := sort.Search(n, func(i int) bool {
-		return o.List[i] > v
+	sort.Slice(pids, func(i, j int) bool {
+		return pids[i] < pids[j]
 	})
-
-	if idx >= n {
-		o.List = append(o.List, v)
-		o.Set[v].Index = n
-		return
-	}
-
-	var e T
-	o.List = append(o.List, e)
-	copy(o.List[idx+1:], o.List[idx:n])
-	o.List[idx] = v
-	o.Set[v].Index = idx
-	for i, l := range o.List[idx+1:] {
-		o.Set[l].Index = i + idx + 1
-	}
-}
-
-func (o *OrderedMap[T]) AddAny(vs ...T) {
-	if n := len(o.List); n > MaxSimpleRecreateSize && len(vs)*2 <= n {
-		for _, v := range vs {
-			o.Add(v)
-		}
-		return
-	}
-
-	for _, v := range vs {
-		if _, ok := o.Set[v]; ok {
-			continue
-		}
-		o.Set[v] = &IndexedValue[T]{
-			Value: v,
-		}
-		o.List = append(o.List, v)
-	}
-	sort.Slice(o.List, func(i, j int) bool {
-		return o.List[i] < o.List[j]
-	})
-	for i, v := range o.List {
-		o.Set[v].Index = i
-	}
+	return pids
 }
 
-func (o *OrderedMap[T]) Delete(v T) bool {
-	e, ok := o.Set[v]
-	if !ok {
-		return false
-	}
-	idx := e.Index
-	delete(o.Set, v)
-	if idx <= 0 {
-		o.List = o.List[idx+1:]
-	} else {
-		o.List = append(o.List[:idx], o.List[idx+1:]...)
-	}
-	for i, l := range o.List[idx:] {
-		o.Set[l].Index = i + idx
+func (p PIDCache) AddAny(pids ...uint32) {
+	for _, pid := range pids {
+		p[pid] = struct{}{}
 	}
-	return true
 }
 
-func (o *OrderedMap[T]) DeleteAny(vs ...T) {
-	if n := len(o.List); n > MaxSimpleRecreateSize && len(vs)*2 <= n {
-		for _, v := range vs {
-			o.Delete(v)
-		}
-		return
-	}
-
-	for _, v := range vs {
-		if _, ok := o.Set[v]; ok {
-			delete(o.Set, v)
-		}
-	}
-	l := make([]T, 0)
-	for v := range o.Set {
-		l = append(l, v)
-	}
-	sort.Slice(l, func(i, j int) bool {
-		return l[i] < l[j]
-	})
-	o.List = l
-	for i, v := range o.List {
-		o.Set[v].Index = i
+func (p PIDCache) DeleteAny(pids ...uint32) {
+	for _, pid := range pids {
+		delete(p, pid)
 	}
 }

pkg/koordlet/runtimehooks/hooks/coresched/cookie_cache_test.go

@@ -21,7 +21,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"golang.org/x/exp/constraints"
 )
 
 func TestCookieCacheEntry(t *testing.T) {
@@ -251,56 +250,6 @@ func TestCookieCacheEntry(t *testing.T) {
 	}
 }
 
-func TestOrderedMap(t *testing.T) {
-	tests := []struct {
-		name   string
-		testFn func(t *testing.T)
-	}{
-		{
-			name:   "test uint32",
-			testFn: testOrderedMap[uint32],
-		},
-		{
-			name:   "test int",
-			testFn: testOrderedMap[int],
-		},
-		{
-			name:   "test float64",
-			testFn: testOrderedMap[float64],
-		},
-		{
-			name:   "test string",
-			testFn: testOrderedMap[string],
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			tt.testFn(t)
-		})
-	}
-}
-
-func testOrderedMap[T constraints.Ordered](t *testing.T) {
-	m := newOrderedMap[T]()
-	assert.NotNil(t, m)
-
-	var e T
-	assert.Equal(t, 0, m.Len())
-	_, gotExist := m.Get(e)
-	assert.False(t, gotExist)
-	m.AddAny(e)
-	assert.Equal(t, 1, m.Len())
-	got, gotExist := m.Get(e)
-	assert.True(t, gotExist)
-	assert.Equal(t, e, got)
-	m.DeleteAny(e)
-	assert.Equal(t, 0, m.Len())
-	_, gotExist = m.Get(e)
-	assert.False(t, gotExist)
-	mc := m.DeepCopy()
-	assert.Equal(t, m, mc)
-}
-
 func testGetOutOfOrderUint32Slice(ss []uint32) []uint32 {
 	s := make([]uint32, len(ss))
 	copy(s, ss)

pkg/koordlet/runtimehooks/hooks/coresched/core_sched_test.go

@@ -1007,7 +1007,7 @@ func TestPlugin_SetContainerCookie(t *testing.T) {
 				assert.True(t, ok)
 				entry, ok := entryIf.(*CookieCacheEntry)
 				assert.True(t, ok)
-				assert.Equal(t, cookieID, entry.CookieID)
+				assert.Equal(t, cookieID, entry.GetCookieID())
 				assert.Equal(t, len(tt.wantFields.cookieToPIDs[cookieID]), len(entry.GetAllPIDs()),
 					"expect [%v] but got [%v]", tt.wantFields.cookieToPIDs[cookieID], entry.GetAllPIDs())
 				for _, pid := range tt.wantFields.cookieToPIDs[cookieID] {
@@ -1829,7 +1829,7 @@ func TestPlugin_LoadAllCookies(t *testing.T) {
 				assert.True(t, ok)
 				entry, ok := entryIf.(*CookieCacheEntry)
 				assert.True(t, ok)
-				assert.Equal(t, cookieID, entry.CookieID)
+				assert.Equal(t, cookieID, entry.GetCookieID())
 				assert.Equal(t, len(tt.wantFields.cookieToPIDs[cookieID]), len(entry.GetAllPIDs()),
 					"expect [%v] but got [%v]", tt.wantFields.cookieToPIDs[cookieID], entry.GetAllPIDs())
 				for _, pid := range tt.wantFields.cookieToPIDs[cookieID] {