Files
milvus/internal/datacoord/services.go
T
e2787d3981 enhance: standardize error handling on merr + Sys/Input classification (#50221)
issue: #47420

## What this PR does

Project-wide migration of raw `fmt.Errorf` / `errors.New` in function
bodies onto
the `merr` framework, plus the Sys-vs-Input error classification and the
machinery it drives (retriability, fine-grained metrics, segcore
unification),
plus the convention docs and a linter that keeps it from regressing.

Scope: storage, proxy, coordinators (root/data/query), query node, data
node,
`pkg/util` & `internal/util`, expression parser, message queue,
streaming, and
misc packages. Bare raw-error usages went from ~3000 to a ~340 allowlist
(package-level sentinels / build-tag / test sites).

---

## How to review this PR

It is large but the vast majority is mechanical. Changes fall into three
tiers;
spend review budget on Part 2 and Part 3.

### Part 1 — Mechanical standardization (low risk, verify by rule)

Each converted call follows one of a small fixed set of rules. To
review, check
that each site obeys the matching rule rather than reading every line:

| Pattern | Rule |
|---|---|
| `fmt.Errorf("...")` originating a new error | →
`merr.WrapErrXxxMsg("...")` with a code matching the failure's meaning |
| Adding context to an existing typed error | → `merr.Wrap(err, "...")`
/ `merr.Wrapf(...)` — **preserves** the inner code (never `WrapErr*Err`,
which overwrites it) |
| Errors inside the streaming subsystem | → `status.New*` factories
(StreamingError), **not** merr — this is the component-internal dialect
(see `docs/dev/error_handling_guide.md`) |
| Low-level / control-flow signal caught by `errors.Is` | → kept as a
package-level `errors.New` sentinel (lowercase, same-package) |

Conventions are documented in `docs/dev/error_handling_guide.md`
(how-to) and
`docs/dev/error_sentinel_convention.md` (rules + audit). A
`gocritic`/`ruleguard`
rule (`rawmerrerror`, in `rules.go`) enforces "no raw `return
errors.New/fmt.Errorf`"
under `make verifiers`.

### Part 2 — Behavior changes (review these closely)

These are the sites where the wire contract or runtime behavior changes,
not just
the source text. Listed by category; representative locations given,
full set in
the diff.

**A. gRPC wire-code shifts: `UnexpectedError(1)/Code 65535` → typed
code.**
Where a handler previously returned a raw error (collapsed to
`Code=65535` on the
wire), it now returns a typed merr, so the client sees a real code. The
most
common shift is to `IllegalArgument(5)/Code 1100` (ParameterInvalid).
Touch
points include datanode task handlers (CreateTask/Query/Drop), proxy
Upsert,
querynode GetMetrics, datacoord CreateIndex, httpserver query-response
builder,
and typeutil schema validation. One code refinement: an index-param
validation
moved `1100` → `1101` (ParameterMissing). **Client/SDK assertions and
any code
that switched on `Code=65535` for these paths must be re-checked** (the
go_client
e2e assertions were already aligned in this PR).

**B. Prometheus `status` label contract change (externally visible).**
The proxy metric's coarse `fail` / `rejected` values are split into
`fail_input` / `fail_system` and `rejected_user` / `rejected_system` (in
`requestutil.ParseMetricLabel`; auth/privilege rejections count as
`rejected_user`), so dashboards can attribute a failure to caller vs
operator.
**Dashboards/alerts querying `status="fail"` must migrate to
`status=~"fail_.*"`, and `status="rejected"` to
`status=~"rejected_.*"`.** The
in-repo Grafana dashboard is already migrated; external dashboards built
on the
old values silently go empty after upgrade. This is the one change that
requires an ops-side migration.

**C. Retriability semantics.**
- C1: `merr.Status(err)` now forces `Retriable=false` when the error is
an
`InputError` — a malformed request can never succeed on blind retry, so
clients
never get the self-contradictory "your input is wrong but you may
retry".
- C2: `retry.Do` short-circuits an `InputError` (non-retriable) — **but
only when
  the caller did not pass a `RetryErr` predicate**. The check is an
`if c.isRetryErr != nil { ... } else if InputError { ... }` *mutually
exclusive*
branch (`pkg/util/retry/retry.go`): an explicit `RetryErr` takes
precedence and
bypasses the InputError abort. `retry.Handle` deliberately does **not**
apply
the InputError abort (its callers signal abort via `shouldRetry=false`).
Four
flusher startup callsites that must retry through transient "not ready"
errors
  were given explicit `RetryErr` escape hatches.

**D. segcore (C++→Go) error classification.**
A single shared Go-side table (`pkg/util/merr/segcore.go`) maps each
segcore code
to a merr sentinel + InputError/signal category, replacing scattered
hand-written
`if errorCode == ...` switches in the cgo wrappers. **Wire `Code` values
change
for every segcore pass-through error, not just the remapped ones.**
Named
sentinels remap (C++ `2003` → merr `2001`, `2033` → `2002`,
Folly/Knowhere codes
likewise); **all remaining pass-through codes (`2004`–`2043`, previously
surfaced to clients as raw C++ enum values) now serialize as `2000`**
(`ErrSegcore`), with the original C++ code preserved in the `Reason`
text
(`segcoreCode=...`); unknown/future codes collapse to `2000` as well
(pinned by
the `wire_code_projection` test). Transient segcore classes (object
storage /
file IO / OOM / mmap / FieldNotLoaded — 11 codes) now report
`Retriable=true`.
**Any client switching on raw segcore codes in the `2004`–`2043` range
must be
re-checked**; the in-Reason code remains available for diagnostics.
Signal
codes (PretendFinished / FollyCancel) are recognized centrally.
`errors.Is`-based
control flow on these (e.g. scheduler skip/retry) is preserved.

**E. InputError classification (25 sentinels + dynamic marks).**
25 sentinels in `errors.go` carry `WithErrorType(InputError)` (the
Collection /
ResourceGroup / Database families, `ErrIndexDuplicate`,
`ErrParameterInvalid`,
`ErrPrivilegeNotAuthenticated`, `ErrImportFailed`, `ErrQueryPlan`, ...),
plus dynamic
marks for the 8 segcore input codes (ExprInvalid, DimNotMatch,
MetricTypeInvalid, FieldIDInvalid, ...) and
`WrapErrAsInputError`. The widest blast radius is `ErrParameterInvalid`
(1100):
~2335 `WrapErrParameterInvalid*` callsites now classify as input /
non-retriable. Because of C1/C2 this changes retriability for
any path that returns these. **The audit to confirm no transient path
was
mis-marked is the single most important review item** (see Part 3). One
reverse
correction: storage field-stats parsing moved from `ErrParameterInvalid`
(input)
to `ErrDataIntegrity` — a corrupted stored stat is data corruption, not
user
input.

### Part 3 — Known risks & traps (called out proactively)

1. **`merr.Wrap` vs `WrapErr*Err` (code-masking).** `WrapErr*Err` builds
a
`wrappedMilvusError{sentinel: ErrServiceInternal}` whose `code()`
returns the
*outer* sentinel — it overwrites the inner typed code and hides the
`errors.Is`
chain. This is intentional (use it to *deliberately* downgrade), but it
was a
recurring conversion defect; the rule "add context with `merr.Wrap`,
downgrade
with `WrapErr*Err`" is enforced by convention and reviewed across the
diff.
2. **InputError × `retry.Do` blast radius.** Marking a sentinel
`InputError` makes
any `retry.Do(...)` without a `RetryErr` predicate stop retrying it.
Reviewers
should sanity-check that no transient use of the 19 newly-marked
sentinels
(especially `ErrParameterInvalid`) sits inside a retry loop that needed
to keep
   spinning. The known flusher cases were handled (see C2).
3. **The ~340 raw-error allowlist.** What remains as bare `errors.New`
is, by
design: package-level sentinels (caught by `errors.Is`), `//go:build
test`
sites, and out-of-band trees (`cmd/`, `tests/`, codegen, walimpls). The
linter
only bans the *direct-return* form; assignment-then-return escapes and
the full
no-exceptions ban are deferred to an AST-based linter (Tier 2,
documented).
4. **segcore C++ second step deferred.** This PR unifies classification
on the Go
side; splitting the dual-semantic C++ codes at the source is a
follow-up.

---

## Validation

- `make verifiers`: Go side clean (gofmt + static-check across modules,
including
  the new `rawmerrerror` rule with a 0-hit baseline repo-wide).
- `make test-go`: passing; the one real regression introduced (a
datanode
`invalid_task_type` assertion shifting `1` → `5` from a ParameterInvalid
  conversion) was fixed in-tree.
- go_client e2e CreateIndex assertions aligned to the new merr messages.

---------

Signed-off-by: zhenshan.cao <zhenshan.cao@zilliz.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-06-12 15:04:51 -07:00

3089 lines
115 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 datacoord
import (
"context"
"fmt"
"math"
"strconv"
"time"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"go.opentelemetry.io/otel"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/internal/coordinator/snmanager"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/metastore/kv/binlog"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster/broadcast"
"github.com/milvus-io/milvus/internal/util/componentutil"
"github.com/milvus-io/milvus/internal/util/importutilv2"
"github.com/milvus-io/milvus/internal/util/segmentutil"
"github.com/milvus-io/milvus/internal/util/streamingutil"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/log"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/metricsinfo"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// GetTimeTickChannel legacy API, returns time tick channel name
func (s *Server) GetTimeTickChannel(ctx context.Context, req *internalpb.GetTimeTickChannelRequest) (*milvuspb.StringResponse, error) {
return &milvuspb.StringResponse{
Status: merr.Success(),
Value: Params.CommonCfg.DataCoordTimeTick.GetValue(),
}, nil
}
// GetStatisticsChannel legacy API, returns statistics channel name
func (s *Server) GetStatisticsChannel(ctx context.Context, req *internalpb.GetStatisticsChannelRequest) (*milvuspb.StringResponse, error) {
return &milvuspb.StringResponse{
Status: merr.Status(merr.WrapErrChannelNotFound("no statistics channel")),
}, nil
}
// Flush notify segment to flush
// this api only guarantees all the segments requested is sealed
// these segments will be flushed only after the Flush policy is fulfilled
func (s *Server) Flush(ctx context.Context, req *datapb.FlushRequest) (*datapb.FlushResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("dbID", req.GetDbID()),
zap.Int64("collectionID", req.GetCollectionID()))
log.Info("receive flush request")
ctx, sp := otel.Tracer(typeutil.DataCoordRole).Start(ctx, "DataCoord-Flush")
defer sp.End()
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.FlushResponse{
Status: merr.Status(err),
}, nil
}
// generate a timestamp timeOfSeal, all data before timeOfSeal is guaranteed to be sealed or flushed
ts, err := s.allocator.AllocTimestamp(ctx)
if err != nil {
log.Warn("unable to alloc timestamp", zap.Error(err))
return nil, err
}
flushResult, err := s.flushCollection(ctx, req.GetCollectionID(), ts, req.GetSegmentIDs())
if err != nil {
return &datapb.FlushResponse{
Status: merr.Status(err),
}, nil
}
return &datapb.FlushResponse{
Status: merr.Success(),
DbID: req.GetDbID(),
CollectionID: req.GetCollectionID(),
SegmentIDs: flushResult.GetSegmentIDs(),
TimeOfSeal: flushResult.GetTimeOfSeal(),
FlushSegmentIDs: flushResult.GetFlushSegmentIDs(),
FlushTs: flushResult.GetFlushTs(),
ChannelCps: flushResult.GetChannelCps(),
}, nil
}
func (s *Server) flushCollection(ctx context.Context, collectionID UniqueID, flushTs uint64, toFlushSegments []UniqueID) (*datapb.FlushResult, error) {
channelCPs := make(map[string]*msgpb.MsgPosition, 0)
coll, err := s.handler.GetCollection(ctx, collectionID)
if err != nil {
log.Warn("fail to get collection", zap.Error(err))
return nil, err
}
if coll == nil {
return nil, merr.WrapErrCollectionNotFound(collectionID)
}
// channel checkpoints must be gotten before sealSegment, make sure checkpoints is earlier than segment's endts
for _, vchannel := range coll.VChannelNames {
cp := s.meta.GetChannelCheckpoint(vchannel)
channelCPs[vchannel] = cp
}
timeOfSeal, _ := tsoutil.ParseTS(flushTs)
sealedSegmentsIDDict := make(map[UniqueID]bool)
if !streamingutil.IsStreamingServiceEnabled() {
for _, channel := range coll.VChannelNames {
sealedSegmentIDs, err := s.segmentManager.SealAllSegments(ctx, channel, toFlushSegments)
if err != nil {
return nil, merr.Wrapf(err, "failed to flush collection %d", collectionID)
}
for _, sealedSegmentID := range sealedSegmentIDs {
sealedSegmentsIDDict[sealedSegmentID] = true
}
}
}
segments := s.meta.GetSegmentsOfCollection(ctx, collectionID)
flushSegmentIDs := make([]UniqueID, 0, len(segments))
for _, segment := range segments {
if segment != nil &&
isFlushState(segment.GetState()) &&
segment.GetLevel() != datapb.SegmentLevel_L0 && // SegmentLevel_Legacy, SegmentLevel_L1, SegmentLevel_L2
!sealedSegmentsIDDict[segment.GetID()] {
flushSegmentIDs = append(flushSegmentIDs, segment.GetID())
}
}
log.Info("flush response with segments",
zap.Int64("collectionID", collectionID),
zap.Int64s("sealSegments", lo.Keys(sealedSegmentsIDDict)),
zap.Int("flushedSegmentsCount", len(flushSegmentIDs)),
zap.Time("timeOfSeal", timeOfSeal),
zap.Uint64("flushTs", flushTs),
zap.Time("flushTs in time", tsoutil.PhysicalTime(flushTs)))
return &datapb.FlushResult{
CollectionID: collectionID,
SegmentIDs: lo.Keys(sealedSegmentsIDDict),
TimeOfSeal: timeOfSeal.Unix(),
FlushSegmentIDs: flushSegmentIDs,
FlushTs: flushTs,
ChannelCps: channelCPs,
DbName: coll.DatabaseName,
CollectionName: coll.Schema.GetName(),
}, nil
}
func (s *Server) FlushAll(ctx context.Context, req *datapb.FlushAllRequest) (*datapb.FlushAllResponse, error) {
log.Ctx(ctx).Info("receive FlushAll request")
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.FlushAllResponse{
Status: merr.Status(err),
}, nil
}
// Create a new broadcaster with exclusive cluster resource key.
broadcaster, err := broadcast.StartBroadcastWithResourceKeys(ctx, message.NewExclusiveClusterResourceKey())
if err != nil {
return &datapb.FlushAllResponse{
Status: merr.Status(err),
}, nil
}
defer broadcaster.Close()
cc := channel.GetClusterChannels()
broadcastFlushAllMsg := message.NewFlushAllMessageBuilderV2().
WithHeader(&message.FlushAllMessageHeader{}).
WithBody(&message.FlushAllMessageBody{}).
WithClusterLevelBroadcast(cc).
MustBuildBroadcast()
res, err := broadcaster.Broadcast(ctx, broadcastFlushAllMsg)
if err != nil {
log.Ctx(ctx).Warn("broadcast FlushAllMessage fail", zap.Error(err))
return &datapb.FlushAllResponse{
Status: merr.Status(err),
}, nil
}
flushAllMsgs := make(map[string]*commonpb.ImmutableMessage, len(res.AppendResults))
msgs := broadcastFlushAllMsg.SplitIntoMutableMessage()
for _, msg := range msgs {
appendResult := res.GetAppendResult(msg.VChannel())
// if is control channel, convert it to physical channel.
pchannel := msg.PChannel()
flushAllMsgs[pchannel] = msg.WithTimeTick(appendResult.TimeTick).
WithLastConfirmed(appendResult.LastConfirmedMessageID).
IntoImmutableMessage(appendResult.MessageID).
IntoImmutableMessageProto()
}
log.Ctx(ctx).Info("FlushAll successfully", log.FieldMessages(msgs))
return &datapb.FlushAllResponse{
Status: merr.Success(),
FlushAllMsgs: flushAllMsgs,
ClusterInfo: &milvuspb.ClusterInfo{
ClusterId: Params.CommonCfg.ClusterID.GetValue(),
Cchannel: cc.ControlChannel,
Pchannels: cc.Channels,
},
}, nil
}
// AssignSegmentID applies for segment ids and make allocation for records.
func (s *Server) AssignSegmentID(ctx context.Context, req *datapb.AssignSegmentIDRequest) (*datapb.AssignSegmentIDResponse, error) {
log := log.Ctx(ctx)
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.AssignSegmentIDResponse{
Status: merr.Status(err),
}, nil
}
assigns := make([]*datapb.SegmentIDAssignment, 0, len(req.SegmentIDRequests))
for _, r := range req.SegmentIDRequests {
log.Info("handle assign segment request",
zap.Int64("collectionID", r.GetCollectionID()),
zap.Int64("partitionID", r.GetPartitionID()),
zap.String("channelName", r.GetChannelName()),
zap.Uint32("count", r.GetCount()),
)
// Load the collection info from Root Coordinator, if it is not found in server meta.
// Note: this request wouldn't be received if collection didn't exist.
_, err := s.handler.GetCollection(ctx, r.GetCollectionID())
if err != nil {
log.Warn("cannot get collection schema", zap.Error(err))
}
// Have segment manager allocate and return the segment allocation info.
segmentAllocations, err := s.segmentManager.AllocSegment(ctx,
r.CollectionID, r.PartitionID, r.ChannelName, int64(r.Count), r.GetStorageVersion())
if err != nil {
log.Warn("failed to alloc segment", zap.Any("request", r), zap.Error(err))
assigns = append(assigns, &datapb.SegmentIDAssignment{
ChannelName: r.ChannelName,
CollectionID: r.CollectionID,
PartitionID: r.PartitionID,
Status: merr.Status(err),
})
continue
}
log.Info("success to assign segments", zap.Int64("collectionID", r.GetCollectionID()), zap.Any("assignments", segmentAllocations))
for _, allocation := range segmentAllocations {
result := &datapb.SegmentIDAssignment{
SegID: allocation.SegmentID,
ChannelName: r.ChannelName,
Count: uint32(allocation.NumOfRows),
CollectionID: r.CollectionID,
PartitionID: r.PartitionID,
ExpireTime: allocation.ExpireTime,
Status: merr.Success(),
}
assigns = append(assigns, result)
}
}
return &datapb.AssignSegmentIDResponse{
Status: merr.Success(),
SegIDAssignments: assigns,
}, nil
}
// AllocSegment alloc a new growing segment, add it into segment meta.
// Only used by Streamingnode, should be deprecated in the future after growing segment fully managed by streaming node.
func (s *Server) AllocSegment(ctx context.Context, req *datapb.AllocSegmentRequest) (*datapb.AllocSegmentResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.AllocSegmentResponse{Status: merr.Status(err)}, nil
}
// !!! SegmentId must be allocated from rootCoord id allocation.
if req.GetCollectionId() == 0 || req.GetPartitionId() == 0 || req.GetVchannel() == "" || req.GetSegmentId() == 0 {
return &datapb.AllocSegmentResponse{Status: merr.Status(merr.ErrParameterInvalid)}, nil
}
// Alloc new growing segment and return the segment info.
segmentInfo, err := s.segmentManager.AllocNewGrowingSegment(
ctx,
AllocNewGrowingSegmentRequest{
CollectionID: req.GetCollectionId(),
PartitionID: req.GetPartitionId(),
SegmentID: req.GetSegmentId(),
ChannelName: req.GetVchannel(),
StorageVersion: req.GetStorageVersion(),
IsCreatedByStreaming: req.GetIsCreatedByStreaming(),
SchemaVersion: req.GetSchemaVersion(),
},
)
if err != nil {
return &datapb.AllocSegmentResponse{Status: merr.Status(err)}, nil
}
clonedSegmentInfo := segmentInfo.Clone()
return &datapb.AllocSegmentResponse{
SegmentInfo: clonedSegmentInfo.SegmentInfo,
Status: merr.Success(),
}, nil
}
// GetSegmentStates returns segments state
func (s *Server) GetSegmentStates(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetSegmentStatesResponse{
Status: merr.Status(err),
}, nil
}
resp := &datapb.GetSegmentStatesResponse{
Status: merr.Success(),
}
for _, segmentID := range req.SegmentIDs {
state := &datapb.SegmentStateInfo{
SegmentID: segmentID,
}
segmentInfo := s.meta.GetHealthySegment(ctx, segmentID)
if segmentInfo == nil {
state.State = commonpb.SegmentState_NotExist
} else {
state.State = segmentInfo.GetState()
state.StartPosition = segmentInfo.GetStartPosition()
}
resp.States = append(resp.States, state)
}
return resp, nil
}
// GetInsertBinlogPaths returns binlog paths info for requested segments
func (s *Server) GetInsertBinlogPaths(ctx context.Context, req *datapb.GetInsertBinlogPathsRequest) (*datapb.GetInsertBinlogPathsResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetInsertBinlogPathsResponse{
Status: merr.Status(err),
}, nil
}
segment := s.meta.GetHealthySegment(ctx, req.GetSegmentID())
if segment == nil {
return &datapb.GetInsertBinlogPathsResponse{
Status: merr.Status(merr.WrapErrSegmentNotFound(req.GetSegmentID())),
}, nil
}
segment = segment.Clone()
err := binlog.DecompressBinLog(storage.InsertBinlog, segment.GetCollectionID(), segment.GetPartitionID(), segment.GetID(), segment.GetBinlogs())
if err != nil {
return &datapb.GetInsertBinlogPathsResponse{
Status: merr.Status(err),
}, nil
}
resp := &datapb.GetInsertBinlogPathsResponse{
Status: merr.Success(),
}
binlogs := segment.GetBinlogs()
fids := make([]UniqueID, 0, len(binlogs))
paths := make([]*internalpb.StringList, 0, len(binlogs))
for _, field := range binlogs {
fids = append(fids, field.GetFieldID())
binlogs := field.GetBinlogs()
p := make([]string, 0, len(binlogs))
for _, log := range binlogs {
p = append(p, log.GetLogPath())
}
paths = append(paths, &internalpb.StringList{Values: p})
}
resp.FieldIDs = fids
resp.Paths = paths
return resp, nil
}
// GetCollectionStatistics returns statistics for collection
// for now only row count is returned
func (s *Server) GetCollectionStatistics(ctx context.Context, req *datapb.GetCollectionStatisticsRequest) (*datapb.GetCollectionStatisticsResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("collectionID", req.GetCollectionID()),
)
log.Info("received request to get collection statistics")
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetCollectionStatisticsResponse{
Status: merr.Status(err),
}, nil
}
resp := &datapb.GetCollectionStatisticsResponse{
Status: merr.Success(),
}
nums := s.meta.GetNumRowsOfCollection(ctx, req.CollectionID)
resp.Stats = append(resp.Stats, &commonpb.KeyValuePair{Key: "row_count", Value: strconv.FormatInt(nums, 10)})
// Calculate schema version consistency proportion
// Only report when schema version > 0 (i.e., AlterCollectionSchema has been called)
collection := s.meta.GetCollection(req.CollectionID)
if collection != nil && collection.Schema != nil && collection.Schema.GetVersion() > 0 {
collectionSchemaVersion := collection.Schema.GetVersion()
// Growing segments are excluded from the consistency gate as a workaround until
// companion PR #48865 lands. Streaming-created growing segments currently carry
// SchemaVersion=0 because the propagation chain in segment_alloc_worker.go and
// msg_handler_impl.go does not yet pass SchemaVersion through. Including them
// would cause the gate to never reach 100% under any write traffic, permanently
// blocking subsequent schema-change DDLs.
//
// This is safe: growing segments will eventually be sealed/flushed, at which
// point the backfill policy picks them up and updates their SchemaVersion. The
// consistency gate only needs to prove that all data eligible for backfill has
// been backfilled — growing segments are not yet eligible.
//
// L0 segments are also excluded: they only contain delete logs, so there is no
// user data to backfill and no schema version consistency to track.
//
// TODO: remove the Growing exclusion once #48865 lands and streaming-created
// segments carry the correct SchemaVersion from creation.
segments := s.meta.SelectSegments(ctx, WithCollection(req.CollectionID), SegmentFilterFunc(func(si *SegmentInfo) bool {
return isSegmentHealthy(si) &&
!si.GetIsImporting() &&
!si.GetIsInvisible() &&
si.GetLevel() != datapb.SegmentLevel_L0 &&
si.GetState() != commonpb.SegmentState_Growing
}))
// When there are no segments the collection is trivially consistent; emit nothing so the
// proxy treats the absent keys as "no backfill in progress" and allows the DDL through.
if len(segments) > 0 {
consistentCount := 0
for _, segment := range segments {
if segment.GetSchemaVersion() == collectionSchemaVersion {
consistentCount++
}
}
log.Info("calculated schema version consistency",
zap.Int32("collectionSchemaVersion", collectionSchemaVersion),
zap.Int("totalSegments", len(segments)),
zap.Int("consistentSegments", consistentCount))
// Emit raw integer counts instead of a floating-point proportion to avoid the rounding
// hazard where e.g. 99999/100000 = 99.999% formats as "100.00" with "%.2f" and would
// falsely satisfy a 100% gate check. Proxy compares these as exact integers.
resp.Stats = append(resp.Stats,
&commonpb.KeyValuePair{
Key: common.SchemaVersionConsistentSegmentsKey,
Value: strconv.Itoa(consistentCount),
},
&commonpb.KeyValuePair{
Key: common.SchemaVersionTotalSegmentsKey,
Value: strconv.Itoa(len(segments)),
},
)
}
}
log.Info("success to get collection statistics", zap.Any("response", resp))
return resp, nil
}
// GetPartitionStatistics returns statistics for partition
// if partID is empty, return statistics for all partitions of the collection
// for now only row count is returned
func (s *Server) GetPartitionStatistics(ctx context.Context, req *datapb.GetPartitionStatisticsRequest) (*datapb.GetPartitionStatisticsResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("collectionID", req.GetCollectionID()),
zap.Int64s("partitionIDs", req.GetPartitionIDs()),
)
resp := &datapb.GetPartitionStatisticsResponse{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetPartitionStatisticsResponse{
Status: merr.Status(err),
}, nil
}
nums := int64(0)
if len(req.GetPartitionIDs()) == 0 {
nums = s.meta.GetNumRowsOfCollection(ctx, req.CollectionID)
}
for _, partID := range req.GetPartitionIDs() {
num := s.meta.GetNumRowsOfPartition(ctx, req.CollectionID, partID)
nums += num
}
resp.Stats = append(resp.Stats, &commonpb.KeyValuePair{Key: "row_count", Value: strconv.FormatInt(nums, 10)})
log.Info("success to get partition statistics", zap.Any("response", resp))
return resp, nil
}
// GetSegmentInfoChannel legacy API, returns segment info statistics channel
func (s *Server) GetSegmentInfoChannel(ctx context.Context, req *datapb.GetSegmentInfoChannelRequest) (*milvuspb.StringResponse, error) {
return &milvuspb.StringResponse{
Status: merr.Success(),
Value: Params.CommonCfg.DataCoordSegmentInfo.GetValue(),
}, nil
}
// GetSegmentInfo returns segment info requested, status, row count, etc included
// Called by: QueryCoord, DataNode, IndexCoord, Proxy.
func (s *Server) GetSegmentInfo(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) {
log := log.Ctx(ctx)
resp := &datapb.GetSegmentInfoResponse{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetSegmentInfoResponse{
Status: merr.Status(err),
}, nil
}
infos := make([]*datapb.SegmentInfo, 0, len(req.GetSegmentIDs()))
channelCPs := make(map[string]*msgpb.MsgPosition)
for _, id := range req.SegmentIDs {
var info *SegmentInfo
if req.IncludeUnHealthy {
info = s.meta.GetSegment(ctx, id)
// info may be not-nil, but ok is false when the segment is being dropped concurrently.
if info == nil {
log.Warn("failed to get segment, this may have been cleaned", zap.Int64("segmentID", id))
err := merr.WrapErrSegmentNotFound(id)
resp.Status = merr.Status(err)
return resp, nil
}
// We should retrieve the deltalog of all child segments,
// but due to the compaction constraint based on indexed segment, there will be at most two generations.
allChildrenDeltalogs, err := s.handler.GetDeltaLogFromCompactTo(ctx, id)
if err != nil {
resp.Status = merr.Status(err)
return resp, nil
}
clonedInfo := info.Clone()
clonedInfo.Deltalogs = append(clonedInfo.Deltalogs, allChildrenDeltalogs...)
segmentutil.ReCalcRowCount(info.SegmentInfo, clonedInfo.SegmentInfo)
infos = append(infos, clonedInfo.SegmentInfo)
} else {
info = s.meta.GetHealthySegment(ctx, id)
if info == nil {
err := merr.WrapErrSegmentNotFound(id)
resp.Status = merr.Status(err)
return resp, nil
}
clonedInfo := info.Clone()
segmentutil.ReCalcRowCount(info.SegmentInfo, clonedInfo.SegmentInfo)
infos = append(infos, clonedInfo.SegmentInfo)
}
vchannel := info.InsertChannel
if _, ok := channelCPs[vchannel]; vchannel != "" && !ok {
channelCPs[vchannel] = s.meta.GetChannelCheckpoint(vchannel)
}
}
resp.Infos = infos
resp.ChannelCheckpoint = channelCPs
return resp, nil
}
// SaveBinlogPaths updates segment related binlog path
// works for Checkpoints and Flush
func (s *Server) SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPathsRequest) (*commonpb.Status, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
var (
nodeID = req.GetBase().GetSourceID()
channelName = req.GetChannel()
)
log := log.Ctx(ctx).With(
zap.Int64("nodeID", nodeID),
zap.String("channel", channelName),
zap.Int64("collectionID", req.GetCollectionID()),
zap.Int64("segmentID", req.GetSegmentID()),
zap.String("level", req.GetSegLevel().String()),
zap.Bool("withFullBinlogs", req.GetWithFullBinlogs()),
)
log.Info("receive SaveBinlogPaths request",
zap.Bool("isFlush", req.GetFlushed()),
zap.Bool("isDropped", req.GetDropped()),
zap.Any("checkpoints", req.GetCheckPoints()))
// for compatibility issue , if len(channelName) not exist, skip the check
// Also avoid to handle segment not found error if not the owner of shard
if len(channelName) != 0 {
// TODO: Current checker implementation is node id based, it cannot strictly promise when ABA channel assignment happens.
// Meanwhile the match operation is not protected with the global segment meta with the same lock,
// So the new recovery can happen with the old match operation concurrently, so it not safe enough to avoid double flush.
// Moreover, the Match operation may be called if the flusher is ready to work, but the channel manager on coord don't see the assignment success.
// So the match operation may be rejected and wait for retry.
// TODO: We need to make an idempotent operation to avoid the double flush strictly.
targetID, err := snmanager.StaticStreamingNodeManager.GetLatestWALLocated(ctx, channelName)
if err != nil || targetID != nodeID {
err := merr.WrapErrChannelNotFound(channelName, fmt.Sprintf("for node %d", nodeID))
log.Warn("failed to get latest wal allocated", zap.Int64("nodeID", nodeID), zap.Int64("channel nodeID", targetID), zap.Error(err))
return merr.Status(err), nil
}
}
// for compatibility issue, before 2.3.4, SaveBinlogPaths has only logpath
// try to parse path and fill logid
err := binlog.CompressSaveBinlogPaths(req)
if err != nil {
log.Warn("fail to CompressSaveBinlogPaths", zap.String("channel", channelName), zap.Error(err))
return merr.Status(err), nil
}
operators := []UpdateOperator{}
if req.GetSegLevel() == datapb.SegmentLevel_L0 {
operators = append(operators, CreateL0Operator(req.GetCollectionID(), req.GetPartitionID(), req.GetSegmentID(), req.GetChannel()))
} else {
segment := s.meta.GetSegment(ctx, req.GetSegmentID())
// validate level one segment
if segment == nil {
err := merr.WrapErrSegmentNotFound(req.GetSegmentID())
log.Warn("failed to get segment", zap.Error(err))
return merr.Status(err), nil
}
if segment.State == commonpb.SegmentState_Dropped {
log.Info("save to dropped segment, ignore this request")
return merr.Success(), nil
}
if !isSegmentHealthy(segment) {
err := merr.WrapErrSegmentNotFound(req.GetSegmentID())
log.Warn("failed to get segment, the segment not healthy", zap.Error(err))
return merr.Status(err), nil
}
if err := s.validateTextSegmentStorage(req); err != nil {
log.Warn("invalid TEXT segment storage format", zap.Error(err))
return merr.Status(err), nil
}
// Set storage version
operators = append(operators, SetStorageVersion(req.GetSegmentID(), req.GetStorageVersion()))
// Set segment state
if req.GetDropped() {
// segmentManager manages growing segments
s.segmentManager.DropSegment(ctx, req.GetChannel(), req.GetSegmentID())
operators = append(operators, UpdateStatusOperator(req.GetSegmentID(), commonpb.SegmentState_Dropped))
} else if req.GetFlushed() {
s.segmentManager.DropSegment(ctx, req.GetChannel(), req.GetSegmentID())
if enableSortCompaction() && req.GetSegLevel() != datapb.SegmentLevel_L0 {
operators = append(operators, SetSegmentIsInvisible(req.GetSegmentID(), true))
}
// set segment to SegmentState_Flushed
operators = append(operators, UpdateStatusOperator(req.GetSegmentID(), commonpb.SegmentState_Flushed))
}
}
if req.GetWithFullBinlogs() {
// check checkpoint will be executed at updateSegmentPack validation to ignore the illegal checkpoint update.
operators = append(operators, UpdateBinlogsFromSaveBinlogPathsOperator(
req.GetSegmentID(),
req.GetField2BinlogPaths(),
req.GetField2StatslogPaths(),
req.GetDeltalogs(),
req.GetField2Bm25LogPaths(),
), UpdateCheckPointOperator(req.GetSegmentID(), req.GetCheckPoints(), true))
} else {
operators = append(operators, AddBinlogsOperator(req.GetSegmentID(), req.GetField2BinlogPaths(), req.GetField2StatslogPaths(), req.GetDeltalogs(), req.GetField2Bm25LogPaths()),
UpdateCheckPointOperator(req.GetSegmentID(), req.GetCheckPoints()))
}
// save manifest, start positions and checkpoints
operators = append(operators,
UpdateManifest(req.GetSegmentID(), req.GetManifestPath()),
UpdateStartPosition(req.GetStartPositions()),
UpdateAsDroppedIfEmptyWhenFlushing(req.GetSegmentID()),
)
// Update segment info in memory and meta. Stale updates (segment already
// flushed / outdated time tick) are swallowed inside UpdateSegmentsInfo as
// benign no-ops, so any error here is a real failure.
if err := s.meta.UpdateSegmentsInfo(ctx, operators...); err != nil {
log.Error("save binlog and checkpoints failed", zap.Error(err))
return merr.Status(err), nil
}
s.meta.SetLastWrittenTime(req.GetSegmentID())
log.Info("SaveBinlogPaths sync segment with meta",
zap.Any("checkpoints", req.GetCheckPoints()),
zap.Strings("binlogs", stringifyBinlogs(req.GetField2BinlogPaths())),
zap.Strings("deltalogs", stringifyBinlogs(req.GetDeltalogs())),
zap.Strings("statslogs", stringifyBinlogs(req.GetField2StatslogPaths())),
zap.Strings("bm25logs", stringifyBinlogs(req.GetField2Bm25LogPaths())),
)
// Validate manifest segment after update
if segment := s.meta.GetSegment(ctx, req.GetSegmentID()); segment != nil {
if msg := ValidateManifestSegment(segment); msg != "" {
log.Warn("manifest segment validation warning", zap.String("detail", msg))
}
}
if req.GetSegLevel() == datapb.SegmentLevel_L0 {
metrics.DataCoordSizeStoredL0Segment.WithLabelValues(fmt.Sprint(req.GetCollectionID())).Observe(calculateL0SegmentSize(req.GetField2StatslogPaths()))
s.compactionTriggerManager.OnCollectionUpdate(req.GetCollectionID())
return merr.Success(), nil
}
// notify building index and compaction for "flushing/flushed" level one segment
if req.GetFlushed() {
// notify building index
s.flushCh <- req.SegmentID
// notify compaction
_, err := s.compactionTrigger.TriggerCompaction(ctx,
NewCompactionSignal().
WithWaitResult(false).
WithCollectionID(req.GetCollectionID()).
WithPartitionID(req.GetPartitionID()).
WithChannel(req.GetChannel()))
if err != nil {
log.Warn("failed to trigger single compaction")
}
}
return merr.Success(), nil
}
func (s *Server) validateTextSegmentStorage(req *datapb.SaveBinlogPathsRequest) error {
if req.GetSegLevel() == datapb.SegmentLevel_L0 || req.GetDropped() {
return nil
}
if !s.meta.collectionHasTextFields(req.GetCollectionID()) {
return nil
}
if req.GetStorageVersion() < storage.StorageV3 {
return merr.WrapErrParameterInvalidMsg(
"TEXT segment %d must be saved with StorageV3 manifest, got storage version %d",
req.GetSegmentID(),
req.GetStorageVersion())
}
if req.GetManifestPath() == "" {
return merr.WrapErrParameterInvalidMsg(
"TEXT segment %d requires non-empty StorageV3 manifest path",
req.GetSegmentID())
}
return nil
}
// DropVirtualChannel notifies vchannel dropped
// And contains the remaining data log & checkpoint to update
func (s *Server) DropVirtualChannel(ctx context.Context, req *datapb.DropVirtualChannelRequest) (*datapb.DropVirtualChannelResponse, error) {
log := log.Ctx(ctx)
resp := &datapb.DropVirtualChannelResponse{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.DropVirtualChannelResponse{
Status: merr.Status(err),
}, nil
}
channel := req.GetChannelName()
log.Info("receive DropVirtualChannel request",
zap.String("channelName", channel))
segments := make([]*SegmentInfo, 0, len(req.GetSegments()))
for _, seg2Drop := range req.GetSegments() {
info := &datapb.SegmentInfo{
ID: seg2Drop.GetSegmentID(),
CollectionID: seg2Drop.GetCollectionID(),
InsertChannel: channel,
Binlogs: seg2Drop.GetField2BinlogPaths(),
Statslogs: seg2Drop.GetField2StatslogPaths(),
Deltalogs: seg2Drop.GetDeltalogs(),
StartPosition: seg2Drop.GetStartPosition(),
DmlPosition: seg2Drop.GetCheckPoint(),
NumOfRows: seg2Drop.GetNumOfRows(),
}
segment := NewSegmentInfo(info)
segments = append(segments, segment)
}
err := s.meta.UpdateDropChannelSegmentInfo(ctx, channel, segments)
if err != nil {
log.Error("Update Drop Channel segment info failed", zap.String("channel", channel), zap.Error(err))
resp.Status = merr.Status(err)
return resp, nil
}
s.segmentManager.DropSegmentsOfChannel(ctx, channel)
s.compactionInspector.removeTasksByChannel(channel)
metrics.DataCoordCheckpointUnixSeconds.DeleteLabelValues(paramtable.GetStringNodeID(), channel)
s.meta.MarkChannelCheckpointDropped(ctx, channel)
// no compaction triggered in Drop procedure
return resp, nil
}
// SetSegmentState reset the state of the given segment.
func (s *Server) SetSegmentState(ctx context.Context, req *datapb.SetSegmentStateRequest) (*datapb.SetSegmentStateResponse, error) {
log := log.Ctx(ctx)
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.SetSegmentStateResponse{
Status: merr.Status(err),
}, nil
}
err := s.meta.SetState(ctx, req.GetSegmentId(), req.GetNewState())
if err != nil {
log.Error("failed to updated segment state in dataCoord meta",
zap.Int64("segmentID", req.SegmentId),
zap.String("newState", req.GetNewState().String()))
return &datapb.SetSegmentStateResponse{
Status: merr.Status(err),
}, nil
}
return &datapb.SetSegmentStateResponse{
Status: merr.Success(),
}, nil
}
func (s *Server) GetStateCode() commonpb.StateCode {
code := s.stateCode.Load()
if code == nil {
return commonpb.StateCode_Abnormal
}
return code.(commonpb.StateCode)
}
// UpdateStateCode update state code
func (s *Server) UpdateStateCode(code commonpb.StateCode) {
s.stateCode.Store(code)
log.Ctx(s.ctx).Info("update datacoord state", zap.String("state", code.String()))
}
// GetComponentStates returns DataCoord's current state
func (s *Server) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest) (*milvuspb.ComponentStates, error) {
code := s.GetStateCode()
log.Ctx(ctx).Debug("DataCoord current state", zap.String("StateCode", code.String()))
nodeID := common.NotRegisteredID
if s.session != nil && s.session.Registered() {
nodeID = s.session.GetServerID() // or Params.NodeID
}
resp := &milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
// NodeID: Params.NodeID, // will race with Server.Register()
NodeID: nodeID,
Role: "datacoord",
StateCode: code,
},
Status: merr.Success(),
}
return resp, nil
}
// GetRecoveryInfo get recovery info for segment.
// Called by: QueryCoord.
func (s *Server) GetRecoveryInfo(ctx context.Context, req *datapb.GetRecoveryInfoRequest) (*datapb.GetRecoveryInfoResponse, error) {
log := log.Ctx(ctx)
collectionID := req.GetCollectionID()
partitionID := req.GetPartitionID()
log = log.With(
zap.Int64("collectionID", collectionID),
zap.Int64("partitionID", partitionID),
)
resp := &datapb.GetRecoveryInfoResponse{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetRecoveryInfoResponse{
Status: merr.Status(err),
}, nil
}
dresp, err := s.broker.DescribeCollectionInternal(s.ctx, collectionID)
if err != nil {
log.Error("get collection info from rootcoord failed",
zap.Error(err))
resp.Status = merr.Status(err)
return resp, nil
}
channels := dresp.GetVirtualChannelNames()
channelInfos := make([]*datapb.VchannelInfo, 0, len(channels))
flushedIDs := make(typeutil.UniqueSet)
for _, c := range channels {
channelInfo := s.handler.GetQueryVChanPositions(&channelMeta{Name: c, CollectionID: collectionID}, partitionID)
channelInfos = append(channelInfos, channelInfo)
log.Info("datacoord append channelInfo in GetRecoveryInfo",
zap.String("channel", channelInfo.GetChannelName()),
zap.Int("# of unflushed segments", len(channelInfo.GetUnflushedSegmentIds())),
zap.Int("# of flushed segments", len(channelInfo.GetFlushedSegmentIds())),
zap.Int("# of dropped segments", len(channelInfo.GetDroppedSegmentIds())),
zap.Int("# of indexed segments", len(channelInfo.GetIndexedSegmentIds())),
zap.Int("# of l0 segments", len(channelInfo.GetLevelZeroSegmentIds())),
)
flushedIDs.Insert(channelInfo.GetFlushedSegmentIds()...)
}
segment2Binlogs := make(map[UniqueID][]*datapb.FieldBinlog)
segment2StatsBinlogs := make(map[UniqueID][]*datapb.FieldBinlog)
segment2DeltaBinlogs := make(map[UniqueID][]*datapb.FieldBinlog)
segment2InsertChannel := make(map[UniqueID]string)
segmentsNumOfRows := make(map[UniqueID]int64)
segment2TextStatsLogs := make(map[UniqueID]map[UniqueID]*datapb.TextIndexStats)
for id := range flushedIDs {
segment := s.meta.GetSegment(ctx, id)
if segment == nil {
err := merr.WrapErrSegmentNotFound(id)
log.Warn("failed to get segment", zap.Int64("segmentID", id))
resp.Status = merr.Status(err)
return resp, nil
}
// Skip non-flushing, non-flushed and dropped segments.
if segment.State != commonpb.SegmentState_Flushed && segment.State != commonpb.SegmentState_Flushing && segment.State != commonpb.SegmentState_Dropped {
continue
}
// Also skip bulk insert & fake segments.
if segment.GetIsImporting() || segment.GetIsFake() {
continue
}
segment2InsertChannel[segment.ID] = segment.InsertChannel
binlogs := segment.GetBinlogs()
if len(binlogs) == 0 && segment.GetManifestPath() == "" {
flushedIDs.Remove(id)
continue
}
field2Binlog := make(map[UniqueID]*datapb.FieldBinlog)
for _, field := range binlogs {
fieldBinlog, ok := field2Binlog[field.GetFieldID()]
if !ok {
fieldBinlog = &datapb.FieldBinlog{
FieldID: field.GetFieldID(),
ChildFields: field.GetChildFields(),
Format: field.GetFormat(),
}
field2Binlog[field.GetFieldID()] = fieldBinlog
} else {
if len(fieldBinlog.ChildFields) == 0 {
fieldBinlog.ChildFields = field.GetChildFields()
}
if fieldBinlog.Format == "" {
fieldBinlog.Format = field.GetFormat()
}
}
fieldBinlog.Binlogs = append(fieldBinlog.Binlogs, field.GetBinlogs()...)
}
for _, fieldBinlogs := range field2Binlog {
segment2Binlogs[id] = append(segment2Binlogs[id], fieldBinlogs)
}
if newCount := segmentutil.CalcRowCountFromBinLog(segment.SegmentInfo); newCount != segment.NumOfRows && newCount > 0 {
log.Warn("segment row number meta inconsistent with bin log row count and will be corrected",
zap.Int64("segmentID", segment.GetID()),
zap.Int64("segment meta row count (wrong)", segment.GetNumOfRows()),
zap.Int64("segment bin log row count (correct)", newCount))
segmentsNumOfRows[id] = newCount
} else {
segmentsNumOfRows[id] = segment.NumOfRows
}
statsBinlogs := segment.GetStatslogs()
field2StatsBinlog := make(map[UniqueID][]*datapb.Binlog)
for _, field := range statsBinlogs {
field2StatsBinlog[field.GetFieldID()] = append(field2StatsBinlog[field.GetFieldID()], field.GetBinlogs()...)
}
for f, paths := range field2StatsBinlog {
fieldBinlogs := &datapb.FieldBinlog{
FieldID: f,
Binlogs: paths,
}
segment2StatsBinlogs[id] = append(segment2StatsBinlogs[id], fieldBinlogs)
}
segment2TextStatsLogs[id] = segment.GetTextStatsLogs()
if len(segment.GetDeltalogs()) > 0 {
segment2DeltaBinlogs[id] = append(segment2DeltaBinlogs[id], segment.GetDeltalogs()...)
}
}
binlogs := make([]*datapb.SegmentBinlogs, 0, len(segment2Binlogs))
for segmentID := range flushedIDs {
sbl := &datapb.SegmentBinlogs{
SegmentID: segmentID,
NumOfRows: segmentsNumOfRows[segmentID],
FieldBinlogs: segment2Binlogs[segmentID],
Statslogs: segment2StatsBinlogs[segmentID],
Deltalogs: segment2DeltaBinlogs[segmentID],
InsertChannel: segment2InsertChannel[segmentID],
TextStatsLogs: segment2TextStatsLogs[segmentID],
}
binlogs = append(binlogs, sbl)
}
resp.Channels = channelInfos
resp.Binlogs = binlogs
return resp, nil
}
// GetRecoveryInfoV2 get recovery info for segment
// Called by: QueryCoord.
func (s *Server) GetRecoveryInfoV2(ctx context.Context, req *datapb.GetRecoveryInfoRequestV2) (*datapb.GetRecoveryInfoResponseV2, error) {
log := log.Ctx(ctx)
collectionID := req.GetCollectionID()
partitionIDs := req.GetPartitionIDs()
log = log.With(
zap.Int64("collectionID", collectionID),
zap.Int64s("partitionIDs", partitionIDs),
)
resp := &datapb.GetRecoveryInfoResponseV2{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetRecoveryInfoResponseV2{
Status: merr.Status(err),
}, nil
}
channels, err := s.getChannelsByCollectionID(ctx, collectionID)
if err != nil {
return &datapb.GetRecoveryInfoResponseV2{
Status: merr.Status(err),
}, nil
}
channelInfos := make([]*datapb.VchannelInfo, 0, len(channels))
flushedIDs := make(typeutil.UniqueSet)
for _, ch := range channels {
channelInfo := s.handler.GetQueryVChanPositions(ch, partitionIDs...)
channelInfos = append(channelInfos, channelInfo)
log.Info("datacoord append channelInfo in GetRecoveryInfo",
zap.String("channel", channelInfo.GetChannelName()),
zap.Any("seekPos", channelInfo.GetSeekPosition()),
zap.Int("# of unflushed segments", len(channelInfo.GetUnflushedSegmentIds())),
zap.Int("# of flushed segments", len(channelInfo.GetFlushedSegmentIds())),
zap.Int("# of dropped segments", len(channelInfo.GetDroppedSegmentIds())),
zap.Int("# of indexed segments", len(channelInfo.GetIndexedSegmentIds())),
zap.Int("# of l0 segments", len(channelInfo.GetLevelZeroSegmentIds())),
zap.Time("# of check point", tsoutil.PhysicalTime(channelInfo.GetSeekPosition().GetTimestamp())),
zap.Time("# of delete check point", tsoutil.PhysicalTime(channelInfo.GetDeleteCheckpoint().GetTimestamp())),
)
flushedIDs.Insert(channelInfo.GetFlushedSegmentIds()...)
}
segmentInfos := make([]*datapb.SegmentInfo, 0)
for id := range flushedIDs {
segment := s.meta.GetSegment(ctx, id)
if segment == nil {
err := merr.WrapErrSegmentNotFound(id)
log.Warn("failed to get segment", zap.Int64("segmentID", id))
resp.Status = merr.Status(err)
return resp, nil
}
// Skip non-flushing, non-flushed and dropped segments.
if segment.State != commonpb.SegmentState_Flushed && segment.State != commonpb.SegmentState_Flushing && segment.State != commonpb.SegmentState_Dropped {
continue
}
// Also skip bulk insert segments.
if segment.GetIsImporting() {
continue
}
binlogs := segment.GetBinlogs()
if len(binlogs) == 0 && segment.GetLevel() != datapb.SegmentLevel_L0 && segment.GetManifestPath() == "" {
continue
}
rowCount := segmentutil.CalcRowCountFromBinLog(segment.SegmentInfo)
if rowCount != segment.NumOfRows && rowCount > 0 {
log.Warn("segment row number meta inconsistent with bin log row count and will be corrected",
zap.Int64("segmentID", segment.GetID()),
zap.Int64("segment meta row count (wrong)", segment.GetNumOfRows()),
zap.Int64("segment bin log row count (correct)", rowCount))
} else {
rowCount = segment.NumOfRows
}
segmentInfos = append(segmentInfos, &datapb.SegmentInfo{
ID: segment.ID,
PartitionID: segment.PartitionID,
CollectionID: segment.CollectionID,
InsertChannel: segment.InsertChannel,
NumOfRows: rowCount,
Level: segment.GetLevel(),
IsSorted: segment.GetIsSorted(),
IsSortedByNamespace: segment.GetIsSortedByNamespace(),
ManifestPath: segment.GetManifestPath(),
DataVersion: segment.GetDataVersion(),
})
}
resp.Channels = channelInfos
resp.Segments = segmentInfos
return resp, nil
}
// GetChannelRecoveryInfo get recovery channel info.
// Called by: StreamingNode.
func (s *Server) GetChannelRecoveryInfo(ctx context.Context, req *datapb.GetChannelRecoveryInfoRequest) (*datapb.GetChannelRecoveryInfoResponse, error) {
log := log.Ctx(ctx).With(
zap.String("vchannel", req.GetVchannel()),
)
log.Info("get channel recovery info request received")
resp := &datapb.GetChannelRecoveryInfoResponse{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
resp.Status = merr.Status(err)
return resp, nil
}
collectionID := funcutil.GetCollectionIDFromVChannel(req.GetVchannel())
channelInfo := s.handler.GetDataVChanPositions(&channelMeta{
Name: req.GetVchannel(),
CollectionID: collectionID,
}, allPartitionID)
if channelInfo.SeekPosition == nil {
log.Warn("channel recovery start position is not found, may collection is on creating")
resp.Status = merr.Status(merr.WrapErrChannelNotAvailable(req.GetVchannel(), "start position is nil"))
return resp, nil
}
segmentsNotCreatedByStreaming := make([]*datapb.SegmentNotCreatedByStreaming, 0)
for _, segmentID := range channelInfo.GetUnflushedSegmentIds() {
segment := s.meta.GetSegment(ctx, segmentID)
if segment != nil && !segment.IsCreatedByStreaming {
segmentsNotCreatedByStreaming = append(segmentsNotCreatedByStreaming, &datapb.SegmentNotCreatedByStreaming{
CollectionId: segment.CollectionID,
PartitionId: segment.PartitionID,
SegmentId: segmentID,
})
}
}
log.Info("datacoord get channel recovery info",
zap.String("channel", channelInfo.GetChannelName()),
zap.Int("# of unflushed segments", len(channelInfo.GetUnflushedSegmentIds())),
zap.Int("# of flushed segments", len(channelInfo.GetFlushedSegmentIds())),
zap.Int("# of dropped segments", len(channelInfo.GetDroppedSegmentIds())),
zap.Int("# of indexed segments", len(channelInfo.GetIndexedSegmentIds())),
zap.Int("# of l0 segments", len(channelInfo.GetLevelZeroSegmentIds())),
zap.Int("# of segments not created by streaming", len(segmentsNotCreatedByStreaming)),
)
resp.Info = channelInfo
resp.Schema = nil // schema is managed by streaming node itself now.
resp.SegmentsNotCreatedByStreaming = segmentsNotCreatedByStreaming
return resp, nil
}
// GetFlushedSegments returns all segment matches provided criterion and in state Flushed or Dropped (compacted but not GCed yet)
// If requested partition id < 0, ignores the partition id filter
func (s *Server) GetFlushedSegments(ctx context.Context, req *datapb.GetFlushedSegmentsRequest) (*datapb.GetFlushedSegmentsResponse, error) {
log := log.Ctx(ctx)
resp := &datapb.GetFlushedSegmentsResponse{
Status: merr.Success(),
}
collectionID := req.GetCollectionID()
partitionID := req.GetPartitionID()
log.Info("received get flushed segments request",
zap.Int64("collectionID", collectionID),
zap.Int64("partitionID", partitionID),
)
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetFlushedSegmentsResponse{
Status: merr.Status(err),
}, nil
}
var segmentIDs []UniqueID
if partitionID < 0 {
segmentIDs = s.meta.GetSegmentsIDOfCollectionWithDropped(ctx, collectionID)
} else {
segmentIDs = s.meta.GetSegmentsIDOfPartitionWithDropped(ctx, collectionID, partitionID)
}
ret := make([]UniqueID, 0, len(segmentIDs))
for _, id := range segmentIDs {
segment := s.meta.GetSegment(ctx, id)
// if this segment == nil, we assume this segment has been gc
if segment == nil ||
(segment.GetState() != commonpb.SegmentState_Dropped &&
segment.GetState() != commonpb.SegmentState_Flushed &&
segment.GetState() != commonpb.SegmentState_Flushing) {
continue
}
if !req.GetIncludeUnhealthy() && segment.GetState() == commonpb.SegmentState_Dropped {
continue
}
ret = append(ret, id)
}
resp.Segments = ret
return resp, nil
}
// GetSegmentsByStates returns all segment matches provided criterion and States
// If requested partition id < 0, ignores the partition id filter
func (s *Server) GetSegmentsByStates(ctx context.Context, req *datapb.GetSegmentsByStatesRequest) (*datapb.GetSegmentsByStatesResponse, error) {
log := log.Ctx(ctx)
resp := &datapb.GetSegmentsByStatesResponse{
Status: merr.Success(),
}
collectionID := req.GetCollectionID()
partitionID := req.GetPartitionID()
states := req.GetStates()
log.Info("received get segments by states request",
zap.Int64("collectionID", collectionID),
zap.Int64("partitionID", partitionID),
zap.Any("states", states))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetSegmentsByStatesResponse{
Status: merr.Status(err),
}, nil
}
var segmentIDs []UniqueID
channels, err := s.getChannelsByCollectionID(ctx, collectionID)
if err != nil {
return &datapb.GetSegmentsByStatesResponse{
Status: merr.Status(err),
}, nil
}
for _, channel := range channels {
channelSegmentsView := s.handler.GetCurrentSegmentsView(ctx, channel, partitionID)
if channelSegmentsView == nil {
continue
}
segmentIDs = append(segmentIDs, channelSegmentsView.FlushedSegmentIDs...)
segmentIDs = append(segmentIDs, channelSegmentsView.GrowingSegmentIDs...)
segmentIDs = append(segmentIDs, channelSegmentsView.L0SegmentIDs...)
segmentIDs = append(segmentIDs, channelSegmentsView.ImportingSegmentIDs...)
}
ret := make([]UniqueID, 0, len(segmentIDs))
statesDict := make(map[commonpb.SegmentState]bool)
for _, state := range states {
statesDict[state] = true
}
for _, id := range segmentIDs {
segment := s.meta.GetHealthySegment(ctx, id)
if segment != nil && statesDict[segment.GetState()] {
ret = append(ret, id)
}
}
resp.Segments = ret
return resp, nil
}
// ShowConfigurations returns the configurations of DataCoord matching req.Pattern
func (s *Server) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &internalpb.ShowConfigurationsResponse{
Status: merr.Status(err),
}, nil
}
configList := make([]*commonpb.KeyValuePair, 0)
for key, value := range Params.GetComponentConfigurations("datacoord", req.Pattern) {
configList = append(configList,
&commonpb.KeyValuePair{
Key: key,
Value: value,
})
}
return &internalpb.ShowConfigurationsResponse{
Status: merr.Success(),
Configuations: configList,
}, nil
}
// GetMetrics returns DataCoord metrics info
// it may include SystemMetrics, Topology metrics, etc.
func (s *Server) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
log := log.Ctx(ctx)
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
msg := "failed to get metrics"
log.Warn(msg, zap.Error(err))
return &milvuspb.GetMetricsResponse{
Status: merr.Status(merr.Wrap(err, msg)),
}, nil
}
resp := &milvuspb.GetMetricsResponse{
Status: merr.Success(),
ComponentName: metricsinfo.ConstructComponentName(typeutil.DataCoordRole,
paramtable.GetNodeID()),
}
ret, err := s.metricsRequest.ExecuteMetricsRequest(ctx, req)
if err != nil {
resp.Status = merr.Status(err)
return resp, nil
}
resp.Response = ret
return resp, nil
}
// ManualCompaction triggers a compaction for a collection
func (s *Server) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("collectionID", req.GetCollectionID()),
)
log.Info("received manual compaction", zap.Any("request", req))
resp := &milvuspb.ManualCompactionResponse{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &milvuspb.ManualCompactionResponse{
Status: merr.Status(err),
}, nil
}
if !Params.DataCoordCfg.EnableCompaction.GetAsBool() {
resp.Status = merr.Status(merr.WrapErrServiceUnavailable("compaction disabled"))
return resp, nil
}
var id int64
var err error
if req.GetMajorCompaction() || req.GetL0Compaction() || req.GetTargetSize() != 0 {
id, err = s.compactionTriggerManager.ManualTrigger(ctx, req.CollectionID, req.GetMajorCompaction(), req.GetL0Compaction(), req.GetTargetSize())
} else {
id, err = s.compactionTrigger.TriggerCompaction(ctx, NewCompactionSignal().
WithIsForce(true).
WithCollectionID(req.GetCollectionID()).
WithPartitionID(req.GetPartitionId()).
WithChannel(req.GetChannel()).
WithSegmentIDs(req.GetSegmentIds()...),
)
}
if err != nil {
log.Error("failed to trigger manual compaction", zap.Error(err))
resp.Status = merr.Status(err)
return resp, nil
}
taskCnt := s.compactionInspector.getCompactionTasksNumBySignalID(id)
if taskCnt == 0 {
resp.CompactionID = -1
resp.CompactionPlanCount = 0
} else {
resp.CompactionID = id
resp.CompactionPlanCount = int32(taskCnt)
}
log.Info("success to trigger manual compaction", zap.Bool("isL0Compaction", req.GetL0Compaction()),
zap.Bool("isMajorCompaction", req.GetMajorCompaction()), zap.Int64("targetSize", req.GetTargetSize()), zap.Int64("compactionID", id), zap.Int("taskNum", taskCnt))
return resp, nil
}
// GetCompactionState gets the state of a compaction
func (s *Server) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("compactionID", req.GetCompactionID()),
)
log.Info("received get compaction state request")
resp := &milvuspb.GetCompactionStateResponse{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &milvuspb.GetCompactionStateResponse{
Status: merr.Status(err),
}, nil
}
if !Params.DataCoordCfg.EnableCompaction.GetAsBool() {
resp.Status = merr.Status(merr.WrapErrServiceUnavailable("compaction disabled"))
return resp, nil
}
info := s.compactionInspector.getCompactionInfo(ctx, req.GetCompactionID())
resp.State = info.state
resp.ExecutingPlanNo = int64(info.executingCnt)
resp.CompletedPlanNo = int64(info.completedCnt)
resp.TimeoutPlanNo = int64(info.timeoutCnt)
resp.FailedPlanNo = int64(info.failedCnt)
log.Info("success to get compaction state", zap.Any("state", info.state), zap.Int("executing", info.executingCnt),
zap.Int("completed", info.completedCnt), zap.Int("failed", info.failedCnt), zap.Int("timeout", info.timeoutCnt))
return resp, nil
}
// GetCompactionStateWithPlans returns the compaction state of given plan
func (s *Server) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("compactionID", req.GetCompactionID()),
)
log.Info("received the request to get compaction state with plans")
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &milvuspb.GetCompactionPlansResponse{
Status: merr.Status(err),
}, nil
}
resp := &milvuspb.GetCompactionPlansResponse{
Status: merr.Success(),
}
if !Params.DataCoordCfg.EnableCompaction.GetAsBool() {
resp.Status = merr.Status(merr.WrapErrServiceUnavailable("compaction disabled"))
return resp, nil
}
info := s.compactionInspector.getCompactionInfo(ctx, req.GetCompactionID())
resp.State = info.state
resp.MergeInfos = lo.MapToSlice[int64, *milvuspb.CompactionMergeInfo](info.mergeInfos, func(_ int64, merge *milvuspb.CompactionMergeInfo) *milvuspb.CompactionMergeInfo {
return merge
})
planIDs := lo.MapToSlice[int64, *milvuspb.CompactionMergeInfo](info.mergeInfos, func(planID int64, _ *milvuspb.CompactionMergeInfo) int64 { return planID })
log.Info("success to get state with plans", zap.Any("state", info.state), zap.Any("merge infos", resp.MergeInfos),
zap.Int64s("plans", planIDs))
return resp, nil
}
// WatchChannels notifies DataCoord to watch vchannels of a collection.
// Deprecated: Redundant design by now, remove it in future.
func (s *Server) WatchChannels(ctx context.Context, req *datapb.WatchChannelsRequest) (*datapb.WatchChannelsResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("collectionID", req.GetCollectionID()),
zap.Strings("channels", req.GetChannelNames()),
zap.Any("dbProperties", req.GetDbProperties()),
)
log.Info("receive watch channels request")
resp := &datapb.WatchChannelsResponse{
Status: merr.Success(),
}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.WatchChannelsResponse{
Status: merr.Status(err),
}, nil
}
for _, channelName := range req.GetChannelNames() {
// TODO: redundant channel mark by now, remove it in future.
if err := s.meta.catalog.MarkChannelAdded(ctx, channelName); err != nil {
// TODO: add background task to periodically cleanup the orphaned channel add marks.
log.Error("failed to mark channel added", zap.Error(err))
resp.Status = merr.Status(err)
return resp, nil
}
// try to init channel checkpoint, if failed, we will log it and continue
startPos := toMsgPositionWithWALNames(channelName, req.GetStartPositions(), req.ChannelWalNames)
if startPos != nil {
startPos.Timestamp = req.GetCreateTimestamp()
if err := s.meta.UpdateChannelCheckpoint(ctx, channelName, startPos); err != nil {
log.Warn("failed to init channel checkpoint, meta update error", zap.String("channel", channelName), zap.Error(err))
}
} else {
log.Info("skip to init channel checkpoint for nil startPosition", zap.String("channel", channelName))
}
}
return resp, nil
}
// GetFlushState gets the flush state of the collection based on the provided flush ts and segment IDs.
func (s *Server) GetFlushState(ctx context.Context, req *datapb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) {
log := log.Ctx(ctx).With(zap.Int64("collection", req.GetCollectionID()),
zap.Uint64("flushTs", req.GetFlushTs()),
zap.Time("flushTs in time", tsoutil.PhysicalTime(req.GetFlushTs()))).
WithRateGroup("dc.GetFlushState", 1, 60)
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &milvuspb.GetFlushStateResponse{
Status: merr.Status(err),
}, nil
}
resp := &milvuspb.GetFlushStateResponse{Status: merr.Success()}
if len(req.GetSegmentIDs()) > 0 {
var unflushed []UniqueID
for _, sid := range req.GetSegmentIDs() {
segment := s.meta.GetHealthySegment(ctx, sid)
// segment is nil if it was compacted, or it's an empty segment and is set to dropped
// TODO: Here's a dirty implementation, because a growing segment may cannot be seen right away by mixcoord,
// it can only be seen by streamingnode right away, so we need to check the flush state at streamingnode but not here.
// use timetick for GetFlushState in-future but not segment list.
if segment == nil || isFlushState(segment.GetState()) {
continue
}
unflushed = append(unflushed, sid)
}
if len(unflushed) != 0 {
log.RatedInfo(10, "DataCoord receive GetFlushState request, Flushed is false", zap.Int64s("unflushed", unflushed), zap.Int("len", len(unflushed)))
resp.Flushed = false
return resp, nil
}
}
channels, err := s.getChannelsByCollectionID(ctx, req.GetCollectionID())
if err != nil {
return &milvuspb.GetFlushStateResponse{
Status: merr.Status(err),
}, nil
}
if len(channels) == 0 { // For compatibility with old client
resp.Flushed = true
log.Info("GetFlushState all flushed without checking flush ts")
return resp, nil
}
for _, channel := range channels {
cp := s.meta.GetChannelCheckpoint(channel.GetName())
cpTs := uint64(0)
if cp != nil {
cpTs = cp.GetTimestamp()
}
if cp == nil || cpTs < req.GetFlushTs() {
resp.Flushed = false
log.RatedInfo(10, "GetFlushState failed, channel unflushed", zap.String("channel", channel.GetName()),
zap.Time("CP", tsoutil.PhysicalTime(cpTs)),
zap.Duration("lag", tsoutil.PhysicalTime(req.GetFlushTs()).Sub(tsoutil.PhysicalTime(cpTs))))
return resp, nil
}
}
resp.Flushed = true
log.Info("GetFlushState all flushed")
return resp, nil
}
// getChannelsByCollectionID gets the channels of the collection.
func (s *Server) getChannelsByCollectionID(ctx context.Context, collectionID int64) ([]RWChannel, error) {
describeRsp, err := s.mixCoord.DescribeCollectionInternal(ctx, &milvuspb.DescribeCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DescribeCollection,
},
CollectionID: collectionID,
})
if err != nil {
return nil, err
}
channels := make([]RWChannel, 0, len(describeRsp.GetVirtualChannelNames()))
for _, channel := range describeRsp.GetVirtualChannelNames() {
startPos := toMsgPosition(channel, describeRsp.GetStartPositions())
channels = append(channels, &channelMeta{
Name: channel,
CollectionID: collectionID,
StartPosition: startPos,
})
}
return channels, nil
}
// GetFlushAllState checks if all DML messages before `FlushAllTs` have been flushed.
func (s *Server) GetFlushAllState(ctx context.Context, req *milvuspb.GetFlushAllStateRequest) (*milvuspb.GetFlushAllStateResponse, error) {
log := log.Ctx(ctx).WithRateGroup("dc.GetFlushAllState", 1, 60)
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &milvuspb.GetFlushAllStateResponse{
Status: merr.Status(err),
}, nil
}
resp := &milvuspb.GetFlushAllStateResponse{
Status: merr.Success(),
}
// TODO: Introduce pchannel level flush checkpoint to
// check if the flush is complete.
// Rather than validate every vchannel checkpoint.
dbsRsp, err := s.broker.ListDatabases(ctx)
if err != nil {
log.Warn("failed to ListDatabases", zap.Error(err))
resp.Status = merr.Status(err)
return resp, nil
}
targetDbs := lo.Uniq(dbsRsp.DbNames)
allFlushed := true
OUTER:
for _, dbName := range targetDbs {
showColRsp, err := s.broker.ShowCollections(ctx, dbName)
if err != nil {
log.Warn("failed to ShowCollections", zap.String("db", dbName), zap.Error(err))
resp.Status = merr.Status(err)
return resp, nil
}
for _, collectionID := range showColRsp.GetCollectionIds() {
describeColRsp, err := s.broker.DescribeCollectionInternal(ctx, collectionID)
if err != nil {
log.Warn("failed to DescribeCollectionInternal", zap.Int64("collectionID", collectionID), zap.Error(err))
resp.Status = merr.Status(err)
return resp, nil
}
for _, channel := range describeColRsp.GetVirtualChannelNames() {
if len(req.GetFlushAllTss()) > 0 {
ok, err := s.verifyFlushAllStateByChannelFlushAllTs(log, channel, req.GetFlushAllTss())
if err != nil {
resp.Status = merr.Status(err)
return resp, nil
}
if !ok {
allFlushed = false
break OUTER
}
} else if req.GetFlushAllTs() != 0 {
// For compatibility, if deprecated FlushAllTs is provided, use it to verify the flush state.
if !s.verifyFlushAllStateByLegacyFlushAllTs(log, channel, req.GetFlushAllTs()) {
allFlushed = false
break OUTER
}
} else {
resp.Status = merr.Status(merr.WrapErrParameterMissingMsg("FlushAllTss or FlushAllTs is required"))
return resp, nil
}
}
}
}
if allFlushed {
log.Info("GetFlushAllState all flushed", zap.Any("flushAllTss", req.GetFlushAllTss()), zap.Uint64("FlushAllTs", req.GetFlushAllTs()))
}
resp.Flushed = allFlushed
return resp, nil
}
func (s *Server) verifyFlushAllStateByChannelFlushAllTs(logger *log.MLogger, channel string, flushAllTss map[string]uint64) (bool, error) {
channelCP := s.meta.GetChannelCheckpoint(channel)
pchannel := funcutil.ToPhysicalChannel(channel)
flushAllTs, ok := flushAllTss[pchannel]
if !ok || flushAllTs == 0 {
logger.Warn("FlushAllTs not found for pchannel", zap.String("pchannel", pchannel), zap.Uint64("flushAllTs", flushAllTs))
return false, merr.WrapErrParameterInvalidMsg("FlushAllTs not found for pchannel %s", pchannel)
}
if channelCP == nil || channelCP.GetTimestamp() < flushAllTs {
logger.RatedInfo(10, "channel unflushed",
zap.String("vchannel", channel),
zap.Uint64("flushAllTs", flushAllTs),
zap.Uint64("channelCP", channelCP.GetTimestamp()),
)
return false, nil
}
return true, nil
}
func (s *Server) verifyFlushAllStateByLegacyFlushAllTs(logger *log.MLogger, channel string, flushAllTs uint64) bool {
channelCP := s.meta.GetChannelCheckpoint(channel)
if channelCP == nil || channelCP.GetTimestamp() < flushAllTs {
logger.RatedInfo(10, "channel unflushed",
zap.String("vchannel", channel),
zap.Uint64("flushAllTs", flushAllTs),
zap.Uint64("channelCP", channelCP.GetTimestamp()),
)
return false
}
return true
}
// Deprecated
// UpdateSegmentStatistics updates a segment's stats.
func (s *Server) UpdateSegmentStatistics(ctx context.Context, req *datapb.UpdateSegmentStatisticsRequest) (*commonpb.Status, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
return merr.Success(), nil
}
// UpdateChannelCheckpoint updates channel checkpoint in dataCoord.
func (s *Server) UpdateChannelCheckpoint(ctx context.Context, req *datapb.UpdateChannelCheckpointRequest) (*commonpb.Status, error) {
log := log.Ctx(ctx)
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
nodeID := req.GetBase().GetSourceID()
// For compatibility with old client
if req.GetVChannel() != "" && req.GetPosition() != nil {
channel := req.GetVChannel()
targetID, err := snmanager.StaticStreamingNodeManager.GetLatestWALLocated(ctx, channel)
if err != nil || targetID != nodeID {
err := merr.WrapErrChannelNotFound(channel, fmt.Sprintf("for node %d", nodeID))
log.Warn("failed to get latest wal allocated", zap.Error(err))
return merr.Status(err), nil
}
if err := s.meta.UpdateChannelCheckpoint(ctx, req.GetVChannel(), req.GetPosition()); err != nil {
log.Warn("failed to UpdateChannelCheckpoint", zap.String("vChannel", req.GetVChannel()), zap.Error(err))
return merr.Status(err), nil
}
return merr.Success(), nil
}
checkpoints := lo.Filter(req.GetChannelCheckpoints(), func(cp *msgpb.MsgPosition, _ int) bool {
channel := cp.GetChannelName()
targetID, err := snmanager.StaticStreamingNodeManager.GetLatestWALLocated(ctx, channel)
if err != nil || targetID != nodeID {
err := merr.WrapErrChannelNotFound(channel, fmt.Sprintf("for node %d", nodeID))
log.Warn("failed to get latest wal allocated", zap.Error(err))
return false
}
return true
})
err := s.meta.UpdateChannelCheckpoints(ctx, checkpoints)
if err != nil {
log.Warn("failed to update channel checkpoint", zap.Error(err))
return merr.Status(err), nil
}
for _, pos := range checkpoints {
if pos == nil || pos.GetMsgID() == nil || pos.GetChannelName() == "" {
continue
}
s.segmentManager.CleanZeroSealedSegmentsOfChannel(ctx, pos.GetChannelName(), pos.GetTimestamp())
}
return merr.Success(), nil
}
// ReportDataNodeTtMsgs gets timetick messages from datanode.
func (s *Server) ReportDataNodeTtMsgs(ctx context.Context, req *datapb.ReportDataNodeTtMsgsRequest) (*commonpb.Status, error) {
return merr.Success(), nil
}
// MarkSegmentsDropped marks the given segments as `Dropped`.
// An error status will be returned and error will be logged, if we failed to mark *all* segments.
// Deprecated, do not use it
func (s *Server) MarkSegmentsDropped(ctx context.Context, req *datapb.MarkSegmentsDroppedRequest) (*commonpb.Status, error) {
log.Ctx(ctx).Info("marking segments dropped", zap.Int64s("segments", req.GetSegmentIds()))
var err error
for _, segID := range req.GetSegmentIds() {
if err = s.meta.SetState(ctx, segID, commonpb.SegmentState_Dropped); err != nil {
// Fail-open.
log.Ctx(ctx).Error("failed to set segment state as dropped", zap.Int64("segmentID", segID))
break
}
}
return merr.Status(err), nil
}
func (s *Server) BroadcastAlteredCollection(ctx context.Context, req *datapb.AlterCollectionRequest) (*commonpb.Status, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
// get collection info from cache
clonedColl := s.meta.GetClonedCollectionInfo(req.CollectionID)
properties := make(map[string]string)
for _, pair := range req.Properties {
properties[pair.GetKey()] = pair.GetValue()
}
// cache miss and update cache
if clonedColl == nil {
collInfo := &collectionInfo{
ID: req.GetCollectionID(),
Schema: req.GetSchema(),
Partitions: req.GetPartitionIDs(),
StartPositions: req.GetStartPositions(),
Properties: properties,
DatabaseID: req.GetDbID(),
DatabaseName: req.GetSchema().GetDbName(),
VChannelNames: req.GetVChannels(),
}
s.meta.AddCollection(collInfo)
return merr.Success(), nil
}
clonedColl.Properties = properties
// add field will change the schema
clonedColl.Schema = req.GetSchema()
s.meta.AddCollection(clonedColl)
return merr.Success(), nil
}
func (s *Server) CheckHealth(ctx context.Context, req *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &milvuspb.CheckHealthResponse{
Status: merr.Status(err),
Reasons: []string{err.Error()},
}, nil
}
if err := CheckCheckPointsHealth(s.meta); err != nil {
return componentutil.CheckHealthRespWithErr(err), nil
}
return componentutil.CheckHealthRespWithErr(nil), nil
}
func (s *Server) GcConfirm(ctx context.Context, request *datapb.GcConfirmRequest) (*datapb.GcConfirmResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GcConfirmResponse{
Status: merr.Status(err),
}, nil
}
resp := &datapb.GcConfirmResponse{
Status: merr.Success(),
}
resp.GcFinished = s.meta.GcConfirm(ctx, request.GetCollectionId(), request.GetPartitionId())
return resp, nil
}
func (s *Server) GcControl(ctx context.Context, request *datapb.GcControlRequest) (*commonpb.Status, error) {
status := &commonpb.Status{}
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
switch request.GetCommand() {
case datapb.GcCommand_Pause:
kv := lo.FindOrElse(request.GetParams(), nil, func(kv *commonpb.KeyValuePair) bool {
return kv.GetKey() == "duration"
})
if kv == nil {
status.ErrorCode = commonpb.ErrorCode_UnexpectedError
status.Reason = "pause duration param not found"
return status, nil
}
pauseSeconds, err := strconv.ParseInt(kv.GetValue(), 10, 64)
if err != nil {
status.ErrorCode = commonpb.ErrorCode_UnexpectedError
status.Reason = fmt.Sprintf("pause duration not valid, %s", err.Error())
return status, nil
}
collectionID, err, _ := common.GetInt64Value(request.GetParams(), "collection_id")
if err != nil {
return merr.Status(err), nil
}
ticket, _ := common.GetStringValue(request.GetParams(), "ticket")
if err := s.garbageCollector.Pause(ctx, collectionID, ticket, time.Duration(pauseSeconds)*time.Second); err != nil {
status.ErrorCode = commonpb.ErrorCode_UnexpectedError
status.Reason = fmt.Sprintf("failed to pause gc, %s", err.Error())
return status, nil
}
case datapb.GcCommand_Resume:
collectionID, err, _ := common.GetInt64Value(request.GetParams(), "collection_id")
if err != nil {
return merr.Status(err), nil
}
ticket, _ := common.GetStringValue(request.GetParams(), "ticket")
if err := s.garbageCollector.Resume(ctx, collectionID, ticket); err != nil {
status.ErrorCode = commonpb.ErrorCode_UnexpectedError
status.Reason = fmt.Sprintf("failed to pause gc, %s", err.Error())
return status, nil
}
default:
status.ErrorCode = commonpb.ErrorCode_UnexpectedError
status.Reason = fmt.Sprintf("unknown gc command: %d", request.GetCommand())
return status, nil
}
return status, nil
}
func (s *Server) GetGcStatus(ctx context.Context) (*datapb.GetGcStatusResponse, error) {
status := s.garbageCollector.GetStatus()
var remainingSeconds int32
if status.IsPaused {
// Convert time.Duration to seconds, rounding to the nearest second.
// Using Round() ensures accuracy when converting to an integer.
remainingSeconds = int32(status.TimeRemaining.Round(time.Second).Seconds())
}
return &datapb.GetGcStatusResponse{
IsPaused: status.IsPaused,
TimeRemainingSeconds: remainingSeconds,
}, nil
}
// ImportV2 handles import requests from proxy by broadcasting import messages.
// This is the entry point for all user-initiated imports.
func (s *Server) ImportV2(ctx context.Context, in *internalpb.ImportRequestInternal) (*internalpb.ImportResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &internalpb.ImportResponse{
Status: merr.Status(err),
}, nil
}
resp := &internalpb.ImportResponse{
Status: merr.Success(),
}
log := log.Ctx(ctx).With(zap.Int64("collection", in.GetCollectionID()),
zap.Int64s("partitions", in.GetPartitionIDs()),
zap.Strings("channels", in.GetChannelNames()))
log.Info("receive import request from proxy, will broadcast",
zap.Int("fileNum", len(in.GetFiles())),
zap.Any("files", in.GetFiles()),
zap.Any("options", in.GetOptions()))
// Validate timeout before allocating resources
// Full validation will happen during broadcast
_, err := importutilv2.GetTimeoutTs(in.GetOptions())
if err != nil {
resp.Status = merr.Status(merr.WrapErrImportFailed(err.Error()))
return resp, nil
}
// Use the incoming JobID if provided (backward compat: old proxy allocates jobID
// before sending broadcast RPC, which is forwarded here with the original jobID).
// Otherwise allocate a new one.
jobID := in.GetJobID()
if jobID == 0 {
if s.allocator == nil {
resp.Status = merr.Status(merr.WrapErrServiceUnavailable("allocator not initialized"))
return resp, nil
}
jobID, _, err = s.allocator.AllocN(1)
if err != nil {
resp.Status = merr.Status(merr.Wrap(err, "failed to allocate job ID"))
return resp, nil
}
}
// Broadcast the import message
// dbName is retrieved inside broadcastImport via broker.DescribeCollectionInternal
err = s.broadcastImport(
ctx,
in.GetCollectionName(),
in.GetCollectionID(),
in.GetPartitionIDs(),
in.GetFiles(),
in.GetOptions(),
in.GetSchema(),
jobID,
in.GetChannelNames(),
)
if err != nil {
log.Warn("failed to broadcast import message", zap.Error(err))
resp.Status = merr.Status(merr.Wrap(err, "failed to broadcast import"))
return resp, nil
}
resp.JobID = fmt.Sprint(jobID)
log.Info("import request broadcasted successfully", zap.String("jobID", resp.JobID))
return resp, nil
}
// createImportJobFromAck creates an import job from ack callback.
// This is called internally when broadcast ack is received.
func (s *Server) createImportJobFromAck(ctx context.Context, in *internalpb.ImportRequestInternal) (*internalpb.ImportResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &internalpb.ImportResponse{
Status: merr.Status(err),
}, nil
}
resp := &internalpb.ImportResponse{
Status: merr.Success(),
}
log := log.Ctx(ctx).With(zap.Int64("collection", in.GetCollectionID()),
zap.Int64s("partitions", in.GetPartitionIDs()),
zap.Strings("channels", in.GetChannelNames()))
log.Info("creating import job from ack callback",
zap.Int("fileNum", len(in.GetFiles())),
zap.Any("files", in.GetFiles()),
zap.Any("options", in.GetOptions()))
timeoutTs, err := importutilv2.GetTimeoutTs(in.GetOptions())
if err != nil {
resp.Status = merr.Status(merr.WrapErrImportFailed(err.Error()))
return resp, nil
}
files := in.GetFiles()
isBackup := importutilv2.IsBackup(in.GetOptions())
if isBackup {
files, err = ListBinlogImportRequestFiles(ctx, s.meta.chunkManager, files, in.GetOptions())
if err != nil {
resp.Status = merr.Status(err)
return resp, nil
}
}
// Allocate file ids.
idStart, _, err := s.allocator.AllocN(int64(len(files)) + 1)
if err != nil {
resp.Status = merr.Status(merr.Wrap(err, "alloc id failed"))
return resp, nil
}
files = lo.Map(files, func(importFile *internalpb.ImportFile, i int) *internalpb.ImportFile {
importFile.Id = idStart + int64(i) + 1
return importFile
})
importCollectionInfo, err := s.handler.GetCollection(ctx, in.GetCollectionID())
if errors.Is(err, merr.ErrCollectionNotFound) {
resp.Status = merr.Status(merr.WrapErrCollectionNotFound(in.GetCollectionID()))
return resp, nil
}
if err != nil {
resp.Status = merr.Status(merr.Wrap(err, "get collection failed"))
return resp, nil
}
if importCollectionInfo == nil {
resp.Status = merr.Status(merr.WrapErrCollectionNotFound(in.GetCollectionID()))
return resp, nil
}
jobID := in.GetJobID()
if jobID == 0 {
jobID = idStart
}
createTime := time.Now()
job := &importJob{
ImportJob: &datapb.ImportJob{
JobID: jobID,
CollectionID: in.GetCollectionID(),
CollectionName: in.GetCollectionName(),
PartitionIDs: in.GetPartitionIDs(),
Vchannels: importCollectionInfo.VChannelNames,
Schema: in.GetSchema(),
TimeoutTs: timeoutTs,
CleanupTs: math.MaxUint64,
State: internalpb.ImportJobState_Pending,
Files: files,
Options: in.GetOptions(),
CreateTime: createTime.Format("2006-01-02T15:04:05Z07:00"),
ReadyVchannels: in.GetChannelNames(),
DataTs: in.GetDataTimestamp(),
AutoCommit: importutilv2.IsAutoCommit(in.GetOptions()),
},
tr: timerecord.NewTimeRecorder("import job"),
}
err = s.importMeta.AddJob(ctx, job)
if err != nil {
resp.Status = merr.Status(merr.Wrap(err, "add import job failed"))
return resp, nil
}
resp.JobID = fmt.Sprint(job.GetJobID())
log.Info("add import job done",
zap.Int64("jobID", job.GetJobID()),
zap.Int("fileNum", len(files)),
zap.Any("files", files),
zap.Strings("readyChannels", in.GetChannelNames()),
)
return resp, nil
}
func (s *Server) GetImportProgress(ctx context.Context, in *internalpb.GetImportProgressRequest) (*internalpb.GetImportProgressResponse, error) {
log := log.Ctx(ctx).With(zap.String("jobID", in.GetJobID()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &internalpb.GetImportProgressResponse{
Status: merr.Status(err),
}, nil
}
resp := &internalpb.GetImportProgressResponse{
Status: merr.Success(),
}
jobID, err := strconv.ParseInt(in.GetJobID(), 10, 64)
if err != nil {
resp.Status = merr.Status(merr.WrapErrParameterInvalidMsg("parse job id failed: %v", err))
return resp, nil
}
job := s.importMeta.GetJob(ctx, jobID)
if job == nil {
resp.Status = merr.Status(merr.WrapErrImportSysFailedMsg("import job does not exist, jobID=%d", jobID))
return resp, nil
}
progress, state, importedRows, totalRows, reason := GetJobProgress(ctx, jobID, s.importMeta, s.meta)
resp.State = state
resp.Reason = reason
resp.Progress = progress
resp.CollectionName = job.GetCollectionName()
resp.CreateTime = job.GetCreateTime()
resp.CompleteTime = job.GetCompleteTime()
resp.ImportedRows = importedRows
resp.TotalRows = totalRows
resp.TaskProgresses = GetTaskProgresses(ctx, jobID, s.importMeta, s.meta)
log.Info("GetImportProgress done", zap.String("jobState", job.GetState().String()), zap.Any("resp", resp))
return resp, nil
}
func (s *Server) ListImports(ctx context.Context, req *internalpb.ListImportsRequestInternal) (*internalpb.ListImportsResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &internalpb.ListImportsResponse{
Status: merr.Status(err),
}, nil
}
resp := &internalpb.ListImportsResponse{
Status: merr.Success(),
JobIDs: make([]string, 0),
States: make([]internalpb.ImportJobState, 0),
Reasons: make([]string, 0),
Progresses: make([]int64, 0),
}
var jobs []ImportJob
if req.GetCollectionID() != 0 {
jobs = s.importMeta.GetJobBy(ctx, WithCollectionID(req.GetCollectionID()))
} else {
jobs = s.importMeta.GetJobBy(ctx)
}
for _, job := range jobs {
progress, state, _, _, reason := GetJobProgress(ctx, job.GetJobID(), s.importMeta, s.meta)
resp.JobIDs = append(resp.JobIDs, fmt.Sprintf("%d", job.GetJobID()))
resp.States = append(resp.States, state)
resp.Reasons = append(resp.Reasons, reason)
resp.Progresses = append(resp.Progresses, progress)
resp.CollectionNames = append(resp.CollectionNames, job.GetCollectionName())
}
return resp, nil
}
// NotifyDropPartition notifies DataCoord to drop segments of specified partition
func (s *Server) NotifyDropPartition(ctx context.Context, channel string, partitionIDs []int64) error {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return err
}
log.Ctx(ctx).Info("receive NotifyDropPartition request",
zap.String("channelname", channel),
zap.Any("partitionID", partitionIDs))
s.segmentManager.DropSegmentsOfPartition(ctx, channel, partitionIDs)
// release all segments of the partition.
return s.meta.DropSegmentsOfPartition(ctx, partitionIDs)
}
// DropSegmentsByTime drop segments that were updated before the flush timestamp for TruncateCollection
func (s *Server) DropSegmentsByTime(ctx context.Context, collectionID int64, flushTsList map[string]uint64) error {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return err
}
log.Ctx(ctx).Info("receive DropSegmentsByTime request",
zap.Int64("collectionID", collectionID))
for channelName, flushTs := range flushTsList {
// wait until the checkpoint reaches or exceeds the flush timestamp
err := s.meta.WatchChannelCheckpoint(ctx, channelName, flushTs)
if err != nil {
log.Ctx(ctx).Warn("WatchChannelCheckpoint failed", zap.Error(err))
return err
}
// drop segments that were updated before the flush timestamp
err = s.meta.TruncateChannelByTime(ctx, channelName, flushTs)
if err != nil {
log.Warn("TruncateChannelByTime failed", zap.Error(err))
return err
}
}
return nil
}
func (s *Server) CreateSnapshot(ctx context.Context, req *datapb.CreateSnapshotRequest) (*commonpb.Status, error) {
log := log.Ctx(ctx).With(zap.Int64("collectionID", req.GetCollectionId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
log.Info("receive CreateSnapshot request", zap.String("name", req.GetName()),
zap.String("description", req.GetDescription()),
zap.Int64("compactionProtectionSeconds", req.GetCompactionProtectionSeconds()))
// Defense-in-depth: re-validate compaction_protection_seconds on the DataCoord side.
// Proxy also validates this, but a buggy or malicious client could bypass Proxy by
// calling this RPC directly. Range validation must be enforced by the owner of the feature.
if req.GetCompactionProtectionSeconds() < 0 {
return merr.Status(merr.WrapErrParameterInvalidMsg("compaction_protection_seconds must be non-negative")), nil
}
maxCompactionProtectionSeconds := paramtable.Get().DataCoordCfg.SnapshotMaxCompactionProtectionSeconds.GetAsInt64()
if req.GetCompactionProtectionSeconds() > maxCompactionProtectionSeconds {
return merr.Status(merr.WrapErrParameterInvalidMsg(
fmt.Sprintf("compaction_protection_seconds must not exceed %d", maxCompactionProtectionSeconds))), nil
}
// Check if snapshot name already exists within this collection.
// Distinguish ErrSnapshotNotFound (the only good case to proceed) from any
// other error (etcd timeout, decode failure, ctx cancel) — those must surface
// rather than be silently swallowed by falling through to the broadcast path.
if _, err := s.snapshotManager.GetSnapshot(ctx, req.GetCollectionId(), req.GetName()); err == nil {
log.Warn("CreateSnapshot failed: snapshot name already exists in collection")
return merr.Status(merr.WrapErrParameterInvalidMsg("snapshot name %s already exists in collection %d", req.GetName(), req.GetCollectionId())), nil
} else if !errors.Is(err, merr.ErrSnapshotNotFound) {
log.Warn("CreateSnapshot: failed to check snapshot existence", zap.Error(err))
return merr.Status(err), nil
}
// Resolve collection identity for the broadcast lock set (also validates
// collection existence). Read from the datacoord-local meta cache via
// handler.GetCollection — avoids a cross-component RPC to MixCoord on every
// CreateSnapshot, and on cache miss handler.GetCollection transparently
// falls back to rootcoord with bounded retries.
coll, err := s.handler.GetCollection(ctx, req.GetCollectionId())
if err != nil {
log.Warn("CreateSnapshot failed to resolve collection", zap.Error(err))
return merr.Status(err), nil
}
if coll == nil {
log.Warn("CreateSnapshot: collection not found")
return merr.Status(merr.WrapErrCollectionNotFound(req.GetCollectionId())), nil
}
dbName := coll.DatabaseName
collectionName := coll.Schema.GetName()
broadcaster, err := broadcast.StartBroadcastWithResourceKeys(ctx,
message.NewSharedDBNameResourceKey(dbName),
message.NewExclusiveCollectionNameResourceKey(dbName, collectionName),
message.NewExclusiveSnapshotNameResourceKey(req.GetCollectionId(), req.GetName()),
)
if err != nil {
log.Warn("CreateSnapshot failed to start broadcast", zap.Error(err))
return merr.Status(err), nil
}
defer broadcaster.Close()
// Re-check collection availability while holding the collection resource
// lock. DropCollection may win the race between the pre-lock collection
// resolution above and lock acquisition; in that case this request must
// terminate before broadcasting a CreateSnapshot message whose ack callback
// can no longer generate a valid snapshot.
hasCollection, err := s.broker.HasCollection(ctx, req.GetCollectionId())
if err != nil {
log.Warn("CreateSnapshot: failed to re-check collection existence after lock", zap.Error(err))
return merr.Status(err), nil
}
if !hasCollection {
log.Warn("CreateSnapshot: collection not found after lock")
return merr.Status(merr.WrapErrCollectionNotFound(req.GetCollectionId())), nil
}
// Double-check after acquiring lock — another goroutine may have created it.
// Same error-handling discipline as the pre-lock check above: only treat
// ErrSnapshotNotFound as "good to proceed"; surface every other error.
if _, err := s.snapshotManager.GetSnapshot(ctx, req.GetCollectionId(), req.GetName()); err == nil {
log.Warn("CreateSnapshot failed: snapshot name already exists in collection")
return merr.Status(merr.WrapErrParameterInvalidMsg("snapshot name %s already exists in collection %d", req.GetName(), req.GetCollectionId())), nil
} else if !errors.Is(err, merr.ErrSnapshotNotFound) {
log.Warn("CreateSnapshot: failed to re-check snapshot existence after lock", zap.Error(err))
return merr.Status(err), nil
}
// Broadcast CreateSnapshot message via DDL framework
// Snapshot ID is allocated in the callback
if _, err := broadcaster.Broadcast(ctx, message.NewCreateSnapshotMessageBuilderV2().
WithHeader(&message.CreateSnapshotMessageHeader{
CollectionId: req.GetCollectionId(),
Name: req.GetName(),
Description: req.GetDescription(),
CompactionProtectionSeconds: req.GetCompactionProtectionSeconds(),
}).
WithBody(&message.CreateSnapshotMessageBody{}).
WithBroadcast([]string{streaming.WAL().ControlChannel()}).
MustBuildBroadcast(),
); err != nil {
log.Error("CreateSnapshot broadcast failed", zap.Error(err))
return merr.Status(err), nil
}
log.Info("CreateSnapshot completed successfully")
return merr.Success(), nil
}
func (s *Server) BatchUpdateManifest(ctx context.Context, req *datapb.BatchUpdateManifestRequest) (*commonpb.Status, error) {
log := log.Ctx(ctx).With(zap.Int64("collectionID", req.GetCollectionId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
log.Info("receive BatchUpdateManifest request", zap.Int("itemCount", len(req.GetItems())))
coll, err := s.broker.DescribeCollectionInternal(ctx, req.GetCollectionId())
if err != nil {
log.Warn("BatchUpdateManifest failed to describe collection", zap.Error(err))
return merr.Status(err), nil
}
dbName := coll.GetDbName()
collectionName := coll.GetCollectionName()
broadcaster, err := broadcast.StartBroadcastWithResourceKeys(ctx,
message.NewSharedDBNameResourceKey(dbName),
message.NewSharedCollectionNameResourceKey(dbName, collectionName),
)
if err != nil {
log.Warn("BatchUpdateManifest failed to start broadcast", zap.Error(err))
return merr.Status(err), nil
}
defer broadcaster.Close()
items := make([]*messagespb.BatchUpdateManifestItem, 0, len(req.GetItems()))
for _, item := range req.GetItems() {
items = append(items, &messagespb.BatchUpdateManifestItem{
SegmentId: item.GetSegmentId(),
ManifestVersion: item.GetManifestVersion(),
})
}
if _, err := broadcaster.Broadcast(ctx, message.NewBatchUpdateManifestMessageBuilderV2().
WithHeader(&message.BatchUpdateManifestMessageHeader{
CollectionId: req.GetCollectionId(),
}).
WithBody(&message.BatchUpdateManifestMessageBody{
Items: items,
}).
WithBroadcast([]string{streaming.WAL().ControlChannel()}).
MustBuildBroadcast(),
); err != nil {
log.Error("BatchUpdateManifest broadcast failed", zap.Error(err))
return merr.Status(err), nil
}
log.Info("BatchUpdateManifest completed successfully")
return merr.Success(), nil
}
func (s *Server) DropSnapshot(ctx context.Context, req *datapb.DropSnapshotRequest) (*commonpb.Status, error) {
log := log.Ctx(ctx).With(zap.String("snapshot", req.GetName()), zap.Int64("collectionID", req.GetCollectionId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
log.Info("receive DropSnapshot request")
// Check if snapshot exists - if not, return success (idempotent).
// Only treat ErrSnapshotNotFound as idempotent success; other errors (etcd timeout,
// decode failure, ctx cancel) must surface so they can be retried or reported.
if _, err := s.snapshotManager.GetSnapshot(ctx, req.GetCollectionId(), req.GetName()); err != nil {
if errors.Is(err, merr.ErrSnapshotNotFound) {
log.Info("DropSnapshot: snapshot not found, returning success (idempotent)")
return merr.Success(), nil
}
log.Warn("DropSnapshot: failed to check snapshot existence", zap.Error(err))
return merr.Status(err), nil
}
// Resolve collection identity so the broadcast can acquire proper DB/collection locks.
// Dropping a snapshot must serialize against concurrent DropCollection / AlterCollection
// on the owning collection, otherwise DropSnapshot can race with a cascade drop path.
//
// Read from the datacoord-local meta cache (populated at startup and kept in sync
// via BroadcastAlteredCollection). This avoids a cross-component RPC to MixCoord
// on every DropSnapshot, and on cache miss handler.GetCollection transparently
// falls back to rootcoord with bounded retries.
coll, err := s.handler.GetCollection(ctx, req.GetCollectionId())
if err != nil {
log.Warn("DropSnapshot failed to resolve collection", zap.Error(err))
return merr.Status(err), nil
}
if coll == nil {
log.Warn("DropSnapshot: collection not found")
return merr.Status(merr.WrapErrCollectionNotFound(req.GetCollectionId())), nil
}
dbName := coll.DatabaseName
collectionName := coll.Schema.GetName()
// Start broadcast with DB + collection + per-collection snapshot locks.
// The snapshot resource key is namespaced by collectionID so DropSnapshot on
// (collA, "backup") does not falsely block DropSnapshot on (collB, "backup").
broadcaster, err := broadcast.StartBroadcastWithResourceKeys(ctx,
message.NewSharedDBNameResourceKey(dbName),
message.NewExclusiveCollectionNameResourceKey(dbName, collectionName),
message.NewExclusiveSnapshotNameResourceKey(req.GetCollectionId(), req.GetName()),
)
if err != nil {
log.Error("DropSnapshot failed to start broadcast", zap.Error(err))
return merr.Status(err), nil
}
defer broadcaster.Close()
// Double-check after acquiring lock - another goroutine may have dropped it.
if _, err := s.snapshotManager.GetSnapshot(ctx, req.GetCollectionId(), req.GetName()); err != nil {
if errors.Is(err, merr.ErrSnapshotNotFound) {
log.Info("DropSnapshot: snapshot not found after lock, returning success (idempotent)")
return merr.Success(), nil
}
log.Warn("DropSnapshot: failed to re-check snapshot existence after lock", zap.Error(err))
return merr.Status(err), nil
}
// Pre-flight pin check UNDER the broadcast lock.
//
// Why it must be under the lock: PinSnapshotData acquires the same
// (collectionID, snapshotName) resource key in shared mode. Because we
// hold the key in EXCLUSIVE mode here, no concurrent Pin can add new pins
// between this check and the ack callback execution — the pin count we
// observe now is authoritative for the rest of this broadcast. This closes
// the race that would otherwise let a Pin slip in, cause the ack callback
// to observe ErrSnapshotPinned, and trigger retry-forever (holding this
// very lock → throughput deadlock on the snapshot control plane).
//
// Rejecting here (instead of in the ack callback) means clients get a
// synchronous, non-retrying error and the broadcast is never initiated.
pinned, err := s.snapshotManager.HasActivePins(ctx, req.GetCollectionId(), req.GetName())
if err != nil {
log.Warn("DropSnapshot: failed to check active pins", zap.Error(err))
return merr.Status(err), nil
}
if pinned {
log.Warn("DropSnapshot rejected: snapshot has active pins")
return merr.Status(merr.WrapErrSnapshotPinned(req.GetName(), "active pins exist, unpin before dropping")), nil
}
// Broadcast DropSnapshot message via DDL framework
if _, err := broadcaster.Broadcast(ctx, message.NewDropSnapshotMessageBuilderV2().
WithHeader(&message.DropSnapshotMessageHeader{
Name: req.GetName(),
CollectionId: req.GetCollectionId(),
}).
WithBody(&message.DropSnapshotMessageBody{}).
WithBroadcast([]string{streaming.WAL().ControlChannel()}).
MustBuildBroadcast(),
); err != nil {
log.Error("DropSnapshot broadcast failed", zap.Error(err))
return merr.Status(err), nil
}
log.Info("DropSnapshot completed successfully")
return merr.Success(), nil
}
func (s *Server) DescribeSnapshot(ctx context.Context, req *datapb.DescribeSnapshotRequest) (*datapb.DescribeSnapshotResponse, error) {
log := log.Ctx(ctx).With(zap.String("snapshotName", req.GetName()), zap.Int64("collectionID", req.GetCollectionId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.DescribeSnapshotResponse{
Status: merr.Status(err),
}, nil
}
log.Info("receive DescribeSnapshot request")
// Delegate to SnapshotManager
snapshotData, err := s.snapshotManager.DescribeSnapshot(ctx, req.GetCollectionId(), req.GetName())
if err != nil {
log.Error("failed to describe snapshot", zap.Error(err))
return &datapb.DescribeSnapshotResponse{
Status: merr.Status(err),
}, nil
}
resp := &datapb.DescribeSnapshotResponse{
Status: merr.Success(),
SnapshotInfo: snapshotData.SnapshotInfo,
}
if req.GetIncludeCollectionInfo() {
resp.CollectionInfo = snapshotData.Collection
resp.IndexInfos = snapshotData.Indexes
}
return resp, nil
}
// RestoreSnapshot restores snapshot data to a new collection.
// This method validates parameters and delegates to snapshotManager for the actual restore.
func (s *Server) RestoreSnapshot(ctx context.Context, req *datapb.RestoreSnapshotRequest) (*datapb.RestoreSnapshotResponse, error) {
log := log.Ctx(ctx).With(
zap.String("snapshot", req.GetName()),
zap.Int64("sourceCollectionID", req.GetSourceCollectionId()),
zap.String("targetDbName", req.GetTargetDbName()),
zap.String("targetCollectionName", req.GetTargetCollectionName()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.RestoreSnapshotResponse{
Status: merr.Status(err),
}, nil
}
log.Info("receive RestoreSnapshot request")
// Validate parameters
if req.GetName() == "" {
err := merr.WrapErrParameterMissingMsg("snapshot name is required")
log.Warn("invalid request", zap.Error(err))
return &datapb.RestoreSnapshotResponse{
Status: merr.Status(err),
}, nil
}
if req.GetTargetCollectionName() == "" {
err := merr.WrapErrParameterMissingMsg("target collection name is required")
log.Warn("invalid request", zap.Error(err))
return &datapb.RestoreSnapshotResponse{
Status: merr.Status(err),
}, nil
}
// Delegate to snapshot manager
jobID, err := s.snapshotManager.RestoreSnapshot(
ctx,
req.GetSourceCollectionId(),
req.GetName(),
req.GetTargetCollectionName(),
req.GetTargetDbName(),
s.startRestoreSnapshotLock,
s.startBroadcastForRestoreSnapshot,
s.rollbackRestoreSnapshot,
s.validateRestoreSnapshotResources,
)
if err != nil {
log.Error("restore snapshot failed", zap.Error(err))
return &datapb.RestoreSnapshotResponse{
Status: merr.Status(err),
}, nil
}
log.Info("restore snapshot completed", zap.Int64("jobID", jobID))
return &datapb.RestoreSnapshotResponse{
Status: merr.Success(),
JobId: jobID,
}, nil
}
// rollbackRestoreSnapshot drops the newly created collection when restore fails.
func (s *Server) rollbackRestoreSnapshot(ctx context.Context, dbName, collectionName string) error {
log := log.Ctx(ctx).With(
zap.String("dbName", dbName),
zap.String("collectionName", collectionName),
)
log.Info("rolling back restore snapshot, dropping collection")
if err := s.broker.DropCollection(ctx, dbName, collectionName); err != nil {
if errors.Is(err, merr.ErrCollectionNotFound) {
log.Debug("collection not found, skipping rollback")
return nil
}
log.Error("failed to drop collection during rollback", zap.Error(err))
return err
}
log.Info("rollback completed, collection dropped")
return nil
}
func (s *Server) GetRestoreSnapshotState(ctx context.Context, req *datapb.GetRestoreSnapshotStateRequest) (*datapb.GetRestoreSnapshotStateResponse, error) {
log := log.Ctx(ctx).With(zap.Int64("jobID", req.GetJobId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetRestoreSnapshotStateResponse{
Status: merr.Status(err),
}, nil
}
// Delegate to SnapshotManager
restoreInfo, err := s.snapshotManager.GetRestoreState(ctx, req.GetJobId())
if err != nil {
log.Warn("failed to get restore state", zap.Error(err))
return &datapb.GetRestoreSnapshotStateResponse{
Status: merr.Status(err),
}, nil
}
log.Info("get restore state completed",
zap.String("state", restoreInfo.GetState().String()),
zap.Int32("progress", restoreInfo.GetProgress()))
return &datapb.GetRestoreSnapshotStateResponse{
Status: merr.Success(),
Info: restoreInfo,
}, nil
}
func (s *Server) ListRestoreSnapshotJobs(ctx context.Context, req *datapb.ListRestoreSnapshotJobsRequest) (*datapb.ListRestoreSnapshotJobsResponse, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.ListRestoreSnapshotJobsResponse{
Status: merr.Status(err),
}, nil
}
// Delegate to SnapshotManager
restoreInfos, err := s.snapshotManager.ListRestoreJobs(ctx, req.GetCollectionId(), req.GetDbId())
if err != nil {
log.Ctx(ctx).Error("failed to list restore jobs", zap.Error(err))
return &datapb.ListRestoreSnapshotJobsResponse{
Status: merr.Status(err),
}, nil
}
log.Ctx(ctx).Info("list restore jobs completed",
zap.Int("totalJobs", len(restoreInfos)),
zap.Int64("filterCollectionId", req.GetCollectionId()))
return &datapb.ListRestoreSnapshotJobsResponse{
Status: merr.Success(),
Jobs: restoreInfos,
}, nil
}
func (s *Server) ListSnapshots(ctx context.Context, req *datapb.ListSnapshotsRequest) (*datapb.ListSnapshotsResponse, error) {
log := log.Ctx(ctx).With(zap.Int64("collectionID", req.GetCollectionId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.ListSnapshotsResponse{
Status: merr.Status(err),
}, nil
}
log.Info("receive ListSnapshots request")
// Delegate to SnapshotManager
snapshots, err := s.snapshotManager.ListSnapshots(ctx, req.GetCollectionId(), req.GetPartitionId(), req.GetDbId())
if err != nil {
log.Error("failed to list snapshots", zap.Error(err))
return &datapb.ListSnapshotsResponse{
Status: merr.Status(err),
}, nil
}
return &datapb.ListSnapshotsResponse{
Status: merr.Success(),
Snapshots: snapshots,
}, nil
}
func (s *Server) PinSnapshotData(ctx context.Context, req *datapb.PinSnapshotDataRequest) (*datapb.PinSnapshotDataResponse, error) {
log := log.Ctx(ctx).With(
zap.String("snapshotName", req.GetName()),
zap.Int64("collectionID", req.GetCollectionId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.PinSnapshotDataResponse{
Status: merr.Status(err),
}, nil
}
log.Info("receive PinSnapshotData request")
// Resolve collection identity for the resource key lock. We need the
// collection to exist before we can acquire the locks that serialize Pin
// against DropSnapshot / DropCollection.
//
// Read from the datacoord-local meta cache (populated at startup and kept in sync
// via BroadcastAlteredCollection). This avoids a cross-component RPC to MixCoord
// on every PinSnapshotData, and on cache miss handler.GetCollection transparently
// falls back to rootcoord with bounded retries.
coll, err := s.handler.GetCollection(ctx, req.GetCollectionId())
if err != nil {
log.Warn("PinSnapshotData failed to resolve collection", zap.Error(err))
return &datapb.PinSnapshotDataResponse{
Status: merr.Status(err),
}, nil
}
if coll == nil {
log.Warn("PinSnapshotData: collection not found")
return &datapb.PinSnapshotDataResponse{
Status: merr.Status(merr.WrapErrCollectionNotFound(req.GetCollectionId())),
}, nil
}
dbName := coll.DatabaseName
collectionName := coll.Schema.GetName()
// Acquire a SHARED snapshot-name lock so that:
// - multiple concurrent Pins on the same snapshot can proceed in parallel
// (shared vs shared),
// - concurrent DropSnapshot (which takes EXCLUSIVE) blocks until all
// in-flight Pins release their shared hold, and vice versa.
//
// This closes the race where a Pin slips in between DropSnapshot's
// pre-flight pin check and the ack callback actually deleting the data —
// without this lock, a racing Pin could make the ack callback observe an
// active pin and either (a) leave an orphan pin on a deleted snapshot, or
// (b) trigger ErrSnapshotPinned inside the retry-forever loop.
//
// Lock-only mode: no Broadcast() is called, Close() just releases the keys.
locker, err := broadcast.StartBroadcastWithResourceKeys(ctx,
message.NewSharedDBNameResourceKey(dbName),
message.NewSharedCollectionNameResourceKey(dbName, collectionName),
message.NewSharedSnapshotNameResourceKey(req.GetCollectionId(), req.GetName()),
)
if err != nil {
log.Warn("PinSnapshotData failed to acquire resource key lock", zap.Error(err))
return &datapb.PinSnapshotDataResponse{
Status: merr.Status(err),
}, nil
}
defer locker.Close()
pinID, err := s.snapshotManager.PinSnapshotData(ctx, req.GetCollectionId(), req.GetName(), req.GetTtlSeconds())
if err != nil {
log.Error("failed to pin snapshot data", zap.Error(err))
return &datapb.PinSnapshotDataResponse{
Status: merr.Status(err),
}, nil
}
log.Info("PinSnapshotData completed successfully", zap.Int64("pinID", pinID))
return &datapb.PinSnapshotDataResponse{
Status: merr.Success(),
PinId: pinID,
}, nil
}
func (s *Server) UnpinSnapshotData(ctx context.Context, req *datapb.UnpinSnapshotDataRequest) (*commonpb.Status, error) {
log := log.Ctx(ctx).With(
zap.Int64("pinID", req.GetPinId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
log.Info("receive UnpinSnapshotData request")
if err := s.snapshotManager.UnpinSnapshotData(ctx, req.GetPinId()); err != nil {
log.Error("failed to unpin snapshot data", zap.Error(err))
return merr.Status(err), nil
}
log.Info("UnpinSnapshotData completed successfully")
return merr.Success(), nil
}
// RefreshExternalCollection manually triggers a refresh job for an external collection
// This uses WAL Broadcast mechanism for idempotency and distributed consistency.
func (s *Server) RefreshExternalCollection(ctx context.Context, req *datapb.RefreshExternalCollectionRequest) (*datapb.RefreshExternalCollectionResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("collectionID", req.GetCollectionId()),
zap.String("collectionName", req.GetCollectionName()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.RefreshExternalCollectionResponse{
Status: merr.Status(err),
}, nil
}
log.Info("receive RefreshExternalCollection request")
if s.externalCollectionRefreshManager == nil {
log.Warn("external collection refresh manager not initialized")
return &datapb.RefreshExternalCollectionResponse{
Status: merr.Status(merr.WrapErrServiceUnavailable("external collection refresh manager not initialized")),
}, nil
}
// Start broadcaster with resource lock (shared DB + exclusive collection)
b, err := s.startBroadcastWithCollectionID(ctx, req.GetCollectionId())
if err != nil {
log.Warn("failed to start broadcaster", zap.Error(err))
return &datapb.RefreshExternalCollectionResponse{
Status: merr.Status(err),
}, nil
}
defer b.Close()
// Synchronous duplicate detection at the RPC edge. The WAL ack callback
// already rejects duplicates, but it does so AFTER the RPC has already
// returned a fresh jobID to the client; that jobID then dies silently and
// the client polls it forever. Surfacing the in-progress jobID here
// (along with merr.ErrTaskDuplicate) lets the caller switch to polling
// the real job. The remaining TOCTOU window between this read and the
// ack-side AddJob falls back to the existing ack-side rejection, i.e.
// today's behavior — never worse.
if active := s.externalCollectionRefreshManager.GetActiveJobByCollectionID(req.GetCollectionId()); active != nil {
log.Info("refresh job already in progress, rejecting at RPC edge",
zap.Int64("existingJobID", active.GetJobId()),
zap.String("existingState", active.GetState().String()))
return &datapb.RefreshExternalCollectionResponse{
Status: merr.Status(merr.WrapErrTaskDuplicate(
"refresh_external_collection",
fmt.Sprintf("refresh job %d is already in progress for collection %s; poll that jobID or wait for it to complete",
active.GetJobId(), req.GetCollectionName()))),
JobId: active.GetJobId(),
}, nil
}
// Pre-allocate JobID for idempotency (ensures same JobID even if retry after failure)
allocatedJobID, err := s.allocator.AllocID(ctx)
if err != nil {
log.Warn("failed to allocate job ID", zap.Error(err))
return &datapb.RefreshExternalCollectionResponse{
Status: merr.Status(err),
}, nil
}
log.Info("pre-allocated job ID for refresh", zap.Int64("jobID", allocatedJobID))
// Build and broadcast the message
msg := message.NewRefreshExternalCollectionMessageBuilderV2().
WithHeader(&message.RefreshExternalCollectionMessageHeader{
CollectionId: req.GetCollectionId(),
CollectionName: req.GetCollectionName(),
JobId: allocatedJobID,
ExternalSource: req.GetExternalSource(),
ExternalSpec: req.GetExternalSpec(),
}).
WithBody(&message.RefreshExternalCollectionMessageBody{}).
WithBroadcast([]string{streaming.WAL().ControlChannel()}).
MustBuildBroadcast()
if _, err := b.Broadcast(ctx, msg); err != nil {
log.Warn("failed to broadcast refresh message", zap.Error(err))
return &datapb.RefreshExternalCollectionResponse{
Status: merr.Status(err),
}, nil
}
log.Info("refresh external collection job submitted via WAL broadcast", zap.Int64("jobID", allocatedJobID))
return &datapb.RefreshExternalCollectionResponse{
Status: merr.Success(),
JobId: allocatedJobID,
}, nil
}
// GetRefreshExternalCollectionProgress returns the progress of a refresh job
func (s *Server) GetRefreshExternalCollectionProgress(ctx context.Context, req *datapb.GetRefreshExternalCollectionProgressRequest) (*datapb.GetRefreshExternalCollectionProgressResponse, error) {
log := log.Ctx(ctx).With(zap.Int64("jobID", req.GetJobId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.GetRefreshExternalCollectionProgressResponse{
Status: merr.Status(err),
}, nil
}
log.Info("receive GetRefreshExternalCollectionProgress request")
if s.externalCollectionRefreshManager == nil {
log.Warn("external collection refresh manager not initialized")
return &datapb.GetRefreshExternalCollectionProgressResponse{
Status: merr.Status(merr.WrapErrServiceUnavailable("external collection refresh manager not initialized")),
}, nil
}
jobInfo, err := s.externalCollectionRefreshManager.GetJobProgress(ctx, req.GetJobId())
if err != nil {
log.Warn("failed to get job progress", zap.Error(err))
return &datapb.GetRefreshExternalCollectionProgressResponse{
Status: merr.Status(err),
}, nil
}
log.Info("get refresh external collection progress completed",
zap.String("state", jobInfo.GetState().String()),
zap.Int64("progress", jobInfo.GetProgress()))
return &datapb.GetRefreshExternalCollectionProgressResponse{
Status: merr.Success(),
JobInfo: jobInfo,
}, nil
}
// ListRefreshExternalCollectionJobs lists refresh jobs for a collection.
// A zero collectionID lists jobs for all external collections.
func (s *Server) ListRefreshExternalCollectionJobs(ctx context.Context, req *datapb.ListRefreshExternalCollectionJobsRequest) (*datapb.ListRefreshExternalCollectionJobsResponse, error) {
log := log.Ctx(ctx).With(
zap.Int64("collectionID", req.GetCollectionId()))
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return &datapb.ListRefreshExternalCollectionJobsResponse{
Status: merr.Status(err),
}, nil
}
log.Info("receive ListRefreshExternalCollectionJobs request")
if s.externalCollectionRefreshManager == nil {
log.Warn("external collection refresh manager not initialized")
return &datapb.ListRefreshExternalCollectionJobsResponse{
Status: merr.Status(merr.WrapErrServiceUnavailable("external collection refresh manager not initialized")),
}, nil
}
jobs, err := s.externalCollectionRefreshManager.ListJobs(ctx, req.GetCollectionId())
if err != nil {
log.Warn("failed to list jobs", zap.Error(err))
return &datapb.ListRefreshExternalCollectionJobsResponse{
Status: merr.Status(err),
}, nil
}
log.Info("list refresh external collection jobs completed", zap.Int("jobCount", len(jobs)))
return &datapb.ListRefreshExternalCollectionJobsResponse{
Status: merr.Success(),
Jobs: jobs,
}, nil
}
// broadcastCommitImportMessage broadcasts a CommitImport WAL message for the given import job.
// The message is broadcast to the job's data vchannels so each vchannel's WAL flusher
// can observe the commit fence, flush pending DML, and call HandleCommitVchannel.
// (Control-channel-only broadcast is dropped by the flusher's IsControlChannel guard
// before reaching the CommitImport case, so it cannot drive per-vchannel commits.)
func (s *Server) broadcastCommitImportMessage(ctx context.Context, job ImportJob) error {
vchannels := job.GetVchannels()
if len(vchannels) == 0 {
return merr.WrapErrImportSysFailedMsg("job %d has no vchannels", job.GetJobID())
}
broadcaster, err := s.startBroadcastWithCollectionID(ctx, job.GetCollectionID())
if err != nil {
return err
}
defer broadcaster.Close()
msg := message.NewCommitImportMessageBuilderV2().
WithHeader(&message.CommitImportMessageHeader{
CollectionId: job.GetCollectionID(),
JobId: job.GetJobID(),
}).
WithBody(&messagespb.CommitImportMessageBody{}).
WithBroadcast(vchannels).
MustBuildBroadcast()
_, err = broadcaster.Broadcast(ctx, msg)
return err
}
// broadcastRollbackImportMessage broadcasts a RollbackImport WAL message for the given import job.
// Targets the job's data vchannels, matching the CommitImport routing.
func (s *Server) broadcastRollbackImportMessage(ctx context.Context, job ImportJob) error {
vchannels := job.GetVchannels()
if len(vchannels) == 0 {
return merr.WrapErrImportSysFailedMsg("job %d has no vchannels", job.GetJobID())
}
broadcaster, err := s.startBroadcastWithCollectionID(ctx, job.GetCollectionID())
if err != nil {
return err
}
defer broadcaster.Close()
msg := message.NewRollbackImportMessageBuilderV2().
WithHeader(&message.RollbackImportMessageHeader{
CollectionId: job.GetCollectionID(),
JobId: job.GetJobID(),
}).
WithBody(&messagespb.RollbackImportMessageBody{}).
WithBroadcast(vchannels).
MustBuildBroadcast()
_, err = broadcaster.Broadcast(ctx, msg)
return err
}
// validateAndExecuteImportAction handles the boilerplate for commit/abort import operations:
// health check, get job, auto-commit guard, per-job keylock with TOCTOU re-validation, and action execution.
func (s *Server) validateAndExecuteImportAction(
ctx context.Context,
jobID int64,
validateState func(job ImportJob) *commonpb.Status,
action func(ctx context.Context, job ImportJob) error,
) (*commonpb.Status, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
job := s.importMeta.GetJob(ctx, jobID)
if job == nil {
return merr.Status(merr.WrapErrImportSysFailedMsg("job %d not found", jobID)), nil
}
if st := validateState(job); st != nil {
return st, nil
}
if job.GetAutoCommit() {
return merr.Status(merr.WrapErrImportFailed(
fmt.Sprintf("job %d is auto-commit, manual commit/abort not allowed", jobID))), nil
}
s.importJobLock.Lock(jobID)
defer s.importJobLock.Unlock(jobID)
job = s.importMeta.GetJob(ctx, jobID)
if job == nil {
return merr.Status(merr.WrapErrImportFailed("job not found after lock")), nil
}
if st := validateState(job); st != nil {
return st, nil
}
if err := action(ctx, job); err != nil {
return merr.Status(err), nil
}
return merr.Success(), nil
}
// CommitImport commits a 2PC import job so that the imported data becomes visible.
// It transitions the job from Uncommitted → Committing by broadcasting a CommitImport WAL message.
// The call is idempotent: if the job is already Committing or Completed it returns success.
func (s *Server) CommitImport(ctx context.Context, req *datapb.CommitImportRequest) (*commonpb.Status, error) {
log := log.Ctx(ctx).With(zap.Int64("jobID", req.GetJobId()))
return s.validateAndExecuteImportAction(ctx, req.GetJobId(),
func(job ImportJob) *commonpb.Status {
switch job.GetState() {
case internalpb.ImportJobState_Committing, internalpb.ImportJobState_Completed:
return merr.Success()
case internalpb.ImportJobState_Uncommitted:
return nil // proceed
default:
return merr.Status(merr.WrapErrImportFailed(
fmt.Sprintf("job %d is in state %s, expected Uncommitted", req.GetJobId(), job.GetState())))
}
},
func(ctx context.Context, job ImportJob) error {
log.Info("committing import job via WAL broadcast")
return s.broadcastCommitImportMessage(ctx, job)
},
)
}
// AbortImport aborts a 2PC import job that has not yet been committed.
// It broadcasts a RollbackImport WAL message to cancel the job.
// Returns an error if the job is already committed or committing.
func (s *Server) AbortImport(ctx context.Context, req *datapb.AbortImportRequest) (*commonpb.Status, error) {
log := log.Ctx(ctx).With(zap.Int64("jobID", req.GetJobId()))
return s.validateAndExecuteImportAction(ctx, req.GetJobId(),
func(job ImportJob) *commonpb.Status {
state := job.GetState()
if state == internalpb.ImportJobState_Failed ||
state == internalpb.ImportJobState_Committing ||
state == internalpb.ImportJobState_Completed {
return merr.Status(merr.WrapErrImportFailed(
fmt.Sprintf("job %d is in terminal/committed state %s, abort not allowed", req.GetJobId(), state)))
}
return nil // proceed
},
func(ctx context.Context, job ImportJob) error {
log.Info("aborting import job via WAL broadcast")
return s.broadcastRollbackImportMessage(ctx, job)
},
)
}
// HandleCommitVchannel records that a vchannel has processed the commit fence for a 2PC import job.
// When all vchannels have acknowledged, the import job transitions to Completed and segments become visible.
func (s *Server) HandleCommitVchannel(ctx context.Context, req *datapb.HandleCommitVchannelRequest) (*commonpb.Status, error) {
if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
return merr.Status(err), nil
}
jobID := req.GetJobId()
vchannel := req.GetVchannel()
// Pre-fetch segment IDs for this job+vchannel BEFORE calling HandleCommitVchannel.
// The callback must not access importMeta because HandleCommitVchannel holds m.mu (write lock);
// calling GetTaskBy inside the callback would attempt to re-acquire m.mu (read lock) → deadlock.
segIDs := s.getImportSegmentIDsByVchannel(ctx, jobID, vchannel)
commitTs := req.GetCommitTimestamp()
err := s.importMeta.HandleCommitVchannel(ctx, jobID, vchannel, func() error {
// Only access s.meta (segment meta) here, NOT s.importMeta.
// Set CommitTimestamp and clear isImporting in a single call per segment.
ops := make([]UpdateOperator, 0, len(segIDs)*2)
for _, segID := range segIDs {
ops = append(ops,
UpdateCommitTimestamp(segID, commitTs),
UpdateIsImporting(segID, false),
)
}
if len(ops) == 0 {
return nil
}
return s.meta.UpdateSegmentsInfo(ctx, ops...)
})
if err != nil {
return merr.Status(err), nil
}
return merr.Success(), nil
}
// getImportSegmentIDsByVchannel returns all segment IDs (including sorted segments) belonging to
// the given import job that are assigned to the given vchannel.
// This must be called BEFORE acquiring importMeta's mutex (i.e., before HandleCommitVchannel).
func (s *Server) getImportSegmentIDsByVchannel(ctx context.Context, jobID int64, vchannel string) []int64 {
tasks := s.importMeta.GetTaskBy(ctx, WithJob(jobID), WithType(ImportTaskType))
var segIDs []int64
for _, task := range tasks {
it, ok := task.(*importTask)
if !ok {
continue
}
// Collect all candidate segment IDs from this task (safe copies).
candidates := make([]int64, 0, len(it.GetSegmentIDs())+len(it.GetSortedSegmentIDs()))
candidates = append(candidates, it.GetSegmentIDs()...)
candidates = append(candidates, it.GetSortedSegmentIDs()...)
for _, segID := range candidates {
seg := s.meta.GetSegment(ctx, segID)
if seg == nil {
continue
}
if seg.GetInsertChannel() != vchannel {
continue
}
segIDs = append(segIDs, segID)
}
}
return segIDs
}