mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 21:39:19 +08:00
a1232fafda
#30633 Signed-off-by: wayblink <anyang.wang@zilliz.com> Co-authored-by: MrPresent-Han <chun.han@zilliz.com>
228 lines
6.6 KiB
Go
228 lines
6.6 KiB
Go
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you 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 datanode
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/samber/lo"
|
|
"go.uber.org/zap"
|
|
|
|
"github.com/milvus-io/milvus/internal/datanode/compaction"
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/log"
|
|
"github.com/milvus-io/milvus/pkg/util/typeutil"
|
|
)
|
|
|
|
const (
|
|
maxTaskNum = 1024
|
|
)
|
|
|
|
type compactionExecutor struct {
|
|
executing *typeutil.ConcurrentMap[int64, compaction.Compactor] // planID to compactor
|
|
completedCompactor *typeutil.ConcurrentMap[int64, compaction.Compactor] // planID to compactor
|
|
completed *typeutil.ConcurrentMap[int64, *datapb.CompactionPlanResult] // planID to CompactionPlanResult
|
|
taskCh chan compaction.Compactor
|
|
dropped *typeutil.ConcurrentSet[string] // vchannel dropped
|
|
|
|
// To prevent concurrency of release channel and compaction get results
|
|
// all released channel's compaction tasks will be discarded
|
|
resultGuard sync.RWMutex
|
|
}
|
|
|
|
func newCompactionExecutor() *compactionExecutor {
|
|
return &compactionExecutor{
|
|
executing: typeutil.NewConcurrentMap[int64, compaction.Compactor](),
|
|
completedCompactor: typeutil.NewConcurrentMap[int64, compaction.Compactor](),
|
|
completed: typeutil.NewConcurrentMap[int64, *datapb.CompactionPlanResult](),
|
|
taskCh: make(chan compaction.Compactor, maxTaskNum),
|
|
dropped: typeutil.NewConcurrentSet[string](),
|
|
}
|
|
}
|
|
|
|
func (c *compactionExecutor) execute(task compaction.Compactor) {
|
|
c.taskCh <- task
|
|
c.toExecutingState(task)
|
|
}
|
|
|
|
func (c *compactionExecutor) toExecutingState(task compaction.Compactor) {
|
|
c.executing.Insert(task.GetPlanID(), task)
|
|
}
|
|
|
|
func (c *compactionExecutor) toCompleteState(task compaction.Compactor) {
|
|
task.Complete()
|
|
c.executing.GetAndRemove(task.GetPlanID())
|
|
}
|
|
|
|
func (c *compactionExecutor) removeTask(planID UniqueID) {
|
|
c.completed.GetAndRemove(planID)
|
|
task, loaded := c.completedCompactor.GetAndRemove(planID)
|
|
if loaded {
|
|
log.Info("Compaction task removed", zap.Int64("planID", planID), zap.String("channel", task.GetChannelName()))
|
|
}
|
|
}
|
|
|
|
// These two func are bounded for waitGroup
|
|
func (c *compactionExecutor) executeWithState(task compaction.Compactor) {
|
|
go c.executeTask(task)
|
|
}
|
|
|
|
func (c *compactionExecutor) start(ctx context.Context) {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case task := <-c.taskCh:
|
|
c.executeWithState(task)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *compactionExecutor) executeTask(task compaction.Compactor) {
|
|
log := log.With(
|
|
zap.Int64("planID", task.GetPlanID()),
|
|
zap.Int64("Collection", task.GetCollection()),
|
|
zap.String("channel", task.GetChannelName()),
|
|
)
|
|
|
|
defer func() {
|
|
c.toCompleteState(task)
|
|
}()
|
|
|
|
log.Info("start to execute compaction")
|
|
|
|
result, err := task.Compact()
|
|
if err != nil {
|
|
log.Warn("compaction task failed", zap.Error(err))
|
|
return
|
|
}
|
|
c.completed.Insert(result.GetPlanID(), result)
|
|
c.completedCompactor.Insert(result.GetPlanID(), task)
|
|
|
|
log.Info("end to execute compaction")
|
|
}
|
|
|
|
func (c *compactionExecutor) stopTask(planID UniqueID) {
|
|
task, loaded := c.executing.GetAndRemove(planID)
|
|
if loaded {
|
|
log.Warn("compaction executor stop task", zap.Int64("planID", planID), zap.String("vChannelName", task.GetChannelName()))
|
|
task.Stop()
|
|
}
|
|
}
|
|
|
|
func (c *compactionExecutor) isValidChannel(channel string) bool {
|
|
// if vchannel marked dropped, compaction should not proceed
|
|
return !c.dropped.Contain(channel)
|
|
}
|
|
|
|
func (c *compactionExecutor) discardByDroppedChannel(channel string) {
|
|
c.dropped.Insert(channel)
|
|
c.discardPlan(channel)
|
|
}
|
|
|
|
func (c *compactionExecutor) discardPlan(channel string) {
|
|
c.resultGuard.Lock()
|
|
defer c.resultGuard.Unlock()
|
|
|
|
c.executing.Range(func(planID int64, task compaction.Compactor) bool {
|
|
if task.GetChannelName() == channel {
|
|
c.stopTask(planID)
|
|
}
|
|
return true
|
|
})
|
|
|
|
// remove all completed plans of channel
|
|
c.completed.Range(func(planID int64, result *datapb.CompactionPlanResult) bool {
|
|
if result.GetChannel() == channel {
|
|
c.removeTask(planID)
|
|
log.Info("remove compaction plan and results",
|
|
zap.String("channel", channel),
|
|
zap.Int64("planID", planID))
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
|
|
func (c *compactionExecutor) getCompactionResult(planID int64) *datapb.CompactionPlanResult {
|
|
c.resultGuard.RLock()
|
|
defer c.resultGuard.RUnlock()
|
|
_, ok := c.executing.Get(planID)
|
|
if ok {
|
|
result := &datapb.CompactionPlanResult{
|
|
State: datapb.CompactionTaskState_executing,
|
|
PlanID: planID,
|
|
}
|
|
return result
|
|
}
|
|
result, ok2 := c.completed.Get(planID)
|
|
if !ok2 {
|
|
return &datapb.CompactionPlanResult{
|
|
PlanID: planID,
|
|
State: datapb.CompactionTaskState_failed,
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (c *compactionExecutor) getAllCompactionResults() []*datapb.CompactionPlanResult {
|
|
c.resultGuard.RLock()
|
|
defer c.resultGuard.RUnlock()
|
|
var (
|
|
executing []int64
|
|
completed []int64
|
|
completedLevelZero []int64
|
|
)
|
|
results := make([]*datapb.CompactionPlanResult, 0)
|
|
// get executing results
|
|
c.executing.Range(func(planID int64, task compaction.Compactor) bool {
|
|
executing = append(executing, planID)
|
|
results = append(results, &datapb.CompactionPlanResult{
|
|
State: datapb.CompactionTaskState_executing,
|
|
PlanID: planID,
|
|
})
|
|
return true
|
|
})
|
|
|
|
// get completed results
|
|
c.completed.Range(func(planID int64, result *datapb.CompactionPlanResult) bool {
|
|
completed = append(completed, planID)
|
|
results = append(results, result)
|
|
|
|
if result.GetType() == datapb.CompactionType_Level0DeleteCompaction {
|
|
completedLevelZero = append(completedLevelZero, planID)
|
|
}
|
|
return true
|
|
})
|
|
|
|
// remove level zero results
|
|
lo.ForEach(completedLevelZero, func(planID int64, _ int) {
|
|
c.completed.Remove(planID)
|
|
c.completedCompactor.Remove(planID)
|
|
})
|
|
|
|
if len(results) > 0 {
|
|
log.Info("DataNode Compaction results",
|
|
zap.Int64s("executing", executing),
|
|
zap.Int64s("completed", completed),
|
|
zap.Int64s("completed levelzero", completedLevelZero),
|
|
)
|
|
}
|
|
|
|
return results
|
|
}
|