milvus/internal/querynodev2/pipeline/pipeline.go
jaime c9d0c157ec
Move some modules from internal to public package (#22572)
Signed-off-by: jaime <yun.zhang@zilliz.com>
2023-04-06 19:14:32 +08:00

83 lines
2.8 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 pipeline
import (
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/querynodev2/delegator"
base "github.com/milvus-io/milvus/internal/util/pipeline"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/metrics"
"github.com/milvus-io/milvus/pkg/mq/msgdispatcher"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/milvus-io/milvus/pkg/util/typeutil"
)
//pipeline used for querynode
type Pipeline interface {
base.StreamPipeline
ExcludedSegments(segInfos ...*datapb.SegmentInfo)
}
type pipeline struct {
base.StreamPipeline
excludedSegments *typeutil.ConcurrentMap[int64, *datapb.SegmentInfo]
collectionID UniqueID
}
func (p *pipeline) ExcludedSegments(segInfos ...*datapb.SegmentInfo) {
for _, segInfo := range segInfos {
log.Debug("pipeline add exclude info",
zap.Int64("segmentID", segInfo.GetID()),
zap.Uint64("ts", segInfo.GetDmlPosition().GetTimestamp()),
)
p.excludedSegments.Insert(segInfo.GetID(), segInfo)
}
}
func (p *pipeline) Close() {
p.StreamPipeline.Close()
metrics.CleanupQueryNodeCollectionMetrics(paramtable.GetNodeID(), p.collectionID)
}
func NewPipeLine(
collectionID UniqueID,
channel string,
manager *DataManager,
tSafeManager TSafeManager,
dispatcher msgdispatcher.Client,
delegator delegator.ShardDelegator,
) (Pipeline, error) {
pipelineQueueLength := paramtable.Get().QueryNodeCfg.FlowGraphMaxQueueLength.GetAsInt32()
excludedSegments := typeutil.NewConcurrentMap[int64, *datapb.SegmentInfo]()
p := &pipeline{
collectionID: collectionID,
excludedSegments: excludedSegments,
StreamPipeline: base.NewPipelineWithStream(dispatcher, nodeCtxTtInterval, enableTtChecker, channel),
}
filterNode := newFilterNode(collectionID, channel, manager, excludedSegments, pipelineQueueLength)
insertNode := newInsertNode(collectionID, channel, manager, delegator, pipelineQueueLength)
deleteNode := newDeleteNode(collectionID, channel, manager, tSafeManager, delegator, pipelineQueueLength)
p.Add(filterNode, insertNode, deleteNode)
return p, nil
}