Files
wehub-resource-sync 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
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

70 lines
2.5 KiB
Go

package utility
import (
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus/internal/util/streamingutil/status"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
var ErrTimeTickVoilation = errors.New("time tick violation")
// ReOrderByTimeTickBuffer is a buffer that stores messages and pops them in order of time tick.
type ReOrderByTimeTickBuffer struct {
messageIDs typeutil.Set[string] // After enabling write ahead buffer, we has two stream to consume,
// write ahead buffer works with the timetick order, but the walscannerimpl works with the message order.
// so repeated message may generate when the swithing between the two stream.
// The deduplicate is used to avoid the repeated message.
messageHeap typeutil.Heap[message.ImmutableMessage]
lastPopTimeTick uint64
bytes int
}
// NewReOrderBuffer creates a new ReOrderBuffer.
func NewReOrderBuffer() *ReOrderByTimeTickBuffer {
return &ReOrderByTimeTickBuffer{
messageIDs: typeutil.NewSet[string](),
messageHeap: typeutil.NewHeap[message.ImmutableMessage](&immutableMessageHeap{}),
}
}
// Push pushes a message into the buffer.
func (r *ReOrderByTimeTickBuffer) Push(msg message.ImmutableMessage) error {
// !!! Drop the unexpected broken timetick rule message.
// It will be enabled until the first timetick coming.
if msg.TimeTick() < r.lastPopTimeTick {
return errors.Wrapf(ErrTimeTickVoilation, "message time tick is less than last pop time tick: %d", r.lastPopTimeTick)
}
msgID := msg.MessageID().Marshal()
if r.messageIDs.Contain(msgID) {
return status.NewInner("message is duplicated: %s", msgID)
}
r.messageHeap.Push(msg)
r.messageIDs.Insert(msgID)
r.bytes += msg.EstimateSize()
return nil
}
// PopUtilTimeTick pops all messages whose time tick is less than or equal to the given time tick.
// The result is sorted by time tick in ascending order.
func (r *ReOrderByTimeTickBuffer) PopUtilTimeTick(timetick uint64) []message.ImmutableMessage {
var res []message.ImmutableMessage
for r.messageHeap.Len() > 0 && r.messageHeap.Peek().TimeTick() <= timetick {
r.bytes -= r.messageHeap.Peek().EstimateSize()
r.messageIDs.Remove(r.messageHeap.Peek().MessageID().Marshal())
res = append(res, r.messageHeap.Pop())
}
r.lastPopTimeTick = timetick
return res
}
// Len returns the number of messages in the buffer.
func (r *ReOrderByTimeTickBuffer) Len() int {
return r.messageHeap.Len()
}
func (r *ReOrderByTimeTickBuffer) Bytes() int {
return r.bytes
}