2023-11-04 12:10:17 +08:00
|
|
|
package writebuffer
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
|
|
|
|
"go.uber.org/zap"
|
|
|
|
|
2023-11-17 21:46:20 +08:00
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
2023-11-04 12:10:17 +08:00
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
|
|
|
|
"github.com/milvus-io/milvus/internal/allocator"
|
|
|
|
"github.com/milvus-io/milvus/internal/datanode/metacache"
|
|
|
|
"github.com/milvus-io/milvus/internal/datanode/syncmgr"
|
2023-11-17 21:46:20 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
2023-11-04 12:10:17 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
|
|
"github.com/milvus-io/milvus/pkg/log"
|
|
|
|
"github.com/milvus-io/milvus/pkg/mq/msgstream"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util/retry"
|
|
|
|
)
|
|
|
|
|
|
|
|
type l0WriteBuffer struct {
|
|
|
|
*writeBufferBase
|
|
|
|
|
2023-11-17 21:46:20 +08:00
|
|
|
l0Segments map[int64]int64 // partitionID => l0 segment ID
|
|
|
|
l0partition map[int64]int64 // l0 segment id => partition id
|
2023-11-04 12:10:17 +08:00
|
|
|
|
|
|
|
syncMgr syncmgr.SyncManager
|
|
|
|
idAllocator allocator.Interface
|
|
|
|
}
|
|
|
|
|
2023-11-17 21:46:20 +08:00
|
|
|
func NewL0WriteBuffer(channel string, metacache metacache.MetaCache, syncMgr syncmgr.SyncManager, option *writeBufferOption) (WriteBuffer, error) {
|
2023-11-04 12:10:17 +08:00
|
|
|
if option.idAllocator == nil {
|
|
|
|
return nil, merr.WrapErrServiceInternal("id allocator is nil when creating l0 write buffer")
|
|
|
|
}
|
|
|
|
return &l0WriteBuffer{
|
|
|
|
l0Segments: make(map[int64]int64),
|
2023-11-17 21:46:20 +08:00
|
|
|
l0partition: make(map[int64]int64),
|
|
|
|
writeBufferBase: newWriteBufferBase(channel, metacache, syncMgr, option),
|
2023-11-04 12:10:17 +08:00
|
|
|
syncMgr: syncMgr,
|
|
|
|
idAllocator: option.idAllocator,
|
|
|
|
}, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (wb *l0WriteBuffer) BufferData(insertMsgs []*msgstream.InsertMsg, deleteMsgs []*msgstream.DeleteMsg, startPos, endPos *msgpb.MsgPosition) error {
|
|
|
|
wb.mut.Lock()
|
|
|
|
defer wb.mut.Unlock()
|
|
|
|
|
|
|
|
// process insert msgs
|
|
|
|
_, err := wb.bufferInsert(insertMsgs, startPos, endPos)
|
|
|
|
if err != nil {
|
|
|
|
log.Warn("failed to buffer insert data", zap.Error(err))
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, msg := range deleteMsgs {
|
2023-11-17 21:46:20 +08:00
|
|
|
l0SegmentID := wb.getL0SegmentID(msg.GetPartitionID(), startPos)
|
2023-11-04 12:10:17 +08:00
|
|
|
pks := storage.ParseIDs2PrimaryKeys(msg.GetPrimaryKeys())
|
|
|
|
err := wb.bufferDelete(l0SegmentID, pks, msg.GetTimestamps(), startPos, endPos)
|
|
|
|
if err != nil {
|
|
|
|
log.Warn("failed to buffer delete data", zap.Error(err))
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// update buffer last checkpoint
|
|
|
|
wb.checkpoint = endPos
|
|
|
|
|
2023-11-17 21:46:20 +08:00
|
|
|
segmentsSync := wb.triggerSync()
|
|
|
|
for _, segment := range segmentsSync {
|
|
|
|
partition, ok := wb.l0partition[segment]
|
|
|
|
if ok {
|
|
|
|
delete(wb.l0partition, segment)
|
|
|
|
delete(wb.l0Segments, partition)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return nil
|
2023-11-04 12:10:17 +08:00
|
|
|
}
|
|
|
|
|
2023-11-17 21:46:20 +08:00
|
|
|
func (wb *l0WriteBuffer) getL0SegmentID(partitionID int64, startPos *msgpb.MsgPosition) int64 {
|
2023-11-04 12:10:17 +08:00
|
|
|
segmentID, ok := wb.l0Segments[partitionID]
|
|
|
|
if !ok {
|
|
|
|
err := retry.Do(context.Background(), func() error {
|
|
|
|
var err error
|
|
|
|
segmentID, err = wb.idAllocator.AllocOne()
|
|
|
|
return err
|
|
|
|
})
|
|
|
|
if err != nil {
|
|
|
|
log.Error("failed to allocate l0 segment ID", zap.Error(err))
|
|
|
|
panic(err)
|
|
|
|
}
|
|
|
|
wb.l0Segments[partitionID] = segmentID
|
2023-11-17 21:46:20 +08:00
|
|
|
wb.l0partition[segmentID] = partitionID
|
|
|
|
wb.metaCache.AddSegment(&datapb.SegmentInfo{
|
|
|
|
ID: segmentID,
|
|
|
|
PartitionID: partitionID,
|
|
|
|
CollectionID: wb.collectionID,
|
|
|
|
InsertChannel: wb.channelName,
|
|
|
|
StartPosition: startPos,
|
|
|
|
State: commonpb.SegmentState_Growing,
|
|
|
|
Level: datapb.SegmentLevel_L0,
|
|
|
|
}, func(_ *datapb.SegmentInfo) *metacache.BloomFilterSet { return metacache.NewBloomFilterSet() }, metacache.SetStartPosRecorded(false))
|
2023-11-04 12:10:17 +08:00
|
|
|
}
|
|
|
|
return segmentID
|
|
|
|
}
|