milvus/internal/datanode/syncmgr/key_lock_dispatcher.go
congqixia a06f601c6e
fix: Make syncmgr lock key before returning future (#32865)
See also #32860

SyncMgr did not ensure task key is locked before `SyncData` returning
which may cause concurrent problem during sync wich multiple policies.

This PR change sync mgr implementation to make sure the key is locked
before returning task result `*conc.Future`

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2024-05-09 10:09:30 +08:00

46 lines
1.1 KiB
Go

package syncmgr
import (
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
"github.com/milvus-io/milvus/pkg/util/conc"
"github.com/milvus-io/milvus/pkg/util/lock"
)
type Task interface {
SegmentID() int64
CalcTargetSegment() (int64, error)
Checkpoint() *msgpb.MsgPosition
StartPosition() *msgpb.MsgPosition
ChannelName() string
Run() error
HandleError(error)
}
type keyLockDispatcher[K comparable] struct {
keyLock *lock.KeyLock[K]
workerPool *conc.Pool[struct{}]
}
func newKeyLockDispatcher[K comparable](maxParallel int) *keyLockDispatcher[K] {
dispatcher := &keyLockDispatcher[K]{
workerPool: conc.NewPool[struct{}](maxParallel, conc.WithPreAlloc(false)),
keyLock: lock.NewKeyLock[K](),
}
return dispatcher
}
func (d *keyLockDispatcher[K]) Submit(key K, t Task, callbacks ...func(error) error) *conc.Future[struct{}] {
d.keyLock.Lock(key)
return d.workerPool.Submit(func() (struct{}, error) {
defer d.keyLock.Unlock(key)
err := t.Run()
for _, callback := range callbacks {
err = callback(err)
}
return struct{}{}, err
})
}