Zhen Ye 19e5e9f910
enhance: broadcaster will lock resource until message acked (#44508)
issue: #43897

- Return LastConfirmedMessageID when wal append operation.
- Add resource-key-based locker for broadcast-ack operation to protect
the coord state when executing ddl.
- Resource-key-based locker is held until the broadcast operation is
acked.
- ResourceKey support shared and exclusive lock.
- Add FastAck execute ack right away after the broadcast done to speed
up ddl.
- Ack callback will support broadcast message result now.
- Add tombstone for broadcaster to avoid to repeatedly commit DDL and
ABA issue.

---------

Signed-off-by: chyezh <chyezh@outlook.com>
2025-09-24 20:58:05 +08:00

104 lines
3.2 KiB
Go

package utility
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/protobuf/types/known/anypb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/walimplstest"
)
func TestWithNotPersisted(t *testing.T) {
ctx := context.Background()
hint := &NotPersistedHint{MessageID: walimplstest.NewTestMessageID(1)}
ctx = WithNotPersisted(ctx, hint)
retrievedHint := GetNotPersisted(ctx)
assert.NotNil(t, retrievedHint)
assert.True(t, retrievedHint.MessageID.EQ(hint.MessageID))
}
func TestWithExtraAppendResult(t *testing.T) {
ctx := context.Background()
extra := &anypb.Any{}
txnCtx := &message.TxnContext{
TxnID: 1,
}
result := &ExtraAppendResult{TimeTick: 123, TxnCtx: txnCtx, Extra: extra}
ctx = WithExtraAppendResult(ctx, result)
retrievedResult := ctx.Value(extraAppendResultValue).(*ExtraAppendResult)
assert.NotNil(t, retrievedResult)
assert.Equal(t, uint64(123), retrievedResult.TimeTick)
assert.Equal(t, txnCtx.TxnID, retrievedResult.TxnCtx.TxnID)
assert.Equal(t, extra, retrievedResult.Extra)
}
func TestModifyAppendResultExtra(t *testing.T) {
ctx := context.Background()
extra := &anypb.Any{}
result := &ExtraAppendResult{Extra: extra}
ctx = WithExtraAppendResult(ctx, result)
modifier := func(old *anypb.Any) *anypb.Any {
return &anypb.Any{TypeUrl: "modified"}
}
ModifyAppendResultExtra(ctx, modifier)
retrievedResult := ctx.Value(extraAppendResultValue).(*ExtraAppendResult)
assert.Equal(t, retrievedResult.Extra.(*anypb.Any).TypeUrl, "modified")
ModifyAppendResultExtra(ctx, func(old *anypb.Any) *anypb.Any {
return nil
})
retrievedResult = ctx.Value(extraAppendResultValue).(*ExtraAppendResult)
assert.Nil(t, retrievedResult.Extra)
}
func TestReplaceAppendResultTimeTick(t *testing.T) {
ctx := context.Background()
result := &ExtraAppendResult{TimeTick: 1}
ctx = WithExtraAppendResult(ctx, result)
ReplaceAppendResultTimeTick(ctx, 2)
retrievedResult := ctx.Value(extraAppendResultValue).(*ExtraAppendResult)
assert.Equal(t, retrievedResult.TimeTick, uint64(2))
}
func TestReplaceAppendResultTxnContext(t *testing.T) {
ctx := context.Background()
txnCtx := &message.TxnContext{}
result := &ExtraAppendResult{TxnCtx: txnCtx}
ctx = WithExtraAppendResult(ctx, result)
newTxnCtx := &message.TxnContext{TxnID: 2}
ReplaceAppendResultTxnContext(ctx, newTxnCtx)
retrievedResult := ctx.Value(extraAppendResultValue).(*ExtraAppendResult)
assert.Equal(t, retrievedResult.TxnCtx.TxnID, newTxnCtx.TxnID)
}
func TestReplaceAppendResultLastConfirmedMessageID(t *testing.T) {
ctx := context.Background()
result := &ExtraAppendResult{LastConfirmedMessageID: walimplstest.NewTestMessageID(1)}
ctx = WithExtraAppendResult(ctx, result)
newLastConfirmedMessageID := walimplstest.NewTestMessageID(2)
ReplaceAppendResultLastConfirmedMessageID(ctx, newLastConfirmedMessageID)
retrievedResult := ctx.Value(extraAppendResultValue).(*ExtraAppendResult)
assert.True(t, retrievedResult.LastConfirmedMessageID.EQ(newLastConfirmedMessageID))
}
func TestWithFlushFromOldArch(t *testing.T) {
ctx := context.Background()
assert.False(t, GetFlushFromOldArch(ctx))
ctx = WithFlushFromOldArch(ctx)
assert.True(t, GetFlushFromOldArch(ctx))
}