mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-30 07:25:37 +08:00
99dff06391
issue: #33285 - using streaming service in insert/upsert/flush/delete/querynode - fixup flusher bugs and refactor the flush operation - enable streaming service for dml and ddl - pass the e2e when enabling streaming service - pass the integration tst when enabling streaming service --------- Signed-off-by: chyezh <chyezh@outlook.com> Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
130 lines
3.6 KiB
Go
130 lines
3.6 KiB
Go
package streaming
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/mock"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
|
|
"github.com/milvus-io/milvus/internal/distributed/streaming/internal/producer"
|
|
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/mock_client"
|
|
"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_producer"
|
|
"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/mock_handler"
|
|
"github.com/milvus-io/milvus/pkg/streaming/util/message"
|
|
"github.com/milvus-io/milvus/pkg/streaming/util/types"
|
|
"github.com/milvus-io/milvus/pkg/streaming/walimpls/impls/walimplstest"
|
|
"github.com/milvus-io/milvus/pkg/util/conc"
|
|
"github.com/milvus-io/milvus/pkg/util/lifetime"
|
|
)
|
|
|
|
const (
|
|
vChannel1 = "by-dev-rootcoord-dml_1"
|
|
vChannel2 = "by-dev-rootcoord-dml_2"
|
|
vChannel3 = "by-dev-rootcoord-dml_3"
|
|
)
|
|
|
|
func TestWAL(t *testing.T) {
|
|
coordClient := mock_client.NewMockClient(t)
|
|
coordClient.EXPECT().Close().Return()
|
|
handler := mock_handler.NewMockHandlerClient(t)
|
|
handler.EXPECT().Close().Return()
|
|
|
|
w := &walAccesserImpl{
|
|
lifetime: lifetime.NewLifetime(lifetime.Working),
|
|
streamingCoordAssignmentClient: coordClient,
|
|
handlerClient: handler,
|
|
producerMutex: sync.Mutex{},
|
|
producers: make(map[string]*producer.ResumableProducer),
|
|
appendExecutionPool: conc.NewPool[struct{}](10),
|
|
dispatchExecutionPool: conc.NewPool[struct{}](10),
|
|
}
|
|
defer w.Close()
|
|
|
|
ctx := context.Background()
|
|
|
|
available := make(chan struct{})
|
|
p := mock_producer.NewMockProducer(t)
|
|
p.EXPECT().IsAvailable().RunAndReturn(func() bool {
|
|
select {
|
|
case <-available:
|
|
return false
|
|
default:
|
|
return true
|
|
}
|
|
})
|
|
p.EXPECT().Produce(mock.Anything, mock.Anything).Return(&types.AppendResult{
|
|
MessageID: walimplstest.NewTestMessageID(1),
|
|
TimeTick: 10,
|
|
TxnCtx: &message.TxnContext{
|
|
TxnID: 1,
|
|
Keepalive: 10 * time.Second,
|
|
},
|
|
}, nil)
|
|
p.EXPECT().Available().Return(available)
|
|
p.EXPECT().Close().Return()
|
|
|
|
handler.EXPECT().CreateProducer(mock.Anything, mock.Anything).Return(p, nil)
|
|
result, err := w.RawAppend(ctx, newInsertMessage(vChannel1))
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
|
|
// Test committed txn.
|
|
txn, err := w.Txn(ctx, TxnOption{
|
|
VChannel: vChannel1,
|
|
Keepalive: 10 * time.Second,
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, txn)
|
|
|
|
err = txn.Append(ctx, newInsertMessage(vChannel1))
|
|
assert.NoError(t, err)
|
|
err = txn.Append(ctx, newInsertMessage(vChannel1))
|
|
assert.NoError(t, err)
|
|
|
|
result, err = txn.Commit(ctx)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
|
|
// Test rollback txn.
|
|
txn, err = w.Txn(ctx, TxnOption{
|
|
VChannel: vChannel1,
|
|
Keepalive: 10 * time.Second,
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, txn)
|
|
|
|
err = txn.Append(ctx, newInsertMessage(vChannel1))
|
|
assert.NoError(t, err)
|
|
err = txn.Append(ctx, newInsertMessage(vChannel1))
|
|
assert.NoError(t, err)
|
|
|
|
err = txn.Rollback(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
resp := w.AppendMessages(ctx,
|
|
newInsertMessage(vChannel1),
|
|
newInsertMessage(vChannel2),
|
|
newInsertMessage(vChannel2),
|
|
newInsertMessage(vChannel3),
|
|
newInsertMessage(vChannel3),
|
|
newInsertMessage(vChannel3),
|
|
)
|
|
assert.NoError(t, resp.UnwrapFirstError())
|
|
}
|
|
|
|
func newInsertMessage(vChannel string) message.MutableMessage {
|
|
msg, err := message.NewInsertMessageBuilderV1().
|
|
WithVChannel(vChannel).
|
|
WithHeader(&message.InsertMessageHeader{}).
|
|
WithBody(&msgpb.InsertRequest{}).
|
|
BuildMutable()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return msg
|
|
}
|