Add injection logic for FlushManager (#10580)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

internal/datanode/flush_manager.go

@@ -52,7 +52,10 @@ type segmentFlushPack struct {
 type notifyMetaFunc func(*segmentFlushPack) error
 
 // taskPostFunc clean up function after single flush task done
-type taskPostFunc func()
+type taskPostFunc func(pack *segmentFlushPack, postInjection postInjectionFunc)
+
+// postInjectionFunc post injection pack process logic
+type postInjectionFunc func(pack *segmentFlushPack)
 
 // make sure implementation
 var _ flushManager = (*rendezvousFlushManager)(nil)
@@ -60,20 +63,30 @@ var _ flushManager = (*rendezvousFlushManager)(nil)
 type orderFlushQueue struct {
 	sync.Once
 	segmentID UniqueID
+	injectCh  chan taskInjection
+
 	// MsgID => flushTask
 	working    sync.Map
 	notifyFunc notifyMetaFunc
 
 	tailMut sync.Mutex
 	tailCh  chan struct{}
+
+	injectMut     sync.Mutex
+	runningTasks  int32
+	injectHandler *injectHandler
+	postInjection postInjectionFunc
 }
 
 // newOrderFlushQueue creates a orderFlushQueue
 func newOrderFlushQueue(segID UniqueID, f notifyMetaFunc) *orderFlushQueue {
-	return &orderFlushQueue{
+	q := &orderFlushQueue{
 		segmentID:  segID,
 		notifyFunc: f,
+		injectCh:   make(chan taskInjection, 100),
 	}
+	q.injectHandler = newInjectHandler(q)
+	return q
 }
 
 // init orderFlushQueue use once protect init, init tailCh
@@ -86,19 +99,43 @@ func (q *orderFlushQueue) init() {
 }
 
 func (q *orderFlushQueue) getFlushTaskRunner(pos *internalpb.MsgPosition) *flushTaskRunner {
-	actual, loaded := q.working.LoadOrStore(string(pos.MsgID), newFlushTaskRunner(q.segmentID))
+	actual, loaded := q.working.LoadOrStore(string(pos.MsgID), newFlushTaskRunner(q.segmentID, q.injectCh))
 	t := actual.(*flushTaskRunner)
 	if !loaded {
+
+		q.injectMut.Lock()
+		q.runningTasks++
+		if q.injectHandler != nil {
+			q.injectHandler.close()
+			q.injectHandler = nil
+		}
+		q.injectMut.Unlock()
+
 		q.tailMut.Lock()
-		t.init(q.notifyFunc, func() {
-			q.working.Delete(string(pos.MsgID))
-		}, q.tailCh)
+		t.init(q.notifyFunc, q.postTask, q.tailCh)
 		q.tailCh = t.finishSignal
 		q.tailMut.Unlock()
 	}
 	return t
 }
 
+func (q *orderFlushQueue) postTask(pack *segmentFlushPack, postInjection postInjectionFunc) {
+	q.working.Delete(string(pack.pos.MsgID))
+	q.injectMut.Lock()
+	q.runningTasks--
+	if q.runningTasks == 0 {
+		q.injectHandler = newInjectHandler(q)
+	}
+	if postInjection != nil {
+		q.postInjection = postInjection
+	}
+
+	if q.postInjection != nil {
+		q.postInjection(pack)
+	}
+	q.injectMut.Unlock()
+}
+
 // enqueueInsertBuffer put insert buffer data into queue
 func (q *orderFlushQueue) enqueueInsertFlush(task flushInsertTask, binlogs, statslogs map[UniqueID]string, flushed bool, pos *internalpb.MsgPosition) {
 	q.getFlushTaskRunner(pos).runFlushInsert(task, binlogs, statslogs, flushed, pos)
@@ -109,6 +146,53 @@ func (q *orderFlushQueue) enqueueDelFlush(task flushDeleteTask, deltaLogs *DelDa
 	q.getFlushTaskRunner(pos).runFlushDel(task, deltaLogs)
 }
 
+// inject performs injection for current task queue
+// send into injectCh in there is running task
+// or perform injection logic here if there is no injection
+func (q *orderFlushQueue) inject(inject taskInjection) {
+	q.injectCh <- inject
+}
+
+type injectHandler struct {
+	once sync.Once
+	wg   sync.WaitGroup
+	done chan struct{}
+}
+
+func newInjectHandler(q *orderFlushQueue) *injectHandler {
+	h := &injectHandler{
+		done: make(chan struct{}),
+	}
+	h.wg.Add(1)
+	go h.handleInjection(q)
+	return h
+}
+
+func (h *injectHandler) handleInjection(q *orderFlushQueue) {
+	defer h.wg.Done()
+	for {
+		select {
+		case inject := <-q.injectCh:
+			q.tailMut.Lock() //Maybe double check
+			injectDone := make(chan struct{})
+			q.tailCh = injectDone
+			q.tailMut.Unlock()
+			inject.injected <- struct{}{}
+			<-inject.injectOver
+			close(injectDone)
+		case <-h.done:
+			return
+		}
+	}
+}
+
+func (h *injectHandler) close() {
+	h.once.Do(func() {
+		close(h.done)
+		h.wg.Wait()
+	})
+}
+
 // rendezvousFlushManager makes sure insert & del buf all flushed
 type rendezvousFlushManager struct {
 	allocatorInterface
@@ -252,6 +336,13 @@ func (m *rendezvousFlushManager) flushDelData(data *DelDataBuf, segmentID Unique
 	return nil
 }
 
+// injectFlush inject process before task finishes
+func (m *rendezvousFlushManager) injectFlush(injection taskInjection, segments ...UniqueID) {
+	for _, segmentID := range segments {
+		m.getFlushQueue(segmentID).inject(injection)
+	}
+}
+
 // fetch meta info for segment
 func (m *rendezvousFlushManager) getSegmentMeta(segmentID UniqueID, pos *internalpb.MsgPosition) (UniqueID, UniqueID, *etcdpb.CollectionMeta, error) {
 	if !m.hasSegment(segmentID, true) {

internal/datanode/flush_manager_test.go

@@ -158,6 +158,95 @@ func TestRendezvousFlushManager(t *testing.T) {
 	finish.Wait()
 
 	assert.EqualValues(t, size, counter.Load())
+}
+
+func TestRendezvousFlushManager_Inject(t *testing.T) {
+	kv := memkv.NewMemoryKV()
+
+	size := 1000
+	var counter atomic.Int64
+	finish := sync.WaitGroup{}
+	finish.Add(size)
+	packs := make([]*segmentFlushPack, 0, size+1)
+	m := NewRendezvousFlushManager(&allocator{}, kv, newMockReplica(), func(pack *segmentFlushPack) error {
+		packs = append(packs, pack)
+		counter.Inc()
+		finish.Done()
+		return nil
+	})
+
+	injected := make(chan struct{})
+	injectOver := make(chan bool)
+	m.injectFlush(taskInjection{
+		injected:   injected,
+		injectOver: injectOver,
+		postInjection: func(*segmentFlushPack) {
+		},
+	}, 1)
+	<-injected
+	injectOver <- true
+
+	ids := make([][]byte, 0, size)
+	for i := 0; i < size; i++ {
+		id := make([]byte, 10)
+		rand.Read(id)
+		ids = append(ids, id)
+	}
+
+	wg := sync.WaitGroup{}
+	wg.Add(size)
+	for i := 0; i < size; i++ {
+		m.flushDelData(nil, 1, &internalpb.MsgPosition{
+			MsgID: ids[i],
+		})
+		m.flushBufferData(nil, 1, true, &internalpb.MsgPosition{
+			MsgID: ids[i],
+		})
+		wg.Done()
+	}
+	wg.Wait()
+	finish.Wait()
+
+	assert.EqualValues(t, size, counter.Load())
+
+	finish.Add(1)
+	id := make([]byte, 10)
+	rand.Read(id)
+	m.flushBufferData(nil, 2, true, &internalpb.MsgPosition{
+		MsgID: id,
+	})
+
+	m.injectFlush(taskInjection{
+		injected:   injected,
+		injectOver: injectOver,
+		postInjection: func(pack *segmentFlushPack) {
+			pack.segmentID = 3
+		},
+	}, 2)
+
+	go func() {
+		<-injected
+		injectOver <- true
+	}()
+	m.flushDelData(nil, 2, &internalpb.MsgPosition{
+		MsgID: id,
+	})
+
+	finish.Wait()
+	assert.EqualValues(t, size+1, counter.Load())
+	assert.EqualValues(t, 3, packs[size].segmentID)
+
+	finish.Add(1)
+	rand.Read(id)
+	m.flushBufferData(nil, 2, false, &internalpb.MsgPosition{
+		MsgID: id,
+	})
+	m.flushDelData(nil, 2, &internalpb.MsgPosition{
+		MsgID: id,
+	})
+	finish.Wait()
+	assert.EqualValues(t, size+2, counter.Load())
+	assert.EqualValues(t, 3, packs[size+1].segmentID)
 
 }
 

internal/datanode/flush_task.go

@@ -50,6 +50,7 @@ type flushTaskRunner struct {
 
 	startSignal  <-chan struct{}
 	finishSignal chan struct{}
+	injectSignal <-chan taskInjection
 
 	segmentID  UniqueID
 	insertLogs map[UniqueID]string
@@ -59,6 +60,12 @@ type flushTaskRunner struct {
 	flushed    bool
 }
 
+type taskInjection struct {
+	injected      chan struct{} // channel to notify injected
+	injectOver    chan bool     // indicates injection over
+	postInjection func(pack *segmentFlushPack)
+}
+
 // init initializes flushTaskRunner with provided actions and signal
 func (t *flushTaskRunner) init(f notifyMetaFunc, postFunc taskPostFunc, signal <-chan struct{}) {
 	t.initOnce.Do(func() {
@@ -109,7 +116,23 @@ func (t *flushTaskRunner) waitFinish(notifyFunc notifyMetaFunc, postFunc taskPos
 	t.Wait()
 	// wait previous task done
 	<-t.startSignal
+
 	pack := t.getFlushPack()
+	var postInjection postInjectionFunc = nil
+	select {
+	case injection := <-t.injectSignal:
+		// notify injected
+		injection.injected <- struct{}{}
+		ok := <-injection.injectOver
+		if ok {
+			// apply postInjection func
+			postInjection = injection.postInjection
+		}
+	default:
+	}
+	postFunc(pack, postInjection)
+
+	// execution done, dequeue and make count --
 	err := errStart
 	for err != nil {
 		err = notifyFunc(pack)
@@ -117,7 +140,6 @@ func (t *flushTaskRunner) waitFinish(notifyFunc notifyMetaFunc, postFunc taskPos
 
 	// notify next task
 	close(t.finishSignal)
-	postFunc()
 }
 
 func (t *flushTaskRunner) getFlushPack() *segmentFlushPack {
@@ -134,10 +156,11 @@ func (t *flushTaskRunner) getFlushPack() *segmentFlushPack {
 }
 
 // newFlushTaskRunner create a usable task runner
-func newFlushTaskRunner(segmentID UniqueID) *flushTaskRunner {
+func newFlushTaskRunner(segmentID UniqueID, injectCh <-chan taskInjection) *flushTaskRunner {
 	t := &flushTaskRunner{
-		WaitGroup: sync.WaitGroup{},
-		segmentID: segmentID,
+		WaitGroup:    sync.WaitGroup{},
+		segmentID:    segmentID,
+		injectSignal: injectCh,
 	}
 	// insert & del
 	t.Add(2)

internal/datanode/flush_task_test.go

@@ -23,7 +23,7 @@ import (
 )
 
 func TestFlushTaskRunner(t *testing.T) {
-	task := newFlushTaskRunner(1)
+	task := newFlushTaskRunner(1, nil)
 	signal := make(chan struct{})
 
 	saveFlag := false
@@ -33,7 +33,7 @@ func TestFlushTaskRunner(t *testing.T) {
 	task.init(func(*segmentFlushPack) error {
 		saveFlag = true
 		return nil
-	}, func() {}, signal)
+	}, func(pack *segmentFlushPack, i postInjectionFunc) {}, signal)
 
 	go func() {
 		<-task.finishSignal
@@ -56,3 +56,62 @@ func TestFlushTaskRunner(t *testing.T) {
 	assert.True(t, saveFlag)
 	assert.True(t, nextFlag)
 }
+
+func TestFlushTaskRunner_Injection(t *testing.T) {
+	injectCh := make(chan taskInjection, 1)
+	task := newFlushTaskRunner(1, injectCh)
+	signal := make(chan struct{})
+
+	saveFlag := false
+	nextFlag := false
+	processed := make(chan struct{})
+
+	injected := make(chan struct{})
+	injectOver := make(chan bool)
+
+	injectCh <- taskInjection{
+		injected:   injected,
+		injectOver: injectOver,
+		postInjection: func(pack *segmentFlushPack) {
+			t.Log("task injection executed")
+			pack.segmentID = 2
+		},
+	}
+
+	go func() {
+		<-injected
+		injectOver <- true
+	}()
+
+	task.init(func(pack *segmentFlushPack) error {
+		assert.EqualValues(t, 2, pack.segmentID)
+		saveFlag = true
+		return nil
+	}, func(pack *segmentFlushPack, i postInjectionFunc) {
+		if i != nil {
+			i(pack)
+		}
+	}, signal)
+
+	go func() {
+		<-task.finishSignal
+		nextFlag = true
+		processed <- struct{}{}
+	}()
+
+	assert.False(t, saveFlag)
+	assert.False(t, nextFlag)
+
+	task.runFlushInsert(&emptyFlushTask{}, nil, nil, false, nil)
+	task.runFlushDel(&emptyFlushTask{}, &DelDataBuf{})
+
+	assert.False(t, saveFlag)
+	assert.False(t, nextFlag)
+
+	close(signal)
+	<-processed
+
+	assert.True(t, saveFlag)
+	assert.True(t, nextFlag)
+
+}