mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-28 14:35:27 +08:00
7c575a18b0
issue: #43897 also for issue: #46166 add ack_sync_up flag into broadcast message header, which indicates that whether the broadcast operation is need to be synced up between the streaming node and the coordinator. If the ack_sync_up is false, the broadcast operation will be acked once the recovery storage see the message at current vchannel, the fast ack operation can be applied to speed up the broadcast operation. If the ack_sync_up is true, the broadcast operation will be acked after the checkpoint of current vchannel reach current message. The fast ack operation can not be applied to speed up the broadcast operation, because the ack operation need to be synced up with streaming node. e.g. if truncate collection operation want to call ack once callback after the all segment are flushed at current vchannel, it should set the ack_sync_up to be true. TODO: current implementation doesn't promise the ack sync up semantic, it only promise FastAck operation will not be applied, wait for 3.0 to implement the ack sync up semantic. only for truncate api now. --------- Signed-off-by: chyezh <chyezh@outlook.com>
107 lines
3.7 KiB
Go
107 lines
3.7 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/mock"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
|
|
"github.com/milvus-io/milvus/internal/coordinator/snmanager"
|
|
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_balancer"
|
|
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_broadcaster"
|
|
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer"
|
|
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/balance"
|
|
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster"
|
|
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster/broadcast"
|
|
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
|
|
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
|
|
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
|
|
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/walimplstest"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
|
|
)
|
|
|
|
func TestBroadcastService(t *testing.T) {
|
|
broadcast.ResetBroadcaster()
|
|
snmanager.ResetStreamingNodeManager()
|
|
// Set up the balancer
|
|
b := mock_balancer.NewMockBalancer(t)
|
|
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
|
|
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
|
|
<-ctx.Done()
|
|
return ctx.Err()
|
|
})
|
|
b.EXPECT().Close().Return().Maybe()
|
|
balance.Register(b)
|
|
|
|
fb := syncutil.NewFuture[broadcaster.Broadcaster]()
|
|
mba := mock_broadcaster.NewMockBroadcastAPI(t)
|
|
mba.EXPECT().Close().Return()
|
|
mb := mock_broadcaster.NewMockBroadcaster(t)
|
|
fb.Set(mb)
|
|
mba.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(&types.BroadcastAppendResult{}, nil)
|
|
mb.EXPECT().WithResourceKeys(mock.Anything, mock.Anything).Return(mba, nil)
|
|
mb.EXPECT().Ack(mock.Anything, mock.Anything).Return(nil)
|
|
mb.EXPECT().LegacyAck(mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
|
mb.EXPECT().Close().Return().Maybe()
|
|
broadcast.Register(mb)
|
|
|
|
msg := message.NewCreateCollectionMessageBuilderV1().
|
|
WithHeader(&message.CreateCollectionMessageHeader{}).
|
|
WithBody(&msgpb.CreateCollectionRequest{}).
|
|
WithBroadcast([]string{"v1"}).
|
|
MustBuildBroadcast()
|
|
|
|
service := NewBroadcastService()
|
|
service.Broadcast(context.Background(), &streamingpb.BroadcastRequest{
|
|
Message: msg.IntoMessageProto(),
|
|
})
|
|
service.Ack(context.Background(), &streamingpb.BroadcastAckRequest{
|
|
BroadcastId: 1,
|
|
Vchannel: "v1",
|
|
})
|
|
service.Ack(context.Background(), &streamingpb.BroadcastAckRequest{
|
|
BroadcastId: 1,
|
|
Vchannel: "v1",
|
|
Message: &commonpb.ImmutableMessage{
|
|
Id: walimplstest.NewTestMessageID(1).IntoProto(),
|
|
Payload: []byte("payload"),
|
|
Properties: map[string]string{"key": "value"},
|
|
},
|
|
})
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
reached := make(chan struct{})
|
|
done := make(chan struct{})
|
|
mb.EXPECT().Ack(mock.Anything, mock.Anything).Unset()
|
|
mb.EXPECT().Ack(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, msg message.ImmutableMessage) error {
|
|
close(reached)
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-done:
|
|
return nil
|
|
}
|
|
})
|
|
go func() {
|
|
<-reached
|
|
cancel()
|
|
time.Sleep(10 * time.Millisecond)
|
|
close(done)
|
|
}()
|
|
_, err := service.Ack(ctx, &streamingpb.BroadcastAckRequest{
|
|
BroadcastId: 1,
|
|
Vchannel: "v1",
|
|
Message: &commonpb.ImmutableMessage{
|
|
Id: walimplstest.NewTestMessageID(1).IntoProto(),
|
|
Payload: []byte("payload"),
|
|
Properties: map[string]string{"key": "value"},
|
|
},
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
}
|