milvus/internal/rootcoord/dml_channels.go
Bingyi Sun 92eff1565c
Fix unlocked mutex in rootcoord (#12772)
issue: #12297
Signed-off-by: sunby <bingyi.sun@zilliz.com>

Co-authored-by: sunby <bingyi.sun@zilliz.com>
2021-12-06 17:25:35 +08:00

185 lines
4.5 KiB
Go

// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
package rootcoord
import (
"context"
"fmt"
"sync"
"go.uber.org/atomic"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/msgstream"
)
type dmlMsgStream struct {
ms msgstream.MsgStream
mutex sync.RWMutex
refcnt int64
}
type dmlChannels struct {
ctx context.Context
factory msgstream.Factory
namePrefix string
capacity int64
idx *atomic.Int64
pool sync.Map
}
func newDmlChannels(ctx context.Context, factory msgstream.Factory, chanNamePrefix string, chanNum int64) *dmlChannels {
d := &dmlChannels{
ctx: ctx,
factory: factory,
namePrefix: chanNamePrefix,
capacity: chanNum,
idx: atomic.NewInt64(0),
pool: sync.Map{},
}
for i := int64(0); i < chanNum; i++ {
name := genChannelName(d.namePrefix, i)
ms, err := factory.NewMsgStream(ctx)
if err != nil {
log.Error("Failed to add msgstream", zap.String("name", name), zap.Error(err))
panic("Failed to add msgstream")
}
d.pool.Store(name, &dmlMsgStream{
ms: ms,
mutex: sync.RWMutex{},
refcnt: 0,
})
}
log.Debug("init dml channels", zap.Int64("num", chanNum))
return d
}
func (d *dmlChannels) getChannelName() string {
cnt := d.idx.Inc()
return genChannelName(d.namePrefix, (cnt-1)%d.capacity)
}
func (d *dmlChannels) listChannels() []string {
var chanNames []string
d.pool.Range(
func(k, v interface{}) bool {
dms := v.(*dmlMsgStream)
dms.mutex.RLock()
if dms.refcnt > 0 {
chanNames = append(chanNames, k.(string))
}
dms.mutex.RUnlock()
return true
})
return chanNames
}
func (d *dmlChannels) getChannelNum() int {
return len(d.listChannels())
}
func (d *dmlChannels) broadcast(chanNames []string, pack *msgstream.MsgPack) error {
for _, chanName := range chanNames {
v, ok := d.pool.Load(chanName)
if !ok {
log.Error("invalid channel name", zap.String("chanName", chanName))
panic("invalid channel name: " + chanName)
}
dms := v.(*dmlMsgStream)
dms.mutex.RLock()
if dms.refcnt > 0 {
if err := dms.ms.Broadcast(pack); err != nil {
log.Error("Broadcast failed", zap.String("chanName", chanName))
dms.mutex.RUnlock()
return err
}
}
dms.mutex.RUnlock()
}
return nil
}
func (d *dmlChannels) broadcastMark(chanNames []string, pack *msgstream.MsgPack) (map[string][]byte, error) {
result := make(map[string][]byte)
for _, chanName := range chanNames {
v, ok := d.pool.Load(chanName)
if !ok {
log.Error("invalid channel name", zap.String("chanName", chanName))
panic("invalid channel name: " + chanName)
}
dms := v.(*dmlMsgStream)
dms.mutex.RLock()
if dms.refcnt > 0 {
ids, err := dms.ms.BroadcastMark(pack)
if err != nil {
log.Error("BroadcastMark failed", zap.String("chanName", chanName))
dms.mutex.RUnlock()
return result, err
}
for cn, idList := range ids {
// idList should have length 1, just flat by iteration
for _, id := range idList {
result[cn] = id.Serialize()
}
}
}
dms.mutex.RUnlock()
}
return result, nil
}
func (d *dmlChannels) addChannels(names ...string) {
for _, name := range names {
v, ok := d.pool.Load(name)
if !ok {
log.Error("invalid channel name", zap.String("chanName", name))
panic("invalid channel name: " + name)
}
dms := v.(*dmlMsgStream)
dms.mutex.Lock()
if dms.refcnt == 0 {
dms.ms.AsProducer([]string{name})
}
dms.refcnt++
dms.mutex.Unlock()
}
}
func (d *dmlChannels) removeChannels(names ...string) {
for _, name := range names {
v, ok := d.pool.Load(name)
if !ok {
log.Error("invalid channel name", zap.String("chanName", name))
panic("invalid channel name: " + name)
}
dms := v.(*dmlMsgStream)
dms.mutex.Lock()
if dms.refcnt > 0 {
dms.refcnt--
if dms.refcnt == 0 {
dms.ms.Close()
}
}
dms.mutex.Unlock()
}
}
func genChannelName(prefix string, idx int64) string {
return fmt.Sprintf("%s_%d", prefix, idx)
}