mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-08 01:58:34 +08:00
106 lines
3.2 KiB
Go
106 lines
3.2 KiB
Go
package writebuffer
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"go.uber.org/zap"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
"github.com/milvus-io/milvus/internal/datanode/metacache"
|
|
"github.com/milvus-io/milvus/internal/datanode/syncmgr"
|
|
"github.com/milvus-io/milvus/pkg/log"
|
|
"github.com/milvus-io/milvus/pkg/mq/msgstream"
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
|
)
|
|
|
|
// Manager is the interface for WriteBuffer management.
|
|
type Manager interface {
|
|
// Register adds a WriteBuffer with provided schema & options.
|
|
Register(channel string, schema *schemapb.CollectionSchema, metacache metacache.MetaCache, opts ...WriteBufferOption) error
|
|
// FlushSegments notifies writeBuffer corresponding to provided channel to flush segments.
|
|
FlushSegments(ctx context.Context, channel string, segmentIDs []int64) error
|
|
// RemoveChannel removes a write buffer from manager.
|
|
RemoveChannel(channel string)
|
|
// BufferData put data into channel write buffer.
|
|
BufferData(channel string, insertMsgs []*msgstream.InsertMsg, deleteMsgs []*msgstream.DeleteMsg, startPos, endPos *msgpb.MsgPosition) error
|
|
}
|
|
|
|
// NewManager returns initialized manager as `Manager`
|
|
func NewManager(syncMgr syncmgr.SyncManager) Manager {
|
|
return &manager{
|
|
syncMgr: syncMgr,
|
|
buffers: make(map[string]WriteBuffer),
|
|
}
|
|
}
|
|
|
|
type manager struct {
|
|
syncMgr syncmgr.SyncManager
|
|
buffers map[string]WriteBuffer
|
|
mut sync.RWMutex
|
|
}
|
|
|
|
// Register a new WriteBuffer for channel.
|
|
func (m *manager) Register(channel string, schema *schemapb.CollectionSchema, metacache metacache.MetaCache, opts ...WriteBufferOption) error {
|
|
m.mut.Lock()
|
|
defer m.mut.Unlock()
|
|
|
|
_, ok := m.buffers[channel]
|
|
if ok {
|
|
return merr.WrapErrChannelReduplicate(channel)
|
|
}
|
|
buf, err := NewWriteBuffer(schema, metacache, m.syncMgr, opts...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
m.buffers[channel] = buf
|
|
return nil
|
|
}
|
|
|
|
// FlushSegments call sync segment and change segments state to Flushed.
|
|
func (m *manager) FlushSegments(ctx context.Context, channel string, segmentIDs []int64) error {
|
|
m.mut.RLock()
|
|
buf, ok := m.buffers[channel]
|
|
m.mut.RUnlock()
|
|
|
|
if !ok {
|
|
log.Ctx(ctx).Warn("write buffer not found when flush segments",
|
|
zap.String("channel", channel),
|
|
zap.Int64s("segmentIDs", segmentIDs))
|
|
return merr.WrapErrChannelNotFound(channel)
|
|
}
|
|
|
|
return buf.FlushSegments(ctx, segmentIDs)
|
|
}
|
|
|
|
// BufferData put data into channel write buffer.
|
|
func (m *manager) BufferData(channel string, insertMsgs []*msgstream.InsertMsg, deleteMsgs []*msgstream.DeleteMsg, startPos, endPos *msgpb.MsgPosition) error {
|
|
m.mut.RLock()
|
|
buf, ok := m.buffers[channel]
|
|
m.mut.RUnlock()
|
|
|
|
if !ok {
|
|
log.Ctx(context.Background()).Warn("write buffer not found when buffer data",
|
|
zap.String("channel", channel))
|
|
return merr.WrapErrChannelNotFound(channel)
|
|
}
|
|
|
|
return buf.BufferData(insertMsgs, deleteMsgs, startPos, endPos)
|
|
}
|
|
|
|
// RemoveChannel remove channel WriteBuffer from manager.
|
|
func (m *manager) RemoveChannel(channel string) {
|
|
m.mut.Lock()
|
|
buf, ok := m.buffers[channel]
|
|
delete(m.buffers, channel)
|
|
m.mut.Unlock()
|
|
|
|
if !ok {
|
|
log.Warn("failed to remove channel, channel not maintained in manager", zap.String("channel", channel))
|
|
return
|
|
}
|
|
|
|
buf.Close()
|
|
}
|