fix: async warmup will be blocked by state lock (#33686)

issue: #33685

Signed-off-by: chyezh <chyezh@outlook.com>

internal/querynodev2/segments/segment.go

@@ -1415,6 +1415,10 @@ func (s *LocalSegment) WarmupChunkCache(ctx context.Context, fieldID int64) {
 		}).Await()
 	case "async":
 		GetWarmupPool().Submit(func() (any, error) {
+			// bad implemtation, warmup is async at another goroutine and hold the rlock.
+			// the state transition of segment in segment loader will blocked.
+			// add a waiter to avoid it.
+			s.ptrLock.BlockUntilDataLoadedOrReleased()
 			if !s.ptrLock.RLockIf(state.IsNotReleased) {
 				return nil, nil
 			}

internal/querynodev2/segments/state/load_state_lock.go

@@ -174,6 +174,16 @@ func (ls *LoadStateLock) StartReleaseAll() (g LoadStateLockGuard) {
 	}
 }
 
+// blockUntilDataLoadedOrReleased blocks until the segment is loaded or released.
+func (ls *LoadStateLock) BlockUntilDataLoadedOrReleased() {
+	ls.cv.L.Lock()
+	defer ls.cv.L.Unlock()
+
+	for ls.state != LoadStateDataLoaded && ls.state != LoadStateReleased {
+		ls.cv.Wait()
+	}
+}
+
 // waitUntilCanReleaseData waits until segment is release data able.
 func (ls *LoadStateLock) waitUntilCanReleaseData() {
 	state := ls.state

internal/querynodev2/segments/state/load_state_lock_test.go

@@ -103,6 +103,24 @@ func TestStartReleaseData(t *testing.T) {
 	assert.Equal(t, LoadStateOnlyMeta, l.state)
 }
 
+func TestBlockUntilDataLoadedOrReleased(t *testing.T) {
+	l := NewLoadStateLock(LoadStateOnlyMeta)
+	ch := make(chan struct{})
+	go func() {
+		l.BlockUntilDataLoadedOrReleased()
+		close(ch)
+	}()
+	select {
+	case <-ch:
+		t.Errorf("should be blocked")
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	g, _ := l.StartLoadData()
+	g.Done(nil)
+	<-ch
+}
+
 func TestStartReleaseAll(t *testing.T) {
 	l := NewLoadStateLock(LoadStateOnlyMeta)
 	// Test Release All, nothing to do on only meta.