From db0ff9c6b4a93bb13c13a6688d33fc8f8309c172 Mon Sep 17 00:00:00 2001 From: sijie-ni-0214 Date: Thu, 10 Sep 2026 13:10:50 +0000 Subject: [PATCH] enhance: pipeline sealed BM25 stats materialization Merge sealed segment statistics as individual loads complete so lazy IDF materialization overlaps object storage reads with aggregation. Reuse immutable cache-owned statistics and adopt the first deserialized field value to avoid redundant full copies. Add coverage for pipelined consumption, limiter cancellation, lease cleanup, and shared cache statistics. Signed-off-by: sijie-ni-0214 --- .../server/wal/vchannel/idf/oracle.go | 60 +++++++++--- .../server/wal/vchannel/idf/oracle_test.go | 95 +++++++++++++++++++ .../server/wal/vchannel/idf/segment_cache.go | 14 +-- .../wal/vchannel/idf/segment_cache_test.go | 35 +++++++ 4 files changed, 183 insertions(+), 21 deletions(-) diff --git a/internal/streamingnode/server/wal/vchannel/idf/oracle.go b/internal/streamingnode/server/wal/vchannel/idf/oracle.go index 79287bc96b5..ffad3710333 100644 --- a/internal/streamingnode/server/wal/vchannel/idf/oracle.go +++ b/internal/streamingnode/server/wal/vchannel/idf/oracle.go @@ -593,15 +593,23 @@ func (r *oracleRuntime) materialize(call *materializationCall) { resultErr = merr.Wrapf(err, "get sealed BM25 resources for data version %s", call.target.String()) return } - sealed, err = r.provider.acquireSealedContributions(call.ctx, resources) + var stats bm25Stats + sealed, err = r.provider.acquireSealedContributionsWithConsumer( + call.ctx, + resources, + func(contribution sealedContribution) { + if stats == nil { + stats = newBM25StatsFromSchema(r.schema) + } + stats.merge(contribution.stats) + }, + ) if err != nil { resultErr = merr.Wrapf(err, "load sealed BM25 stats for data version %s", call.target.String()) return } - - stats := newBM25StatsFromSchema(r.schema) - for _, contribution := range sealed { - stats.merge(contribution.stats) + if stats == nil { + stats = newBM25StatsFromSchema(r.schema) } var oldSealed map[int64]sealedContribution @@ -1081,6 +1089,14 @@ func (p *Provider) getSealedBM25Resources( func (p *Provider) acquireSealedContributions( ctx context.Context, resources []*datapb.StreamingNodeBM25Resource, +) (map[int64]sealedContribution, error) { + return p.acquireSealedContributionsWithConsumer(ctx, resources, nil) +} + +func (p *Provider) acquireSealedContributionsWithConsumer( + ctx context.Context, + resources []*datapb.StreamingNodeBM25Resource, + consume func(sealedContribution), ) (map[int64]sealedContribution, error) { loaded := make([]sealedContribution, len(resources)) keepLeases := false @@ -1099,13 +1115,25 @@ func (p *Provider) acquireSealedContributions( if limiter == nil { limiter = getGlobalSealedStatsLoadLimiter() } + results := make(chan sealedContribution, len(resources)) + contributions := make(map[int64]sealedContribution, len(resources)) + collectorDone := make(chan struct{}) + go func() { + defer close(collectorDone) + for contribution := range results { + contributions[contribution.segmentID] = contribution + if consume != nil { + consume(contribution) + } + } + }() + group, groupCtx := errgroup.WithContext(ctx) + var acquireErr error for i, resource := range resources { if err := limiter.Acquire(groupCtx, 1); err != nil { - if groupErr := group.Wait(); groupErr != nil { - return nil, groupErr - } - return nil, err + acquireErr = err + break } i, resource := i, resource group.Go(func() error { @@ -1120,16 +1148,18 @@ func (p *Provider) acquireSealedContributions( stats: stats, lease: lease, } + results <- loaded[i] return nil }) } - if err := group.Wait(); err != nil { - return nil, err + groupErr := group.Wait() + close(results) + <-collectorDone + if groupErr != nil { + return nil, groupErr } - - contributions := make(map[int64]sealedContribution, len(loaded)) - for _, contribution := range loaded { - contributions[contribution.segmentID] = contribution + if acquireErr != nil { + return nil, acquireErr } keepLeases = true return contributions, nil diff --git a/internal/streamingnode/server/wal/vchannel/idf/oracle_test.go b/internal/streamingnode/server/wal/vchannel/idf/oracle_test.go index 246c8d66753..58f138df7a1 100644 --- a/internal/streamingnode/server/wal/vchannel/idf/oracle_test.go +++ b/internal/streamingnode/server/wal/vchannel/idf/oracle_test.go @@ -190,6 +190,101 @@ func TestAcquireSealedContributionsUsesSharedLimit(t *testing.T) { require.Equal(t, int32(2), maxActive.Load()) } +func TestAcquireSealedContributionsConsumesCompletedLoadsImmediately(t *testing.T) { + stats := storage.NewBM25Stats() + stats.Append(map[uint32]float32{1: 1}) + statsBytes, err := stats.Serialize() + require.NoError(t, err) + + secondReadStarted := make(chan struct{}) + releaseSecondRead := make(chan struct{}) + chunkManager := mocks.NewChunkManager(t) + chunkManager.EXPECT().Read(mock.Anything, mock.Anything). + RunAndReturn(func(_ context.Context, path string) ([]byte, error) { + if path == "stats-2" { + close(secondReadStarted) + <-releaseSecondRead + } + return statsBytes, nil + }).Times(2) + + provider := &Provider{ + chunkManager: chunkManager, + sealedCache: newSegmentCache(), + sealedStatsLoadLimiter: semaphore.NewWeighted(2), + } + consumed := make(chan int64, 2) + var aggregate bm25Stats + type result struct { + contributions map[int64]sealedContribution + err error + } + resultCh := make(chan result, 1) + go func() { + contributions, err := provider.acquireSealedContributionsWithConsumer( + context.Background(), + testSealedBM25Resources(1, 2), + func(contribution sealedContribution) { + if aggregate == nil { + aggregate = newBM25StatsFromSchema(testBM25WALView(qviews.DataVersion{}).Schema) + } + aggregate.merge(contribution.stats) + consumed <- contribution.segmentID + }, + ) + resultCh <- result{contributions: contributions, err: err} + }() + + <-secondReadStarted + select { + case segmentID := <-consumed: + require.Equal(t, int64(1), segmentID) + case <-time.After(time.Second): + t.Fatal("completed sealed stats were not consumed while another load was pending") + } + close(releaseSecondRead) + acquired := <-resultCh + require.NoError(t, acquired.err) + require.Len(t, acquired.contributions, 2) + require.Equal(t, int64(2), aggregate[testBM25OutputFieldID].NumRow()) + for _, contribution := range acquired.contributions { + contribution.lease.Close() + } +} + +func TestAcquireSealedContributionsCancellationWhileWaitingForLimit(t *testing.T) { + readStarted := make(chan struct{}) + chunkManager := mocks.NewChunkManager(t) + chunkManager.EXPECT().Read(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, _ string) ([]byte, error) { + close(readStarted) + <-ctx.Done() + return nil, ctx.Err() + }).Once() + + provider := &Provider{ + chunkManager: chunkManager, + sealedCache: newSegmentCache(), + sealedStatsLoadLimiter: semaphore.NewWeighted(1), + } + ctx, cancel := context.WithCancel(context.Background()) + resultCh := make(chan error, 1) + go func() { + _, err := provider.acquireSealedContributions(ctx, testSealedBM25Resources(1, 2)) + resultCh <- err + }() + + <-readStarted + cancel() + select { + case err := <-resultCh: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("sealed stats acquisition did not stop after cancellation") + } + require.Empty(t, provider.sealedCache.entries) +} + func TestAcquireSealedContributionsReleasesLeasesAfterError(t *testing.T) { stats := storage.NewBM25Stats() stats.Append(map[uint32]float32{1: 1}) diff --git a/internal/streamingnode/server/wal/vchannel/idf/segment_cache.go b/internal/streamingnode/server/wal/vchannel/idf/segment_cache.go index 38d9f1b44f2..5cb0e944919 100644 --- a/internal/streamingnode/server/wal/vchannel/idf/segment_cache.go +++ b/internal/streamingnode/server/wal/vchannel/idf/segment_cache.go @@ -36,9 +36,8 @@ func (c *segmentCache) acquire( chunkManager storage.ChunkManager, resource *datapb.StreamingNodeBM25Resource, ) (bm25Stats, *segmentCacheLease, error) { - aggregate := make(bm25Stats) if chunkManager == nil || resource == nil { - return aggregate, nil, nil + return make(bm25Stats), nil, nil } key, err := buildSealedCacheKey(resource) if err != nil { @@ -48,8 +47,7 @@ func (c *segmentCache) acquire( if err != nil { return nil, nil, err } - aggregate.merge(stats) - return aggregate, &segmentCacheLease{cache: c, keys: []sealedCacheKey{key}}, nil + return stats, &segmentCacheLease{cache: c, keys: []sealedCacheKey{key}}, nil } func (c *segmentCache) retain( @@ -146,7 +144,6 @@ func loadSealedSegmentStats( } stats := make(bm25Stats) for fieldID, paths := range pathsByField { - fieldStats := stats.getOrCreate(fieldID) for _, path := range paths { bytes, err := chunkManager.Read(ctx, path) if err != nil { @@ -156,7 +153,12 @@ func loadSealedSegmentStats( if err != nil { return nil, err } - fieldStats.Merge(loaded) + fieldStats := stats[fieldID] + if fieldStats == nil { + stats[fieldID] = loaded + } else { + fieldStats.Merge(loaded) + } } } return stats, nil diff --git a/internal/streamingnode/server/wal/vchannel/idf/segment_cache_test.go b/internal/streamingnode/server/wal/vchannel/idf/segment_cache_test.go index 4a6048bdbda..a215b92b556 100644 --- a/internal/streamingnode/server/wal/vchannel/idf/segment_cache_test.go +++ b/internal/streamingnode/server/wal/vchannel/idf/segment_cache_test.go @@ -87,3 +87,38 @@ func TestLoadSealedSegmentStatsUsesLegacyBinlogsForStorageV2(t *testing.T) { require.NoError(t, err) require.Equal(t, int64(1), loaded[102].NumRow()) } + +func TestSegmentCacheAcquireReturnsCachedStats(t *testing.T) { + paramtable.Init() + ctx := context.Background() + chunkManager := storage.NewLocalChunkManager() + stats := storage.NewBM25Stats() + stats.Append(map[uint32]float32{1: 1}) + bytes, err := stats.Serialize() + require.NoError(t, err) + statsPath := t.TempDir() + "/bm25-stats" + require.NoError(t, chunkManager.Write(ctx, statsPath, bytes)) + resource := &datapb.StreamingNodeBM25Resource{ + SegmentId: 3, + StorageVersion: storage.StorageV2, + Bm25Binlogs: []*datapb.FieldBinlog{{ + FieldID: 102, + Binlogs: []*datapb.Binlog{{LogPath: statsPath}}, + }}, + } + + cache := newSegmentCache() + first, firstLease, err := cache.acquire(ctx, chunkManager, resource) + require.NoError(t, err) + require.NotNil(t, firstLease) + second, secondLease, err := cache.acquire(ctx, chunkManager, resource) + require.NoError(t, err) + require.NotNil(t, secondLease) + key, err := buildSealedCacheKey(resource) + require.NoError(t, err) + require.Same(t, cache.entries[key].stats[102], first[102]) + require.Same(t, first[102], second[102]) + + firstLease.Close() + secondLease.Close() +}