未验证 提交 1357ef70 编写于 作者: C congqixia 提交者: GitHub

Merge get worker process for same nodeID (#25185)

Signed-off-by: NCongqi Xia <congqi.xia@zilliz.com>
上级 63b6a4a6
......@@ -18,10 +18,12 @@ package cluster
import (
"fmt"
"strconv"
"go.uber.org/zap"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/conc"
"github.com/milvus-io/milvus/pkg/util/typeutil"
)
......@@ -36,6 +38,7 @@ type WorkerBuilder func(nodeID int64) (Worker, error)
type grpcWorkerManager struct {
workers *typeutil.ConcurrentMap[int64, Worker]
builder WorkerBuilder
sf conc.Singleflight[Worker] //singleflight.Group
}
// GetWorker returns worker with specified nodeID.
......@@ -43,20 +46,25 @@ func (m *grpcWorkerManager) GetWorker(nodeID int64) (Worker, error) {
worker, ok := m.workers.Get(nodeID)
var err error
if !ok {
//TODO merge request?
worker, err = m.builder(nodeID)
worker, err, _ = m.sf.Do(strconv.FormatInt(nodeID, 10), func() (Worker, error) {
worker, err = m.builder(nodeID)
if err != nil {
log.Warn("failed to build worker",
zap.Int64("nodeID", nodeID),
zap.Error(err),
)
return nil, err
}
old, exist := m.workers.GetOrInsert(nodeID, worker)
if exist {
worker.Stop()
worker = old
}
return worker, nil
})
if err != nil {
log.Warn("failed to build worker",
zap.Int64("nodeID", nodeID),
zap.Error(err),
)
return nil, err
}
old, exist := m.workers.GetOrInsert(nodeID, worker)
if exist {
worker.Stop()
worker = old
}
}
if !worker.IsHealthy() {
// TODO wrap error
......
package conc
import "golang.org/x/sync/singleflight"
// Singleflight wraps golang.org/x/sync/singleflight.Group into generic one.
type Singleflight[T any] struct {
internal singleflight.Group
}
// SingleflightResult is a generic Result wrapper for DoChan.
type SingleflightResult[T any] struct {
Val T
Err error
Shared bool
}
func (sf *Singleflight[T]) Do(key string, fn func() (T, error)) (T, error, bool) {
raw, err, shared := sf.internal.Do(key, func() (any, error) {
return fn()
})
var t T
if raw != nil {
t = raw.(T)
}
return t, err, shared
}
func (sf *Singleflight[T]) DoChan(key string, fn func() (T, error)) <-chan SingleflightResult[T] {
ch := make(chan SingleflightResult[T], 1)
go func() {
val, err, shared := sf.Do(key, fn)
ch <- SingleflightResult[T]{
Val: val,
Err: err,
Shared: shared,
}
}()
return ch
}
func (sf *Singleflight[T]) Forget(key string) {
sf.internal.Forget(key)
}
package conc
import (
"sync"
"testing"
"github.com/stretchr/testify/suite"
"go.uber.org/atomic"
)
type SingleflightSuite struct {
suite.Suite
}
func (s *SingleflightSuite) TestDo() {
counter, hasShared := atomic.Int32{}, atomic.Bool{}
sf := Singleflight[any]{}
ch := make(chan struct{})
var wg sync.WaitGroup
wg.Add(10)
for i := 0; i < 10; i++ {
go func(i int) {
defer wg.Done()
_, _, shared := sf.Do("test_do", func() (any, error) {
<-ch
counter.Add(1)
return struct{}{}, nil
})
if shared {
hasShared.Store(true)
}
}(i)
}
close(ch)
wg.Wait()
if hasShared.Load() {
s.Less(counter.Load(), int32(10))
}
}
func (s *SingleflightSuite) TestDoChan() {
counter, hasShared := atomic.Int32{}, atomic.Bool{}
sf := Singleflight[any]{}
ch := make(chan struct{})
var wg sync.WaitGroup
wg.Add(10)
for i := 0; i < 10; i++ {
go func(i int) {
defer wg.Done()
ch := sf.DoChan("test_dochan", func() (any, error) {
<-ch
counter.Add(1)
return struct{}{}, nil
})
result := <-ch
if result.Shared {
hasShared.Store(true)
}
}(i)
}
close(ch)
wg.Wait()
if hasShared.Load() {
s.Less(counter.Load(), int32(10))
}
}
func (s *SingleflightSuite) TestForget() {
counter, hasShared := atomic.Int32{}, atomic.Bool{}
sf := Singleflight[any]{}
ch := make(chan struct{})
var wg sync.WaitGroup
wg.Add(10)
for i := 0; i < 10; i++ {
go func(i int) {
defer wg.Done()
_, _, shared := sf.Do("test_forget", func() (any, error) {
<-ch
counter.Add(1)
return struct{}{}, nil
})
if shared {
hasShared.Store(true)
}
}(i)
}
flag := false
sf.Forget("test_forget")
sf.Do("test_forget", func() (any, error) {
flag = true
return struct{}{}, nil
})
close(ch)
wg.Wait()
s.True(flag)
}
func TestSingleFlight(t *testing.T) {
suite.Run(t, new(SingleflightSuite))
}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册