498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
384 lines
14 KiB
Go
384 lines
14 KiB
Go
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package delegator
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
|
|
"github.com/milvus-io/milvus/internal/querynodev2/segments"
|
|
)
|
|
|
|
func TestDelegatorGrowingFlushSourcePassesTaskSchema(t *testing.T) {
|
|
ctx := context.Background()
|
|
schema := &schemapb.CollectionSchema{Name: "task-schema"}
|
|
segment := segments.NewMockSegment(t)
|
|
segment.EXPECT().
|
|
FlushData(ctx, int64(3), int64(7), mock.Anything).
|
|
RunAndReturn(func(_ context.Context, startOffset int64, endOffset int64, config *segments.FlushConfig) (*segments.FlushResult, error) {
|
|
require.EqualValues(t, 3, startOffset)
|
|
require.EqualValues(t, 7, endOffset)
|
|
require.True(t, config.Schema == schema)
|
|
return &segments.FlushResult{
|
|
ManifestPath: "manifest",
|
|
NumRows: 4,
|
|
TimestampFrom: 100,
|
|
TimestampTo: 200,
|
|
ColumnGroupMemorySizes: map[int64]int64{100: 64},
|
|
FieldNullCounts: map[int64]int64{100: 1},
|
|
}, nil
|
|
})
|
|
|
|
source := &delegatorGrowingFlushSource{segment: segment}
|
|
result, err := source.FlushGrowingData(ctx, 3, 7, &syncmgr.GrowingFlushConfig{
|
|
Schema: schema,
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "manifest", result.ManifestPath)
|
|
require.EqualValues(t, 4, result.NumRows)
|
|
require.EqualValues(t, 100, result.TimestampFrom)
|
|
require.EqualValues(t, 200, result.TimestampTo)
|
|
require.EqualValues(t, 64, result.ColumnGroupMemorySizes[100])
|
|
require.EqualValues(t, 1, result.FieldNullCounts[100])
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderCloseWaitsForSourceRelease(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(10)).Once()
|
|
segment.EXPECT().Unpin().Once()
|
|
|
|
source, state := provider.GetGrowingFlushSource(1001, 10, nil)
|
|
require.Equal(t, syncmgr.GrowingSourceUsable, state)
|
|
require.NotNil(t, source)
|
|
|
|
closeDone := make(chan struct{})
|
|
go func() {
|
|
provider.Close()
|
|
close(closeDone)
|
|
}()
|
|
|
|
select {
|
|
case <-closeDone:
|
|
t.Fatal("provider close should wait for active growing source lease")
|
|
case <-time.After(50 * time.Millisecond):
|
|
}
|
|
|
|
source.Release()
|
|
require.Eventually(t, func() bool {
|
|
select {
|
|
case <-closeDone:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}, time.Second, 10*time.Millisecond)
|
|
|
|
source, state = provider.GetGrowingFlushSource(1001, 10, nil)
|
|
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
|
|
require.Nil(t, source)
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderRetainedSource(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(10)).Once()
|
|
|
|
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 0, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 10},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Twice()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Twice()
|
|
segment.EXPECT().InsertCount().Return(int64(10)).Twice()
|
|
segment.EXPECT().Unpin().Times(3)
|
|
segmentManager.EXPECT().ReleaseDetached(context.Background(), segment).Once()
|
|
|
|
source, state := provider.GetGrowingFlushSource(1001, 10, nil)
|
|
require.Equal(t, syncmgr.GrowingSourceUsable, state)
|
|
require.NotNil(t, source)
|
|
source.Release()
|
|
|
|
source, state = provider.GetGrowingFlushSource(1001, 10, nil)
|
|
require.Equal(t, syncmgr.GrowingSourceUsable, state)
|
|
require.NotNil(t, source)
|
|
source.Release()
|
|
source.(syncmgr.GrowingFlushSourceCommitter).CommitGrowingFlush(10)
|
|
segmentManager.AssertNotCalled(t, "ReleaseDetached", context.Background(), segment)
|
|
provider.MarkReleaseDetached(1001)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
|
|
source, state = provider.GetGrowingFlushSource(1001, 10, nil)
|
|
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
|
|
require.Nil(t, source)
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderUsesInsertCountAsOffset(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(12)).Twice()
|
|
segment.EXPECT().Unpin().Once()
|
|
|
|
source, state := provider.GetGrowingFlushSource(1001, 12, nil)
|
|
require.Equal(t, syncmgr.GrowingSourceUsable, state)
|
|
require.NotNil(t, source)
|
|
require.EqualValues(t, 12, source.CurrentOffset())
|
|
source.Release()
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderMissingSegmentPendingUntilTSafeCaughtUp(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
currentTSafe := uint64(100)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, nil, func() uint64 {
|
|
return currentTSafe
|
|
})
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
|
|
source, state := provider.GetGrowingFlushSource(1001, 10, &msgpb.MsgPosition{Timestamp: 200})
|
|
require.Equal(t, syncmgr.GrowingSourcePending, state)
|
|
require.Nil(t, source)
|
|
|
|
currentTSafe = 200
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
|
|
source, state = provider.GetGrowingFlushSource(1001, 10, &msgpb.MsgPosition{Timestamp: 200})
|
|
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
|
|
require.Nil(t, source)
|
|
|
|
currentTSafe = 100
|
|
provider.Deactivate()
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
|
|
source, state = provider.GetGrowingFlushSource(1001, 10, &msgpb.MsgPosition{Timestamp: 200})
|
|
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
|
|
require.Nil(t, source)
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderPrepareWaitsFence(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
waitedFence := uint64(0)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, func(ctx context.Context, fenceTs uint64) error {
|
|
waitedFence = fenceTs
|
|
return nil
|
|
})
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(10)).Once()
|
|
|
|
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 10},
|
|
{SegmentID: 1002},
|
|
})
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, 200, waitedFence)
|
|
require.True(t, provider.IsReleasePrepared(1001, 0))
|
|
require.True(t, provider.IsReleaseAllowed(1001, 0))
|
|
require.False(t, provider.IsReleasePrepared(1002, 200))
|
|
require.True(t, provider.IsReleaseAllowed(1002, 200))
|
|
require.False(t, provider.IsReleaseAllowed(1002, 201))
|
|
require.False(t, provider.IsReleasePrepared(1003, 200))
|
|
require.False(t, provider.IsReleaseAllowed(1003, 200))
|
|
|
|
segment.EXPECT().Unpin().Once()
|
|
segmentManager.EXPECT().ReleaseDetached(context.Background(), segment).Once()
|
|
provider.releaseRetainedIfComplete(1001, 10)
|
|
segmentManager.AssertNotCalled(t, "ReleaseDetached", context.Background(), segment)
|
|
provider.MarkReleaseDetached(1001)
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderHandoffOnlyRejectsNewSegmentsBeforeFence(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
waitFenceEntered := make(chan struct{})
|
|
releaseFence := make(chan struct{})
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, func(ctx context.Context, fenceTs uint64) error {
|
|
close(waitFenceEntered)
|
|
<-releaseFence
|
|
return nil
|
|
})
|
|
|
|
prepareDone := make(chan error, 1)
|
|
go func() {
|
|
prepareDone <- provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 10},
|
|
})
|
|
}()
|
|
|
|
select {
|
|
case <-waitFenceEntered:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("prepare did not enter fence wait")
|
|
}
|
|
|
|
source, state := provider.GetGrowingFlushSource(1002, 1, nil)
|
|
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
|
|
require.Nil(t, source)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(10)).Once()
|
|
close(releaseFence)
|
|
require.NoError(t, <-prepareDone)
|
|
|
|
segment.EXPECT().Unpin().Once()
|
|
provider.Close()
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderDeactivatedOnlyServesRetainedSources(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
waitCount := 0
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, func(ctx context.Context, fenceTs uint64) error {
|
|
waitCount++
|
|
return nil
|
|
})
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(10)).Once()
|
|
|
|
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 100, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 10},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, waitCount)
|
|
|
|
provider.Deactivate()
|
|
|
|
err = provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1002, TargetOffset: 5},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, waitCount)
|
|
require.False(t, provider.IsReleaseAllowed(1002, 200))
|
|
require.False(t, provider.IsReleasePrepared(1002, 200))
|
|
require.True(t, provider.IsReleaseAllowed(1001, 100))
|
|
require.True(t, provider.IsReleasePrepared(1001, 100))
|
|
|
|
err = provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 11},
|
|
})
|
|
require.NoError(t, err)
|
|
require.False(t, provider.IsReleasePrepared(1001, 200))
|
|
|
|
segment.EXPECT().Unpin().Once()
|
|
segmentManager.EXPECT().ReleaseDetached(context.Background(), segment).Once()
|
|
provider.releaseRetainedIfComplete(1001, 10)
|
|
segmentManager.AssertNotCalled(t, "ReleaseDetached", context.Background(), segment)
|
|
provider.MarkReleaseDetached(1001)
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderReleasesWhenDetachedBeforeCommit(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(10)).Once()
|
|
|
|
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 0, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 10},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
provider.MarkReleaseDetached(1001)
|
|
segmentManager.AssertNotCalled(t, "ReleaseDetached", context.Background(), segment)
|
|
|
|
segment.EXPECT().Unpin().Once()
|
|
segmentManager.EXPECT().ReleaseDetached(context.Background(), segment).Once()
|
|
provider.releaseRetainedIfComplete(1001, 10)
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderPrepareRollbackOnFailure(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
behindSegment := segments.NewMockSegment(t)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(10)).Once()
|
|
segmentManager.EXPECT().GetGrowing(int64(1002)).Return(behindSegment).Once()
|
|
behindSegment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
behindSegment.EXPECT().InsertCount().Return(int64(5)).Once()
|
|
behindSegment.EXPECT().Unpin().Once()
|
|
segment.EXPECT().Unpin().Once()
|
|
|
|
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 10},
|
|
{SegmentID: 1002, TargetOffset: 10},
|
|
})
|
|
require.Error(t, err)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
|
|
source, state := provider.GetGrowingFlushSource(1001, 10, nil)
|
|
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
|
|
require.Nil(t, source)
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderRegisterRetainedSegmentNotFound(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
|
|
registry := syncmgr.NewGrowingSourceRegistry()
|
|
registration := registry.Register("ch", provider)
|
|
defer registry.Unregister(registration)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
|
|
|
|
err := registry.PrepareGrowingSourceReleaseHandoff(context.Background(), "ch", 0, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 10},
|
|
})
|
|
require.Error(t, err)
|
|
}
|
|
|
|
func TestDelegatorGrowingSourceProviderRegisterRetainedBehindTarget(t *testing.T) {
|
|
segmentManager := segments.NewMockSegmentManager(t)
|
|
segment := segments.NewMockSegment(t)
|
|
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
|
|
|
|
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
|
|
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
|
|
segment.EXPECT().InsertCount().Return(int64(5)).Once()
|
|
segment.EXPECT().Unpin().Once()
|
|
|
|
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 0, []syncmgr.GrowingSourceReleaseHandoffSegment{
|
|
{SegmentID: 1001, TargetOffset: 10},
|
|
})
|
|
require.Error(t, err)
|
|
}
|