fix: use the object heap to keep the min ddl ts order

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

Author: SimFG

Diff for: internal/rootcoord/scheduler.go

@@ -28,6 +28,7 @@ import (
 	"github.com/milvus-io/milvus/internal/tso"
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/lock"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 type IScheduler interface {
@@ -46,6 +47,7 @@ type scheduler struct {
 	tsoAllocator tso.Allocator
 
 	taskChan chan task
+	taskHeap typeutil.Heap[task]
 
 	lock sync.Mutex
 
@@ -56,16 +58,22 @@ type scheduler struct {
 	lockMapping    map[LockLevel]*lock.KeyLock[string]
 }
 
+func GetTaskHeapOrder(t task) Timestamp {
+	return t.GetTs()
+}
+
 func newScheduler(ctx context.Context, idAllocator allocator.Interface, tsoAllocator tso.Allocator) *scheduler {
 	ctx1, cancel := context.WithCancel(ctx)
 	// TODO
 	n := 1024 * 10
+	taskArr := make([]task, 0)
 	s := &scheduler{
 		ctx:            ctx1,
 		cancel:         cancel,
 		idAllocator:    idAllocator,
 		tsoAllocator:   tsoAllocator,
 		taskChan:       make(chan task, n),
+		taskHeap:       typeutil.NewObjectArrayBasedMinimumHeap[task, Timestamp](taskArr, GetTaskHeapOrder),
 		minDdlTs:       *atomic.NewUint64(0),
 		clusterLock:    lock.NewKeyLock[string](),
 		databaseLock:   lock.NewKeyLock[string](),
@@ -93,7 +101,7 @@ func (s *scheduler) Stop() {
 }
 
 func (s *scheduler) execute(task task) {
-	defer s.setMinDdlTs(task.GetTs()) // we should update ts, whatever task succeeds or not.
+	defer s.setMinDdlTs() // we should update ts, whatever task succeeds or not.
 	task.SetInQueueDuration()
 	if err := task.Prepare(task.GetCtx()); err != nil {
 		task.NotifyDone(err)
@@ -153,6 +161,7 @@ func (s *scheduler) setTs(task task) error {
 		return err
 	}
 	task.SetTs(ts)
+	s.taskHeap.Push(task)
 	return nil
 }
 
@@ -186,18 +195,29 @@ func (s *scheduler) GetMinDdlTs() Timestamp {
 	return s.minDdlTs.Load()
 }
 
-func (s *scheduler) setMinDdlTs(ts Timestamp) {
-	s.minDdlTs.Store(ts)
+func (s *scheduler) setMinDdlTs() {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	header := s.taskHeap.Peek()
+	for header != nil && header.IsFinished() {
+		s.minDdlTs.Store(header.GetTs())
+		s.taskHeap.Pop()
+		header = s.taskHeap.Peek()
+	}
 }
 
 func (s *scheduler) executeTaskWithLock(task task, lockerKey LockerKey) error {
 	if lockerKey == nil {
 		if err := s.setID(task); err != nil {
 			return err
 		}
+		s.lock.Lock()
 		if err := s.setTs(task); err != nil {
+			s.lock.Unlock()
 			return err
 		}
+		s.lock.Unlock()
 		s.execute(task)
 		return nil
 	}

Diff for: internal/rootcoord/task.go

@@ -19,6 +19,7 @@ package rootcoord
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	"go.uber.org/zap"
@@ -53,16 +54,18 @@ type task interface {
 	Execute(ctx context.Context) error
 	WaitToFinish() error
 	NotifyDone(err error)
+	IsFinished() bool
 	SetInQueueDuration()
 	GetLockerKey() LockerKey
 }
 
 type baseTask struct {
-	ctx  context.Context
-	core *Core
-	done chan error
-	ts   Timestamp
-	id   UniqueID
+	ctx        context.Context
+	core       *Core
+	done       chan error
+	isFinished atomic.Bool
+	ts         Timestamp
+	id         UniqueID
 
 	tr       *timerecord.TimeRecorder
 	queueDur time.Duration
@@ -116,12 +119,17 @@ func (b *baseTask) WaitToFinish() error {
 
 func (b *baseTask) NotifyDone(err error) {
 	b.done <- err
+	b.isFinished.Store(true)
 }
 
 func (b *baseTask) SetInQueueDuration() {
 	b.queueDur = b.tr.ElapseSpan()
 }
 
+func (b *baseTask) IsFinished() bool {
+	return b.isFinished.Load()
+}
+
 func (b *baseTask) GetLockerKey() LockerKey {
 	return nil
 }

Diff for: pkg/util/typeutil/heap.go

@@ -69,6 +69,39 @@ func (h *heapArray[E]) Peek() interface{} {
 	return (*h)[0]
 }
 
+type objectHeapArray[O any, E constraints.Ordered] struct {
+	objects      []O
+	getOrderFunc func(O) E
+}
+
+func (h *objectHeapArray[O, E]) Len() int {
+	return len(h.objects)
+}
+
+func (h *objectHeapArray[O, E]) Less(i, j int) bool {
+	return h.getOrderFunc(h.objects[i]) < h.getOrderFunc(h.objects[j])
+}
+
+func (h *objectHeapArray[O, E]) Swap(i, j int) {
+	h.objects[i], h.objects[j] = h.objects[j], h.objects[i]
+}
+
+func (h *objectHeapArray[O, E]) Push(x interface{}) {
+	h.objects = append(h.objects, x.(O))
+}
+
+func (h *objectHeapArray[O, E]) Pop() interface{} {
+	old := h.objects
+	n := len(old)
+	x := old[n-1]
+	h.objects = old[0 : n-1]
+	return x
+}
+
+func (h *objectHeapArray[O, E]) Peek() interface{} {
+	return h.objects[0]
+}
+
 // reverseOrderedInterface is a heap base interface that reverses the order of the elements.
 type reverseOrderedInterface[E constraints.Ordered] struct {
 	HeapInterface
@@ -107,6 +140,37 @@ func NewArrayBasedMinimumHeap[E constraints.Ordered](initial []E) Heap[E] {
 	}
 }
 
+func NewObjectArrayBasedMaximumHeap[O any, E constraints.Ordered](initial []O, getOrderFunc func(O) E) Heap[O] {
+	if initial == nil {
+		initial = make([]O, 0)
+	}
+	ha := &objectHeapArray[O, E]{
+		objects:      initial,
+		getOrderFunc: getOrderFunc,
+	}
+	reverse := reverseOrderedInterface[E]{
+		HeapInterface: ha,
+	}
+	heap.Init(reverse)
+	return &heapImpl[O, reverseOrderedInterface[E]]{
+		inner: reverse,
+	}
+}
+
+func NewObjectArrayBasedMinimumHeap[O any, E constraints.Ordered](initial []O, getOrderFunc func(O) E) Heap[O] {
+	if initial == nil {
+		initial = make([]O, 0)
+	}
+	ha := &objectHeapArray[O, E]{
+		objects:      initial,
+		getOrderFunc: getOrderFunc,
+	}
+	heap.Init(ha)
+	return &heapImpl[O, *objectHeapArray[O, E]]{
+		inner: ha,
+	}
+}
+
 // heapImpl is a min-heap of E.
 type heapImpl[E any, H HeapInterface] struct {
 	inner H

Diff for: pkg/util/typeutil/heap_test.go

@@ -39,3 +39,53 @@ func TestMaximumHeap(t *testing.T) {
 		assert.Equal(t, i, heap.Pop())
 	}
 }
+
+type FooHeapObject struct {
+	value int
+}
+
+func GetFooHeapObjectOrderFunc(obj *FooHeapObject) int {
+	return obj.value
+}
+
+func TestMinimumObjectHeap(t *testing.T) {
+	h := []*FooHeapObject{
+		{value: 4},
+		{value: 5},
+		{value: 2},
+	}
+	heap := NewObjectArrayBasedMinimumHeap(h, GetFooHeapObjectOrderFunc)
+	assert.Equal(t, 2, heap.Peek().value)
+	assert.Equal(t, 3, heap.Len())
+	heap.Push(&FooHeapObject{value: 3})
+	assert.Equal(t, 2, heap.Peek().value)
+	assert.Equal(t, 4, heap.Len())
+	heap.Push(&FooHeapObject{value: 1})
+	assert.Equal(t, 1, heap.Peek().value)
+	assert.Equal(t, 5, heap.Len())
+	for i := 1; i <= 5; i++ {
+		assert.Equal(t, i, heap.Peek().value)
+		assert.Equal(t, i, heap.Pop().value)
+	}
+}
+
+func TestMaximumObjectHeap(t *testing.T) {
+	h := []*FooHeapObject{
+		{value: 4},
+		{value: 1},
+		{value: 2},
+	}
+	heap := NewObjectArrayBasedMaximumHeap(h, GetFooHeapObjectOrderFunc)
+	assert.Equal(t, 4, heap.Peek().value)
+	assert.Equal(t, 3, heap.Len())
+	heap.Push(&FooHeapObject{value: 3})
+	assert.Equal(t, 4, heap.Peek().value)
+	assert.Equal(t, 4, heap.Len())
+	heap.Push(&FooHeapObject{value: 5})
+	assert.Equal(t, 5, heap.Peek().value)
+	assert.Equal(t, 5, heap.Len())
+	for i := 5; i >= 1; i-- {
+		assert.Equal(t, i, heap.Peek().value)
+		assert.Equal(t, i, heap.Pop().value)
+	}
+}