mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-29 23:15:28 +08:00
260 lines
6.5 KiB
Go
260 lines
6.5 KiB
Go
//go:build test
|
|
// +build test
|
|
|
|
package walimpls
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math/rand"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
"github.com/remeh/sizedwaitgroup"
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
"github.com/milvus-io/milvus/internal/proto/streamingpb"
|
|
"github.com/milvus-io/milvus/internal/util/streamingutil/message"
|
|
"github.com/milvus-io/milvus/internal/util/streamingutil/options"
|
|
)
|
|
|
|
var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
|
|
|
|
func randString(l int) string {
|
|
builder := strings.Builder{}
|
|
for i := 0; i < l; i++ {
|
|
builder.WriteRune(letters[rand.Intn(len(letters))])
|
|
}
|
|
return builder.String()
|
|
}
|
|
|
|
type walImplsTestFramework struct {
|
|
b OpenerBuilderImpls
|
|
t *testing.T
|
|
messageCount int
|
|
}
|
|
|
|
func NewWALImplsTestFramework(t *testing.T, messageCount int, b OpenerBuilderImpls) *walImplsTestFramework {
|
|
return &walImplsTestFramework{
|
|
b: b,
|
|
t: t,
|
|
messageCount: messageCount,
|
|
}
|
|
}
|
|
|
|
// Run runs the test framework.
|
|
// if test failed, a error will be returned.
|
|
func (f walImplsTestFramework) Run() {
|
|
// create opener.
|
|
o, err := f.b.Build()
|
|
assert.NoError(f.t, err)
|
|
assert.NotNil(f.t, o)
|
|
defer o.Close()
|
|
|
|
// construct pChannel
|
|
name := "test_" + randString(4)
|
|
pChannel := &streamingpb.PChannelInfo{
|
|
Name: name,
|
|
Term: 1,
|
|
ServerID: 1,
|
|
VChannelInfos: []*streamingpb.VChannelInfo{},
|
|
}
|
|
ctx := context.Background()
|
|
|
|
// create a wal.
|
|
w, err := o.Open(ctx, &OpenOption{
|
|
Channel: pChannel,
|
|
})
|
|
assert.NoError(f.t, err)
|
|
assert.NotNil(f.t, w)
|
|
defer w.Close()
|
|
|
|
f.testReadAndWrite(ctx, w)
|
|
}
|
|
|
|
func (f walImplsTestFramework) testReadAndWrite(ctx context.Context, w WALImpls) {
|
|
// Test read and write.
|
|
wg := sync.WaitGroup{}
|
|
wg.Add(3)
|
|
|
|
var written []message.ImmutableMessage
|
|
var read1, read2 []message.ImmutableMessage
|
|
go func() {
|
|
defer wg.Done()
|
|
var err error
|
|
written, err = f.testAppend(ctx, w)
|
|
assert.NoError(f.t, err)
|
|
}()
|
|
go func() {
|
|
defer wg.Done()
|
|
var err error
|
|
read1, err = f.testRead(ctx, w, "scanner1")
|
|
assert.NoError(f.t, err)
|
|
}()
|
|
go func() {
|
|
defer wg.Done()
|
|
var err error
|
|
read2, err = f.testRead(ctx, w, "scanner2")
|
|
assert.NoError(f.t, err)
|
|
}()
|
|
|
|
wg.Wait()
|
|
|
|
f.assertSortedMessageList(read1)
|
|
f.assertSortedMessageList(read2)
|
|
sort.Sort(sortByMessageID(written))
|
|
f.assertEqualMessageList(written, read1)
|
|
f.assertEqualMessageList(written, read2)
|
|
|
|
// Test different scan policy, StartFrom.
|
|
readFromIdx := len(read1) / 2
|
|
readFromMsgID := read1[readFromIdx].MessageID()
|
|
s, err := w.Read(ctx, ReadOption{
|
|
Name: "scanner_deliver_start_from",
|
|
DeliverPolicy: options.DeliverPolicyStartFrom(readFromMsgID),
|
|
})
|
|
assert.NoError(f.t, err)
|
|
for i := readFromIdx; i < len(read1); i++ {
|
|
msg, ok := <-s.Chan()
|
|
assert.NotNil(f.t, msg)
|
|
assert.True(f.t, ok)
|
|
assert.True(f.t, msg.MessageID().EQ(read1[i].MessageID()))
|
|
}
|
|
s.Close()
|
|
|
|
// Test different scan policy, StartAfter.
|
|
s, err = w.Read(ctx, ReadOption{
|
|
Name: "scanner_deliver_start_after",
|
|
DeliverPolicy: options.DeliverPolicyStartAfter(readFromMsgID),
|
|
})
|
|
assert.NoError(f.t, err)
|
|
for i := readFromIdx + 1; i < len(read1); i++ {
|
|
msg, ok := <-s.Chan()
|
|
assert.NotNil(f.t, msg)
|
|
assert.True(f.t, ok)
|
|
assert.True(f.t, msg.MessageID().EQ(read1[i].MessageID()))
|
|
}
|
|
s.Close()
|
|
|
|
// Test different scan policy, Latest.
|
|
s, err = w.Read(ctx, ReadOption{
|
|
Name: "scanner_deliver_latest",
|
|
DeliverPolicy: options.DeliverPolicyLatest(),
|
|
})
|
|
assert.NoError(f.t, err)
|
|
timeoutCh := time.After(1 * time.Second)
|
|
select {
|
|
case <-s.Chan():
|
|
f.t.Errorf("should be blocked")
|
|
case <-timeoutCh:
|
|
}
|
|
s.Close()
|
|
}
|
|
|
|
func (f walImplsTestFramework) assertSortedMessageList(msgs []message.ImmutableMessage) {
|
|
for i := 1; i < len(msgs); i++ {
|
|
assert.True(f.t, msgs[i-1].MessageID().LT(msgs[i].MessageID()))
|
|
}
|
|
}
|
|
|
|
func (f walImplsTestFramework) assertEqualMessageList(msgs1 []message.ImmutableMessage, msgs2 []message.ImmutableMessage) {
|
|
assert.Equal(f.t, f.messageCount, len(msgs1))
|
|
assert.Equal(f.t, f.messageCount, len(msgs2))
|
|
for i := 0; i < len(msgs1); i++ {
|
|
assert.True(f.t, msgs1[i].MessageID().EQ(msgs2[i].MessageID()))
|
|
// assert.True(f.t, bytes.Equal(msgs1[i].Payload(), msgs2[i].Payload()))
|
|
id1, ok1 := msgs1[i].Properties().Get("id")
|
|
id2, ok2 := msgs2[i].Properties().Get("id")
|
|
assert.True(f.t, ok1)
|
|
assert.True(f.t, ok2)
|
|
assert.Equal(f.t, id1, id2)
|
|
id1, ok1 = msgs1[i].Properties().Get("const")
|
|
id2, ok2 = msgs2[i].Properties().Get("const")
|
|
assert.True(f.t, ok1)
|
|
assert.True(f.t, ok2)
|
|
assert.Equal(f.t, id1, id2)
|
|
}
|
|
}
|
|
|
|
func (f walImplsTestFramework) testAppend(ctx context.Context, w WALImpls) ([]message.ImmutableMessage, error) {
|
|
ids := make([]message.ImmutableMessage, f.messageCount)
|
|
swg := sizedwaitgroup.New(5)
|
|
for i := 0; i < f.messageCount; i++ {
|
|
swg.Add()
|
|
go func(i int) {
|
|
defer swg.Done()
|
|
// ...rocksmq has a dirty implement of properties,
|
|
// without commonpb.MsgHeader, it can not work.
|
|
header := commonpb.MsgHeader{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Insert,
|
|
MsgID: int64(i),
|
|
},
|
|
}
|
|
payload, err := proto.Marshal(&header)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
properties := map[string]string{
|
|
"id": fmt.Sprintf("%d", i),
|
|
"const": "t",
|
|
}
|
|
typ := message.MessageTypeUnknown
|
|
msg := message.NewBuilder().
|
|
WithPayload(payload).
|
|
WithProperties(properties).
|
|
WithMessageType(typ).
|
|
BuildMutable()
|
|
id, err := w.Append(ctx, msg)
|
|
assert.NoError(f.t, err)
|
|
assert.NotNil(f.t, id)
|
|
ids[i] = message.NewBuilder().
|
|
WithPayload(payload).
|
|
WithProperties(properties).
|
|
WithMessageID(id).
|
|
WithMessageType(typ).
|
|
BuildImmutable()
|
|
}(i)
|
|
}
|
|
swg.Wait()
|
|
return ids, nil
|
|
}
|
|
|
|
func (f walImplsTestFramework) testRead(ctx context.Context, w WALImpls, name string) ([]message.ImmutableMessage, error) {
|
|
s, err := w.Read(ctx, ReadOption{
|
|
Name: name,
|
|
DeliverPolicy: options.DeliverPolicyAll(),
|
|
})
|
|
assert.NoError(f.t, err)
|
|
assert.Equal(f.t, name, s.Name())
|
|
defer s.Close()
|
|
|
|
msgs := make([]message.ImmutableMessage, 0, f.messageCount)
|
|
for i := 0; i < f.messageCount; i++ {
|
|
msg, ok := <-s.Chan()
|
|
assert.NotNil(f.t, msg)
|
|
assert.True(f.t, ok)
|
|
msgs = append(msgs, msg)
|
|
}
|
|
return msgs, nil
|
|
}
|
|
|
|
type sortByMessageID []message.ImmutableMessage
|
|
|
|
func (a sortByMessageID) Len() int {
|
|
return len(a)
|
|
}
|
|
|
|
func (a sortByMessageID) Swap(i, j int) {
|
|
a[i], a[j] = a[j], a[i]
|
|
}
|
|
|
|
func (a sortByMessageID) Less(i, j int) bool {
|
|
return a[i].MessageID().LT(a[j].MessageID())
|
|
}
|