milvus-io/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-10-12 02:19:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add an in-tree extension mechanism for distributions (#52981)

Browse Source 
issue: #52979
design doc:
docs/design-docs/design_docs/20260831-in_tree_extension_mechanism.md

A distribution that compiles its own behavior into the milvus binary
needs little from milvus: a way to install its request hook, the
QueryNode's tuning hook and the cipher without a `.so`, a callback that
starts its control-plane engine when the coordinator becomes active and
stops it on shutdown, and a declaration that it runs the deployment
shape its coordinator behaviors are written for. This PR is that seam
and nothing else. A stock binary installs nothing and behaves exactly as
before.

An earlier revision of this PR shipped a table of typed capabilities
consulted from many places in the tree. Reviewing it against what the
distribution actually needs showed most entries were constants per
deployment (configuration), duplicates of what the request hook already
does, or answers that belonged on the wire. The table is gone; what
remains is small enough to state here.

## Public surface (`pkg/extension`)

```go
func SetHook(h hook.Hook)                       // hookutil prefers it over proxy.soPath
func InstalledHook() hook.Hook

type QueryHook interface { ... }                // the queryNode.soPath plug-in's method set
func SetQueryHook(h QueryHook)                  // the QueryNode prefers it over queryNode.soPath
func InstalledQueryHook() QueryHook

func SetCipher(c hook.Cipher)                   // hookutil prefers it over cipherPlugin.soPathGo
func InstalledCipher() hook.Cipher

type Coordinator interface {                    // the coordinator as its own clients see it
    rootcoordpb.RootCoordClient
    querypb.QueryCoordClient
    datapb.DataCoordClient
}
type CoordinatorEngine interface {
    Start(ctx context.Context, coord Coordinator) error
    Stop() error
}
func SetCoordinatorEngine(e CoordinatorEngine)
func InstalledCoordinatorEngine() CoordinatorEngine

func SetForm()                                  // marks a deployment form installed, in every role
func FormInstalled() bool                       // the coordinators' form-gated behaviors read it

func WithQueryResourceGroup(ctx, rg) context.Context   // set by a hook's Before
func QueryResourceGroupFromContext(ctx) string
```

A distribution calls the setters it needs, then
`cmd/milvus.Main(os.Args)` (exported so a distribution can ship its own
`main`).

- **Request hook.** `hook.Hook` (milvus-proto) is consulted by the
proxy's unary interceptor for every RPC on the gRPC and REST surfaces,
and by hand by the service's two streams, `CreateReplicateStream` and
`DumpMessages`, in the same order and with the handler under the context
`Before` returns; a test fails if the service declares a stream that
does not. A compiled-in hook is used in preference to `proxy.soPath`;
configuring both is refused at start-up. The compiled-in hook is
initialized as a plug-in is: `Init` with the `hook.*` configuration
before it is stored, and the `hook.*` reload watcher re-inits it on a
config change. A failing `Init` is fatal at start-up when a form is
installed, whatever `common.panicWhenPluginFail` says; a hook installed
without a form follows the setting like any plug-in.
- **Query hook and cipher.** The QueryNode's tuning hook
(`SetQueryHook`) and the Go half of the cipher pair (`SetCipher`) can be
compiled in the same way. Each is used in preference to its so path
(`queryNode.soPath`, `cipherPlugin.soPathGo`), is refused beside it, and
gets the same `Init` calls and watchers a plug-in gets.
`cipherPlugin.soPathCpp` still declares that encryption is on, so a
compiled-in cipher in a deployment without it stays idle.
`optimizers.QueryHook` is now an alias of `extension.QueryHook`. None of
the three slots implies another or `FormInstalled`.
- **Coordinator engine.** `mixcoord` starts the installed engine once
the replica is ACTIVE (a standby never starts it), on the coordinator
client it uses itself, and stops it on shutdown. A start failure on
activation is fatal. The engine reaches the coordinator only through
`Coordinator`, so it can do exactly what a proxy can do.
- **Query resource-group mark.** `WithQueryResourceGroup` pins a query
to one resource group: the shard client routes it to the leaders whose
replica lives in that group (`ShardLeadersList.resource_groups`, added
on master by #52716) and the proxy attributes its latency to that group.
Nothing in a stock binary sets it.

## Behaviors a form turns on by declaring itself

These apply only when a form is declared (`extension.SetForm()`, which a
distribution does in every role); a stock binary keeps master's
semantics unchanged. `FormInstalled` is independent of `SetHook`: a
distribution may install a request hook without switching any
coordinator behavior on.

- **Delegators on regular query nodes.** Under a form, a replica's shard
delegators live on its regular RW query nodes, which watch the channel
and read the WAL remotely, instead of on the streaming node's embedded
query node; the streaming node is kept for DDL and the write ahead log.
This lets a form load the same collection into several resource groups,
one streaming node total. The trade is stated in the design doc: a
Strong-consistency search no longer takes the WAL-MVCC shortcut, and the
manual-flush prepare step is skipped with a warning. This shape has been
exercised on a real one-streaming-node / multi-RG cluster, not only in
unit tests: `scripts/walkthrough.py` (create, bulk-import ingest, index,
load into a query cluster, cold and warm query, drop) and the e2e suite
(55 tests, including two resource groups loading the same collection)
run against a local k8s lake-form instance, and five concurrent rounds
of two resource groups querying one collection return correct results
with both replicas intact.
- **A load request speaks only for the resource groups it names.**
Loading a collection into `rg_b` does not restate the whole placement
and evict it from `rg_a`; the groups the request does not name keep
their replica count. A request naming no group states the whole
placement. Per-group load progress is readable through
`ShowLoadCollections` with `resource_group` set (appended proto field;
`-1` means the group holds no replica of the collection). A load that
only adds resource groups or replicas to an already loaded collection
keeps the collection serving (an incremental expansion) instead of
resetting it to loading, whether the added replica lands in a group the
request newly names or one that already serves it; the checkers load the
added replicas, as they do on every other path. An expansion replayed
after a crash stays idempotent.
- **Index engine version with no QueryNode session.** With no QueryNode
registered, datacoord answers the index engine version compiled into the
coordinator itself instead of `0`, and bounds an explicit target version
by that same value, on the assumption that a QueryNode started later
runs the same image.

## Removed from the earlier revision

`common.security.requireAPIKey`, the `proxy.internalDomain.*` listeners,
`queryCoord.resourceGroupScopedLoad`,
`queryCoord.reloadSegmentOnIndexDrop`,
`dataCoord.index.allowVectorIndexDropOnLoadedCollection`,
`dataCoord.index.queryNodesScaleToZero`, the `RefusalResponse` hook
extension, the capability table and `Provider`, and the link-time
primary configuration file name. Each is either a deployment constant
(configuration a distribution ships in its own files), something the
request hook already does, or a default this PR now makes unconditional.

## Compatibility

- `ShowCollectionsRequest.resource_group` is the only proto field this
PR appends; an old peer leaves it empty and both sides read empty as "no
scope". `ShardLeadersList.resource_groups`, which the query mark also
reads, was added on master by #52716, not here; an old peer leaves it
empty too and empty is read as "unknown".
- `hook.Hook` is milvus-proto's and unchanged. No new configuration key.
- `ResetForTest` is built only under the `test` tag.

## Testing

- `pkg/extension`: setters, getters, the context mark, and the split
between `SetHook` and `SetForm` (each answers without the other).
- hookutil: compiled-in hook used, refused beside a plug-in, absent by
default, initialized with the hook configuration, re-initialized by the
reload watcher, a failing `Init` fails start-up when a form is installed
even with `common.panicWhenPluginFail` off but follows the setting
without one; compiled-in cipher used when `soPathCpp` is set, refused
beside `soPathGo`, idle without `soPathCpp`, re-initialized on a
`cipherPlugin.*` edit.
- querynode: compiled-in query hook used, refused beside
`queryNode.soPath`, initialized with the `autoIndex.params.*`
configuration, not installed when it cannot initialize.
- proxy: a hook-pinned resource group reaches the search task; per-group
latency series exist only for pinned requests; both streams consult the
hook, and every stream the service declares is one of them.
- mixcoord: the engine starts on activation only, receives the
coordinator client, is stopped once.
- querycoordv2: scoped `ShowLoadCollections` (100 / -1 / unknown group
refused); scoped load placement with and without a form declared; the
expansion predicate's legs, a request that only adds a replica to an
already-loaded group, and a replayed expansion; delegators placed on a
regular query node under a form and on a streaming query node without
one.
- datacoord: zero-session engine version and its upper bound with a form
declared; `0` without.
- `golangci-lint` with the repo config on every touched package.

One deliberate difference from the "stock binary is untouched" claim,
listed rather than hidden: a stock binary with the streaming service on
bypasses the channel checker's version cache for a tick when a replica's
read-write node set is empty at that tick (a streaming node
mid-restart). The cost is one extra `checkReplica` per tick; the case it
exists for, a deployment that starts a replica's nodes on demand, is the
form's.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_015JwRUhmz28AXHyuAJTSs3E

---------

Signed-off-by: Cai Zhang <[email protected]>
Signed-off-by: cai.zhang <[email protected]>
Co-authored-by: Claude Fable 5.1 <[email protected]>
This commit is contained in:
cai.zhang
2026-09-23 05:45:35 +08:00
committed by GitHub
co-authored by Claude Fable 5.1
parent 47bbf3817c
commit 2cbbb72041
77 changed files with 9807 additions and 2408 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/main.go
+1 -78
View File
@@ -17,88 +17,11 @@
package main
import (
"log"
"os"
"os/exec"
"os/signal"
"strings"
"golang.org/x/exp/slices"
"github.com/milvus-io/milvus/cmd/asan"
"github.com/milvus-io/milvus/cmd/milvus"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/internal/util/streamingutil"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
func main() {
// after 2.6.0, we enable streaming service by default.
// TODO: after remove all streamingutil.IsStreamingServiceEnabled(), we can remove this code.
streamingutil.SetStreamingServiceEnabled()
defer asan.LsanDoLeakCheck()
idx := slices.Index(os.Args, "--run-with-subprocess")
// execute command as a subprocess if the command contains "--run-with-subprocess"
if idx > 0 {
args := slices.Delete(os.Args, idx, idx+1)
log.Println("run subprocess with cmd:", args) //nolint:gosec // args are from os.Args, not user input
/* #nosec G204 */
cmd := exec.Command(args[0], args[1:]...) //nolint:gosec // args are from os.Args, not user input
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Start(); err != nil {
// Command not found on PATH, not executable, &c.
log.Fatal(err)
}
// wait for the command to finish
waitCh := make(chan error, 1)
go func() {
waitCh <- cmd.Wait()
close(waitCh)
}()
sc := make(chan os.Signal, 1)
signal.Notify(sc)
// Need a for loop to handle multiple signals
for {
select {
case sig := <-sc:
if err := cmd.Process.Signal(sig); err != nil {
log.Println("error sending signal", sig, err)
}
case err := <-waitCh:
// clean session
paramtable.Init()
params := paramtable.Get()
if len(args) >= 3 {
metaPath := params.EtcdCfg.MetaRootPath.GetValue()
endpoints := params.EtcdCfg.Endpoints.GetValue()
etcdEndpoints := strings.Split(endpoints, ",")
sessionSuffix := sessionutil.GetSessions(cmd.Process.Pid)
defer sessionutil.RemoveServerInfoFile(cmd.Process.Pid)
if err := milvus.CleanSession(metaPath, etcdEndpoints, sessionSuffix); err != nil {
log.Println("clean session failed", err.Error())
}
}
if err != nil {
log.Println("subprocess exit, ", err.Error())
} else {
log.Println("exit code:", cmd.ProcessState.ExitCode())
}
return
}
}
} else {
milvus.RunMilvus(os.Args)
}
milvus.Main(os.Args)
}
cmd/milvus/main_entry.go
+130
View File
@@ -0,0 +1,130 @@
// 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 milvus
import (
"log"
"os"
"os/exec"
"os/signal"
"strings"
"golang.org/x/exp/slices"
"github.com/milvus-io/milvus/cmd/asan"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/internal/util/streamingutil"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// Main is the process entry point, exported so a distribution can build its own
// main package around it instead of duplicating the startup preamble. The body
// is the former cmd/main.go moved unchanged apart from three adaptations: it
// reads args instead of the os.Args global (with the inner slice renamed to
// subArgs), the original if/else became an if block that always returns
// followed by an unconditional tail call, and the subprocess marker is
// stripped into a copy of args rather than in place.
//
// args is the full argument vector, program name first, as os.Args is. Main
// does not modify the caller's slice: the subprocess branch deletes an element,
// which on os.Args was invisible and on an exported API would hand the caller
// back a shortened vector with a cleared tail, so it deletes from a copy.
//
// A distribution installs its compiled-in hook and coordinator engine
// (pkg/extension) before calling Main; nothing runs earlier than the first
// line here except package initialization.
func Main(args []string) {
// after 2.6.0, we enable streaming service by default.
// TODO: after remove all streamingutil.IsStreamingServiceEnabled(), we can remove this code.
streamingutil.SetStreamingServiceEnabled()
defer asan.LsanDoLeakCheck()
// execute command as a subprocess if the command contains "--run-with-subprocess"
if subArgs, runSubprocess := subprocessArgs(args); runSubprocess {
log.Println("run subprocess with cmd:", subArgs) //nolint:gosec // args are from os.Args, not user input
/* #nosec G204 */
cmd := exec.Command(subArgs[0], subArgs[1:]...) //nolint:gosec // args are from os.Args, not user input
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Start(); err != nil {
// Command not found on PATH, not executable, &c.
log.Fatal(err)
}
// wait for the command to finish
waitCh := make(chan error, 1)
go func() {
waitCh <- cmd.Wait()
close(waitCh)
}()
sc := make(chan os.Signal, 1)
signal.Notify(sc)
// Need a for loop to handle multiple signals
for {
select {
case sig := <-sc:
if err := cmd.Process.Signal(sig); err != nil {
log.Println("error sending signal", sig, err)
}
case err := <-waitCh:
// clean session
paramtable.Init()
params := paramtable.Get()
if len(subArgs) >= 3 {
metaPath := params.EtcdCfg.MetaRootPath.GetValue()
endpoints := params.EtcdCfg.Endpoints.GetValue()
etcdEndpoints := strings.Split(endpoints, ",")
sessionSuffix := sessionutil.GetSessions(cmd.Process.Pid)
defer sessionutil.RemoveServerInfoFile(cmd.Process.Pid)
if err := CleanSession(metaPath, etcdEndpoints, sessionSuffix); err != nil {
log.Println("clean session failed", err.Error())
}
}
if err != nil {
log.Println("subprocess exit, ", err.Error())
} else {
log.Println("exit code:", cmd.ProcessState.ExitCode())
}
return
}
}
}
RunMilvus(args)
}
// subprocessArgs returns args with the --run-with-subprocess marker removed,
// and whether the marker was present. It deletes from a copy, never from args,
// which is what lets Main hand a caller its own argument vector back untouched.
func subprocessArgs(args []string) ([]string, bool) {
idx := slices.Index(args, "--run-with-subprocess")
if idx <= 0 {
// idx == 0 is the program name, which is never the marker.
return args, false
}
stripped := slices.Clone(args)
return slices.Delete(stripped, idx, idx+1), true
}
cmd/milvus/main_entry_test.go
+60
View File
@@ -0,0 +1,60 @@
// 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 milvus
import (
"testing"
"github.com/stretchr/testify/assert"
)
// Main must return normally on too few arguments, taking the usage branch
// rather than panicking or exiting. That is the contract a distribution's own
// main relies on.
func TestMainWithoutSubcommandReturns(t *testing.T) {
Main([]string{"milvus"})
}
// Main must leave the caller's argument vector as it found it: a distribution
// may keep using its own os.Args after the call.
func TestMainDoesNotModifyCallerArgs(t *testing.T) {
args := []string{"milvus", "unknown-command"}
want := append([]string(nil), args...)
Main(args)
assert.Equal(t, want, args)
}
// The one in-place change Main used to make - deleting the subprocess marker -
// now lives in a pure helper, so it is tested here without exec'ing a
// subprocess: the marker is stripped from a copy, and the caller's vector is
// returned untouched.
func TestSubprocessArgsStripsTheMarker(t *testing.T) {
args := []string{"milvus", "--run-with-subprocess", "run", "--config", "x"}
want := append([]string(nil), args...)
subArgs, ok := subprocessArgs(args)
assert.True(t, ok)
assert.Equal(t, []string{"milvus", "run", "--config", "x"}, subArgs)
assert.Equal(t, want, args, "the caller's vector must be untouched")
}
func TestSubprocessArgsWithoutTheMarker(t *testing.T) {
args := []string{"milvus", "run"}
subArgs, ok := subprocessArgs(args)
assert.False(t, ok)
assert.Equal(t, args, subArgs, "no marker means no strip")
}
docs/design-docs/design_docs/20260824-per-resource-group-load-and-shard-leader-readiness.md
+1 -1
View File
@@ -167,7 +167,7 @@ Two constraints are enforced by that placement rather than documented for the co
- The proto change is additive and on the response: one repeated field, one generated accessor. `GetShardLeaders` gains no request parameter; its replica filter and its gate are unchanged, and `utils/util.go` differs from its pre-PR state only by filling the new array.
- An old coordinator fills the other three arrays and leaves `resource_groups` empty; the proxy reads that as "unknown" for every leader, and a scoped `ChannelWorkload` against unknown tags gets the retriable refusal above rather than a node from the wrong group. There is no capability handshake on this path and none is needed: the empty array *is* the signal.
- Nothing on the proxy's request path sets `ChannelWorkload.ResourceGroup` yet, so deployed traffic routes exactly as before; the filter and its error contract land here so that the first consumer inherits them rather than re-deriving them.
- The two `utils` surfaces now refuse an unknown group name with `ErrResourceGroupNotFound`; they are in-process only and have no caller before this PR. That is a boundary worth stating: the per-group figure is reachable **in-process only**. `Proxy.GetLoadingProgress` goes through `ShowLoadCollections` and remains a cross-group average with no resource-group parameter, so a client outside the milvus process cannot ask this question yet.
- The two `utils` surfaces refuse an unknown group name with `ErrResourceGroupNotFound`. The per-group figure is reachable on the wire through `ShowLoadCollectionsRequest.resource_group`: set, `inMemory_percentages` is the progress of the replicas in that group, `-1` for a collection with no replica there; empty keeps the collection-wide figure, so `Proxy.GetLoadingProgress` and every existing caller are unchanged.
## 4. Alternatives considered
docs/design-docs/design_docs/20260831-in_tree_extension_mechanism.md
+329
View File
@@ -0,0 +1,329 @@
# MEP: In-Tree Extension Mechanism for Distributions
- **Created:** 2026-08-31
- **Author(s):** @xiaocai2333
- **Status:** In progress
- **Component:** pkg/extension | cmd/milvus | hookutil | Proxy | Coordinator
- **Related Issues:** #52979
- **Implemented by:** #52981
## Summary
A distribution that compiles its own behavior into the milvus binary needs
little from milvus: a way to install a request hook, the QueryNode's tuning
hook and the cipher without a `.so`, and a callback that starts its
control-plane engine when the coordinator becomes active and stops it on
shutdown. `pkg/extension` is those setters, plus the one context mark the
proxy reads. Everything else such a distribution does is either the hook's
own reach, a coordinator RPC, or a configuration item. A stock binary installs
nothing and behaves exactly as before.
## Motivation
The managed-cloud form of milvus had been a fork carrying a handful of
behavior changes in the proxy and the coordinators, re-applied by hand on
every rebase. The first attempt to replace the fork was a table of eight
typed capabilities consulted from eighteen places in the tree. Reviewing it
against the fork's actual needs showed that most entries were not
capabilities: four were constants per deployment (configuration), two
duplicated what the request hook already does on every proxy RPC, one
duplicated coordinator RPCs, and one was an answer with no RPC that belonged
on the wire. What remains is small enough to state on one page.
## Public interfaces
```go
package extension
func SetHook(h hook.Hook) // hookutil prefers it over proxy.soPath
func InstalledHook() hook.Hook
func SetForm() // declares a form installed, in every role
func FormInstalled() bool // reads SetForm's mark; independent of SetHook
type QueryHook interface { ... } // the queryNode.soPath plug-in's method set
func SetQueryHook(h QueryHook) // the QueryNode prefers it over queryNode.soPath
func InstalledQueryHook() QueryHook
func SetCipher(c hook.Cipher) // hookutil prefers it over cipherPlugin.soPathGo
func InstalledCipher() hook.Cipher
type Coordinator interface { // the coordinator as its own clients see it
rootcoordpb.RootCoordClient
querypb.QueryCoordClient
datapb.DataCoordClient
}
type CoordinatorEngine interface {
Start(ctx context.Context, coord Coordinator) error
Stop() error
}
func SetCoordinatorEngine(e CoordinatorEngine)
func InstalledCoordinatorEngine() CoordinatorEngine
func WithQueryResourceGroup(ctx, rg) context.Context // set by a hook's Before
func QueryResourceGroupFromContext(ctx) string
```
A distribution calls the setters it needs, then `cmd/milvus.Main(os.Args)`.
### The request hook
`hook.Hook` (milvus-proto) is consulted by the proxy's unary interceptor for
every RPC on both the gRPC and the REST surface, and by the service's two
streams, `CreateReplicateStream` and `DumpMessages`, which consult it by hand -
the same way, `Mock`, `Before` and `After` in order, with the stream's request
(`DumpMessagesRequest`, or an empty `ReplicateRequest` for the replicate
stream, which carries none of its own) and the stream run under the context
`Before` returned. A test fails as soon as the service declares a stream that
does not consult it. A compiled-in hook that also
implements `hook.Extension` is stored as the extension, so `Report` and
`ReportAction` reach it as they reach a plug-in's `MilvusExtension`.
`Mock` answers without forwarding, `Before` may rewrite the request in place,
block, and return the context the handler runs under, `After` sees the result.
`VerifyAPIKey` answers the API key. Whether the external listener also accepts
a username and password is a policy the hook itself enforces from `Before`,
which sees the same metadata the authentication interceptor does; a refusal
from `Before` always reaches the client as the interceptor's own
`InvalidArgument`, which is in pymilvus's non-retried set.
One addition to hookutil: a compiled-in hook (`SetHook`) is used in preference
to `proxy.soPath`, and a deployment that configures both is refused at
start-up. A compiled-in hook is otherwise treated exactly as a plug-in is: it
gets the same `Init` call with the `hook.*` configuration before it is
installed, a failure to initialize keeps the proxy from starting, and it is
registered with the same watcher, so editing a `hook.*` key re-initializes it
with the new configuration without a restart. It differs in one respect:
`common.panicWhenPluginFail`, which lets an operator run on without a plug-in
that failed to load, does not reach it. A compiled-in hook that cannot
initialize, or is configured beside a plug-in, stops the proxy whatever the
setting says, because the distribution that compiled it in has switched the
coordinators' behaviors on too, and a proxy serving through the default hook
beside them would run half of that distribution.
### The query hook
The QueryNode's search-parameter tuning hook - the `QueryNodePlugin` symbol a
`queryNode.soPath` plug-in exports - can be compiled in the same way
(`SetQueryHook`). A compiled-in query hook is used in preference to
`queryNode.soPath`, and a deployment that configures both is refused at
start-up: both would tune every search, and only one can. The refusal is
reported the way a plug-in's load failure is: the QueryNode treats it as
fatal when `autoIndex.enable` is on, and with tuning off it starts without
either hook, as it always has. It is otherwise treated exactly as the plug-in
is: it gets the same two `Init` calls with the
`autoIndex.params.search` and `autoIndex.params.tuning` configuration before it
is installed, and the same watchers re-initialize it when those keys change.
`autoIndex.enable` still decides whether tuning happens at all, whichever way
the hook got there. `optimizers.QueryHook` is now an alias of
`extension.QueryHook`, so nothing in the tree changes.
### The cipher
The Go half of the cipher plug-in pair - the `CipherPlugin` symbol a
`cipherPlugin.soPathGo` plug-in exports - can be compiled in the same way
(`SetCipher`). A compiled-in cipher is used in preference to
`cipherPlugin.soPathGo`, and a deployment that configures both is refused at
start-up: both would answer for the encryption keys, and only one can. Only
that half changes hands: the C++ half is still loaded by the core from
`cipherPlugin.soPathCpp`, and that path stays the deployment's declaration
that encryption is on. A compiled-in cipher in a deployment without
`soPathCpp` is left idle and `IsClusterEncryptionEnabled` stays false,
exactly as with no plug-in. When it is used it is treated exactly as the
plug-in is: it gets the same `Init` with the `cipherPlugin.*` configuration
before it is installed, and the same reload callbacks re-initialize it when
those keys change.
### The coordinator engine
`mixcoord` starts the installed engine once the replica is ACTIVE (a standby
never starts it), on the coordinator client it uses itself, and stops it on
shutdown. A start failure on activation is fatal: a coordinator serving without
its engine would accept work nothing accounts for. Start is called at most
once and Stop at most once, only after Start - a standby that shuts down never
has its engine stopped - and a shutdown does not wait for a slow Start: Stop
may overlap it, and must make it return. The engine reaches the
coordinator through nothing but `Coordinator`, so what it can do is exactly
what a proxy can do - including seeding its own accounts through
`CreateCredential` / `OperateUserRole`, and reading per-resource-group load
progress through `ShowLoadCollections` with `resource_group` set.
### Context marks
`WithQueryResourceGroup` pins a query to one resource group. The shard client
routes it to the leaders whose replica lives in that group
(`ShardLeadersList.resource_groups`) and the proxy attributes its latency to
that group. Nothing in a stock binary sets it.
### Form-gated behaviors
A distribution declares the deployment shape its coordinator behaviors are
written for with `extension.SetForm()`; `extension.FormInstalled()` reads that
declaration, independent of `SetHook`. Three behaviors in the coordinators are
switched on by `FormInstalled()` alone. They exist for the
deployment shape a distribution runs - one streaming node kept for DDL and the
write ahead log while queries are served from resource groups of regular query
nodes, one collection loaded into several of those groups independently,
every role rolled from one image - and a stock binary, which
answers false, keeps master's behavior exactly. Because the query coordinator
and the data coordinator read the mark too, a distribution must make the
declaration in every role it runs, not only in the proxy.
- **Delegator placement** (`streamingutil.UseStreamingQueryNodeAsDelegator`).
With the streaming service on, the query coordinator places shard delegators
on streaming query nodes only, gives every replica a streaming query node of
its own, and refuses a collection more replicas than there are streaming
nodes. With a form installed, delegators go onto the replica's regular RW
query nodes instead, which watch the channel and read the WAL remotely, and
the streaming node takes no part in serving: a form runs one, and loads the
same collection into several resource groups. The replica bound in
`utils.AssignReplica`, the channel and leader checkers, the assign policies'
node filter, the channel balance helpers and the manual channel transfer ask
this one question; each already carried both placements. Stock: exactly the
streaming gate's answer, as before.
A delegator that does not share a process with the WAL trades two QueryNode
shortcuts for that separation. On a stock replica, where they do share one:
`delegator.GetLatestMVCCTimestampIfLocal` reads the WAL directly and lets a
Strong-consistency search be answered without waiting for tsafe to catch up;
`PrepareReleaseManualFlushIfLocal` runs the manual-flush prepare step in
place. A form's delegator, reading the WAL remotely, gets an error from the
first (`it does not run beside the WAL`), so every Strong search waits for
the channel's tsafe instead; and the second hits
`isReleaseManualFlushPrepareUnavailable`, which logs a warning and
unsubscribes without the prepare step. Neither is a bug - the WAL is being
read remotely by design - but both are the difference between a delegator on
a streaming node and one on a regular query node, and both paths run on the
2.5-compat branches (`utils.GetChannelRWAndRONodesFor260`,
`utils.filterNodeLessThan260` and the checkers' non-streaming-query-node
branches) that `cmd/milvus/main_entry.go` already marks for removal.
- **Resource-group-scoped load placement** (`completePlacementForOutOfScopeResourceGroups`).
With a form installed, a load request naming resource groups only ever
changes the placement in those groups: the groups it does not name keep the
replicas they hold; the querycoord load job recognizes a request that only
adds groups to a Loaded collection as a pure expansion (the fast path in
`job_load.go`) and keeps the collection serving instead of resetting it to
Loading; and a scoped `LoadPartitions` on a Loaded collection is judged on
the groups it names - a named group that holds no replica yet is an
addition, not a replica-number change, while changing the count of a named
group that already holds replicas is refused as before. Stock, all three
are master's: every request states the whole placement, so a second
`LoadCollection(resource_groups=[rg_1])` on a collection loaded in `rg_0`
moves the replica to `rg_1` and the request's `replica_number` is the
total; every load of an already-loaded collection - a two-group
`LoadCollection` included - writes it back to Loading/0 and the
collection-wide observer walks it up again; and `LoadPartitions` refuses any
change of the total replica count.
The scope comes from the request as the caller wrote it. A request naming
no group at all states the whole placement even for a form, which is what a
plain `load_collection` or `load_partitions` has always meant, so nothing is
carried over for it; a request that names `__default_resource_group`
explicitly is a scoped request like any other. A cluster-level load
override states the whole placement too, whatever the request named.
The expansion path is replay-safe: the spawn is persisted before the rest
of it runs, and a message replayed after a failure finds the same replica
set already stored. An identical replica set on a Loaded collection is
therefore read as a replayed expansion (a request that changes nothing is
never broadcast), and every step is idempotent.
An expansion registers no load task of its own. The collection stays
Loaded, so there is no status to write back and no aggregate to complete;
how far a group has loaded and whether it can serve are answered from the
live target and distribution (`utils.LoadPercentageByResourceGroup` and
`utils.ShardLeaderReadinessByResourceGroup`), and the checkers load the
added replicas on their own schedule. This is what upstream's
`UpdateLoadConfig` already does when it adds a replica to a loaded
collection: spawn, write the count, pull the target, and leave the loading
to the checkers. A group whose replicas never load is the caller's to
release - the deployment that asked for the group is the one that knows it
no longer wants it - exactly as an added replica that never loads is
upstream.
- **Index engine version with no QueryNode session** (datacoord's version
manager). With a form installed, datacoord answers its own compiled-in
index engine version (knowhere's for vectors, the constant for scalars)
rather than zero, which knowhere would otherwise read as "only DISKANN
loads off disk" and misroute other disk indexes onto the in-memory path.
This assumes a QueryNode started later runs the same image as the
coordinator - true for a distribution that rolls every role from one image
- and the same assumption sets the upper bound with no session: an
operator's `dataCoord.targetVecIndexVersion` (or the scalar one) above what
this image can load is clamped down to it rather than written into index
builds unchecked, with a rate-limited warning. A QueryNode that does
register replaces both figures with the cluster-wide ones. Stock: 0 and no
upper bound, as on master - the current version is the minimum over every
QueryNode's, a datacoord that comes up first in a rolling upgrade must not
build indexes an older QueryNode cannot load, and an override is written
through.
## Configuration this mechanism relies on
Both items are code defaults, not `milvus.yaml` entries: a distribution sets
them through `user.yaml` or the environment.
| Key | Default | What a distribution sets it for |
|---|---|---|
| `dataCoord.externalCollection.refreshWaitForIndex` | false | hold a refresh until its segments are indexed |
| `builtinRoles.*` | - | the roles the distribution's accounts bind to |
## Compatibility
- `ShowCollectionsRequest.resource_group` is the only proto field this PR
appends; an old peer leaves it empty and the reader treats empty as "no
scope". `ShardLeadersList.resource_groups`, which the query mark also reads,
was added on master by #52716, not here; an old peer leaves it empty too and
empty is read as "unknown".
- Every configuration item defaults to the stock behavior.
- `hook.Hook` is milvus-proto's and unchanged.
- `hook.Cipher` is milvus-proto's and unchanged; `cipherPlugin.soPathCpp`
keeps its meaning.
- `optimizers.QueryHook` is an alias of `extension.QueryHook` with the same
method set; a `QueryNodePlugin` built against either loads unchanged.
- The default configuration file is `milvus.yaml`; a distribution that ships
its configuration under another name places it as `user.yaml` in
`MILVUSCONF`, which the existing file list already reads last.
## Test plan
- `pkg/extension`: the setters and getters, and the context mark.
- hookutil: a compiled-in hook is used, refused beside a plug-in, absent by
default, initialized with the `hook.*` configuration, re-initialized when
that configuration changes, and fatal when it cannot initialize even with
`common.panicWhenPluginFail` off.
- hookutil: a compiled-in cipher is used when `cipherPlugin.soPathCpp` is
set, refused beside `cipherPlugin.soPathGo`, idle without `soPathCpp`,
absent by default, not installed when it cannot initialize, and
re-initialized on a `cipherPlugin.*` edit.
- querynode: a compiled-in query hook is used, refused beside a plug-in,
absent by default, initialized with the `autoIndex.params.*` configuration,
and not installed when it cannot initialize.
- proxy: a hook-pinned resource group reaches the search task; per-resource-
group latency series exist only for pinned requests; both streams consult
the hook, and every stream the service declares is one of them.
- mixcoord: the engine starts on activation only, receives the coordinator
client, and is stopped once.
- querycoord / datacoord: each form-gated behavior with and without a form
installed - the stock cases assert master's answers (a second replica
refused for want of a streaming node, a delegator placed on the streaming
query node, a scoped load that moves the replica, version 0
with no session) - plus a request naming no resource group; the resource
group scope of `completePlacementForOutOfScopeResourceGroups` (a named group
is an addition, an unnamed one is carried through); the expansion predicate
in `job_load.go`, its legs, a request that only adds a replica to an
already-loaded group, and a replayed expansion; a scoped
`ShowLoadCollections` answered from the group's own figures (100 / -1 /
unknown group refused); a delegator placed on a regular query node under a
form; and an index engine version override clamped with no QueryNode
registered.
## Rejected alternatives
- **A capability table with a `Provider`, `Requires()` and typed-nil checks.**
Shipped first as #52981. Eight capabilities, nine cross-boundary interfaces,
eighteen read points; most entries were constants or duplicates of existing
mechanisms. Replaced by this document.
- **A link-time primary configuration file name.** An entrypoint that renames
the file solves it with no milvus change.
- **Per-callback interfaces for the proxy request path.** Every proxy RPC
already passes the hook's `Mock`, `Before` and `After`, and `Before` returns
the handler's context.
internal/coordinator/mix_coord.go
+38 -6
View File
@@ -63,15 +63,21 @@ type mixCoordImpl struct {
tikvCli *txnkv.Client
address string
proxyCreator proxyutil.ProxyCreator
proxyWatcher *proxyutil.ProxyWatcher
proxyClientManager proxyutil.ProxyClientManagerInterface
proxyCreator proxyutil.ProxyCreator
proxyWatcher *proxyutil.ProxyWatcher
metricsCacheManager *metricsinfo.MetricsCacheManager
stateCode atomic.Int32
initOnce sync.Once
startOnce sync.Once
session *sessionutil.Session
// onActive holds callbacks to run once this replica becomes ACTIVE -
// after initInternal has built the sub-coordinators and the state moved
// to Healthy. A standby replica never fires them; see OnActive.
onActiveMu sync.Mutex
onActive []func()
activated bool
initOnce sync.Once
startOnce sync.Once
session *sessionutil.Session
factory dependency.Factory
@@ -293,6 +299,32 @@ func (s *mixCoordImpl) startAndUpdateHealthy() {
s.UpdateStateCode(commonpb.StateCode_Healthy)
s.startPosixCleanupTask()
RegisterMgrRoute(s)
s.onActiveMu.Lock()
s.activated = true
fns := s.onActive
s.onActive = nil
s.onActiveMu.Unlock()
for _, fn := range fns {
fn()
}
}
// OnActive runs fn once this replica is ACTIVE: sub-coordinators initialized,
// state Healthy. A callback registered after activation runs immediately; one
// registered on a replica that stays standby never runs, which is the point -
// work hung off this hook (a coordinator engine doing resource-group
// accounting, say) must run on exactly one replica, and before activation the
// sub-coordinators it would read are not initialized at all.
func (s *mixCoordImpl) OnActive(fn func()) {
s.onActiveMu.Lock()
if s.activated {
s.onActiveMu.Unlock()
fn()
return
}
s.onActive = append(s.onActive, fn)
s.onActiveMu.Unlock()
}
func (s *mixCoordImpl) IsServerActive(serverID int64) bool {
internal/datacoord/index_engine_version_manager.go
+77 -6
View File
@@ -10,8 +10,10 @@ import (
"github.com/samber/lo"
"golang.org/x/time/rate"
"github.com/milvus-io/milvus/internal/util/segcore"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v3/util/lock"
@@ -202,7 +204,7 @@ func (m *versionManagerImpl) GetCurrentIndexEngineVersion() int32 {
func (m *versionManagerImpl) getCurrentVersion() int32 {
if len(m.versions) == 0 {
return 0
return noSessionVersion(func() int32 { return segcore.GetIndexEngineInfo().CurrentIndexVersion })
}
current := int32(math.MaxInt32)
@@ -223,7 +225,11 @@ func (m *versionManagerImpl) GetMinimalIndexEngineVersion() int32 {
func (m *versionManagerImpl) getMinimalVersion() int32 {
if len(m.versions) == 0 {
return 0
// The same assumption noSessionVersion makes for the current and
// maximum versions covers the lower bound too: a QueryNode started
// later runs this image, so an override below what this image's
// segcore can load is clamped up rather than built.
return noSessionVersion(func() int32 { return segcore.GetIndexEngineInfo().MinIndexVersion })
}
minimal := int32(0)
@@ -244,7 +250,7 @@ func (m *versionManagerImpl) GetCurrentScalarIndexEngineVersion() int32 {
func (m *versionManagerImpl) getCurrentScalarVersion() int32 {
if len(m.scalarIndexVersions) == 0 {
return 0
return noSessionVersion(func() int32 { return common.CurrentScalarIndexEngineVersion })
}
current := int32(math.MaxInt32)
@@ -285,7 +291,16 @@ func (m *versionManagerImpl) GetMaximumIndexEngineVersion() int32 {
}
func (m *versionManagerImpl) getMaximumVersion() int32 {
return getMaximumVersionFrom(m.versions)
// The highest version this image can LOAD, which is what an upper bound
// asks for and is not the same as the version it builds at. A registered
// QueryNode is read as max(Current, Maximum) below, so answering with the
// current version alone would make the same image bound an operator's
// target lower with no session than with one, and clamp index builds
// during a restart to a version this very image can read past.
return getMaximumVersionFrom(m.versions, func() int32 {
info := segcore.GetIndexEngineInfo()
return max(info.CurrentIndexVersion, info.MaxIndexVersion)
})
}
func (m *versionManagerImpl) GetMaximumScalarIndexEngineVersion() int32 {
@@ -296,11 +311,67 @@ func (m *versionManagerImpl) GetMaximumScalarIndexEngineVersion() int32 {
}
func (m *versionManagerImpl) getMaximumScalarVersion() int32 {
return getMaximumVersionFrom(m.scalarIndexVersions)
// As getMaximumVersion: the bound with no session is what this image can
// LOAD, which is how a registered node is read too.
return getMaximumVersionFrom(m.scalarIndexVersions, func() int32 {
return max(common.CurrentScalarIndexEngineVersion, common.MaximumScalarIndexEngineVersion)
})
}
func getMaximumVersionFrom(versions map[int64]sessionutil.IndexEngineVersion) int32 {
// noSessionVersion is the current index engine version with no QueryNode
// session registered, given the version compiled into this coordinator.
//
// A stock binary answers 0, as it always has. The current version is the MIN
// over every QueryNode's, the highest version all of them can load, and with
// none registered there is nothing to take a minimum of: datacoord coming up
// before any QueryNode is ordinary during a full restart or a rolling upgrade,
// and an index or compaction output built in that window with this
// coordinator's own version could not be loaded by an older QueryNode that
// registers afterwards. Zero is the answer that assumes nothing.
//
// With a form installed (extension.FormInstalled) the answer is the version
// compiled into the coordinator. Such a deployment rolls every role from one
// image, so a QueryNode started later runs this same version, and it loads
// collections on demand, where version 0 costs it something: knowhere reads 0
// as "only DISKANN loads off disk" and misroutes other disk indexes onto the
// in-memory path. The same assumption sets the ceiling with no session
// (getMaximumVersionFrom), so an operator override above it is clamped rather
// than written into index builds nothing can load. If the assumption is wrong
// - a QueryNode on an older image joins - its session replaces both figures
// the moment it registers, and the answers become cluster-wide again.
//
// compiledIn is called only on that path; see getMaximumVersionFrom.
func noSessionVersion(compiledIn func() int32) int32 {
if extension.FormInstalled() {
return compiledIn()
}
return 0
}
// getMaximumVersionFrom returns the highest index version every registered
// QueryNode can load. compiledIn is the answer when none is registered and a
// form is installed: the version this coordinator's own image can load. A
// stock binary keeps the unbounded MaxInt32 it always had.
//
// The bound is the same assumption noSessionVersion makes with no session -- a
// QueryNode started later runs this image -- and it has a consequence worth
// stating, because it is the whole point of returning it here. This is the
// ceiling clampVersion applies to dataCoord.targetVecIndexVersion and
// dataCoord.targetScalarIndexVersion, so an operator override above what this
// image can load is clamped down to it (with the rate-limited warning) instead
// of being written into index builds unchecked. An unbounded MaxInt32 makes
// the clamp a no-op precisely when there is no QueryNode to disprove the
// override, which is what a stock binary does and keeps doing.
//
// compiledIn is called only when it is the answer. Asking this image for its
// version is three cgo calls, made under the manager's lock and for every
// segment index a compaction checks (ResolveVecIndexVersion), while the answer
// almost always comes from the registered sessions instead.
func getMaximumVersionFrom(versions map[int64]sessionutil.IndexEngineVersion, compiledIn func() int32) int32 {
if len(versions) == 0 {
if extension.FormInstalled() {
return compiledIn()
}
return math.MaxInt32
}
internal/datacoord/index_engine_version_manager_test.go
+252
View File
@@ -2,13 +2,17 @@ package datacoord
import (
"math"
"strconv"
"testing"
"github.com/blang/semver/v4"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/internal/util/segcore"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/v3/common"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
@@ -833,3 +837,251 @@ func Test_IndexEngineVersionManager_SessionVersionCleanupOnStartup(t *testing.T)
_, exists := vm.sessionVersion[2]
assert.False(t, exists, "offline node should be removed from sessionVersion map")
}
// installForm turns this test's binary into a form, and turns it back into a
// stock binary when the test ends.
func installForm(t *testing.T) {
t.Helper()
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
ext.SetForm()
}
// A stock binary with no QueryNode session answers exactly what master does:
// version 0, no upper bound, and an operator override written through. The
// current version is the MIN over every QueryNode's; with none registered
// there is nothing to take it over, and a coordinator that came up first in a
// rolling upgrade must not build indexes an older QueryNode cannot load.
func TestAStockBinaryAnswersZeroWithNoSession(t *testing.T) {
paramtable.Init()
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
m := newIndexEngineVersionManager()
assert.Zero(t, m.GetCurrentIndexEngineVersion())
assert.Zero(t, m.GetCurrentScalarIndexEngineVersion())
assert.Equal(t, int32(math.MaxInt32), m.GetMaximumIndexEngineVersion())
assert.Equal(t, int32(math.MaxInt32), m.GetMaximumScalarIndexEngineVersion())
// With no upper bound, an override above what this image can load is
// written through - master's behavior, kept on a stock binary.
p := paramtable.Get()
compiledInVec := segcore.GetIndexEngineInfo().CurrentIndexVersion
p.Save(Params.DataCoordCfg.TargetVecIndexVersion.Key, strconv.Itoa(int(compiledInVec)+5))
defer p.Reset(Params.DataCoordCfg.TargetVecIndexVersion.Key)
p.Save(Params.DataCoordCfg.TargetScalarIndexVersion.Key,
strconv.Itoa(int(common.CurrentScalarIndexEngineVersion)+5))
defer p.Reset(Params.DataCoordCfg.TargetScalarIndexVersion.Key)
assert.Equal(t, compiledInVec+5, m.ResolveVecIndexVersion(),
"with no session a stock binary writes an override through, as it always has")
assert.Equal(t, common.CurrentScalarIndexEngineVersion+5, m.ResolveScalarIndexVersion())
}
// With a form installed, an empty session set is the resting state and the
// versions come from this process's own engine - the same values the absent
// query nodes, running this same image, would have reported. Version zero here
// is what misroutes disk indexes in knowhere, so the assertion pins non-zero
// as well as source equality. The store-path gate keeps its native reading
// regardless: no session means no evidence any reader is on an older layout,
// but the coordinator's own engine version is not that evidence either.
func TestEmptySessionSetComesFromThisBinary(t *testing.T) {
installForm(t)
m := newIndexEngineVersionManager()
vec := m.GetCurrentIndexEngineVersion()
assert.Equal(t, segcore.GetIndexEngineInfo().CurrentIndexVersion, vec,
"the fallback must be this binary's own knowhere version")
assert.NotZero(t, vec, "version zero is the misrouting answer the fallback exists to avoid")
assert.Equal(t, common.CurrentScalarIndexEngineVersion, m.GetCurrentScalarIndexEngineVersion())
// The lower bound comes from this binary too. The assumption that a
// QueryNode started later runs this image covers the floor the same way
// it covers the ceiling: an override below what this image's segcore can
// load would build an index this same image cannot read.
assert.Equal(t, segcore.GetIndexEngineInfo().MinIndexVersion, m.GetMinimalIndexEngineVersion())
paramtable.Get().Save(Params.DataCoordCfg.IndexStorePathVersion.Key, "1")
defer paramtable.Get().Reset(Params.DataCoordCfg.IndexStorePathVersion.Key)
assert.Equal(t, indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_BUILD_ROOTED,
m.GetClusterMinIndexStorePathVersion())
}
// A query node that IS reporting always wins over the no-session fallback:
// the fallback is about an empty set, never about overriding a session that
// exists.
func TestAReportingQueryNodeWinsOverTheFallback(t *testing.T) {
installForm(t)
m := newIndexEngineVersionManager()
m.Startup(map[string]*sessionutil.Session{
"qn1": {SessionRaw: sessionutil.SessionRaw{
ServerID: 1,
IndexEngineVersion: sessionutil.IndexEngineVersion{MinimalIndexVersion: 0, CurrentIndexVersion: 1},
ScalarIndexEngineVersion: sessionutil.IndexEngineVersion{MinimalIndexVersion: 0, CurrentIndexVersion: 1},
}},
})
assert.Equal(t, int32(1), m.GetCurrentIndexEngineVersion())
assert.Equal(t, int32(1), m.GetCurrentScalarIndexEngineVersion())
}
// TestEmptySessionSetBoundsAnOverrideByThisBinary is the other half of the
// form's no-session fallback. The clamp at the end of the resolve functions is
// what stops dataCoord.targetVecIndexVersion from asking for an index nothing
// in the cluster can load; with no QueryNode session the upper bound is
// MaxInt32 on a stock binary, so the clamp does nothing and the override is
// written through whatever it says. The bound with no session is the same
// assumption the form's current version makes: a QueryNode started later runs
// this image.
func TestEmptySessionSetBoundsAnOverrideByThisBinary(t *testing.T) {
paramtable.Init()
installForm(t)
p := paramtable.Get()
m := newIndexEngineVersionManager()
info := segcore.GetIndexEngineInfo()
loadableVec := max(info.CurrentIndexVersion, info.MaxIndexVersion)
assert.Equal(t, loadableVec, m.GetMaximumIndexEngineVersion(),
"with no session the ceiling is the highest version this image can load")
assert.Equal(t, common.CurrentScalarIndexEngineVersion, m.GetMaximumScalarIndexEngineVersion())
p.Save(Params.DataCoordCfg.TargetVecIndexVersion.Key, strconv.Itoa(int(loadableVec)+5))
defer p.Reset(Params.DataCoordCfg.TargetVecIndexVersion.Key)
p.Save(Params.DataCoordCfg.TargetScalarIndexVersion.Key,
strconv.Itoa(int(common.CurrentScalarIndexEngineVersion)+5))
defer p.Reset(Params.DataCoordCfg.TargetScalarIndexVersion.Key)
assert.Equal(t, loadableVec, m.ResolveVecIndexVersion(),
"an override above what this image can load must be clamped to it")
assert.Equal(t, common.CurrentScalarIndexEngineVersion, m.ResolveScalarIndexVersion())
}
// An override BELOW the bound is untouched: the clamp is a ceiling, not a
// rewrite, and asking for an older index version is a legitimate thing to do.
func TestEmptySessionSetLeavesAnOverrideBelowThisBinaryAlone(t *testing.T) {
paramtable.Init()
installForm(t)
p := paramtable.Get()
m := newIndexEngineVersionManager()
p.Save(Params.DataCoordCfg.ForceRebuildSegmentIndex.Key, "true")
defer p.Reset(Params.DataCoordCfg.ForceRebuildSegmentIndex.Key)
p.Save(Params.DataCoordCfg.TargetVecIndexVersion.Key, "1")
defer p.Reset(Params.DataCoordCfg.TargetVecIndexVersion.Key)
p.Save(Params.DataCoordCfg.ForceRebuildScalarSegmentIndex.Key, "true")
defer p.Reset(Params.DataCoordCfg.ForceRebuildScalarSegmentIndex.Key)
p.Save(Params.DataCoordCfg.TargetScalarIndexVersion.Key, "0")
defer p.Reset(Params.DataCoordCfg.TargetScalarIndexVersion.Key)
assert.EqualValues(t, 1, m.ResolveVecIndexVersion())
assert.EqualValues(t, 0, m.ResolveScalarIndexVersion())
}
// The bound with no session must be the same figure a registered QueryNode
// running this image would give: every session is read as max(Current,
// Maximum), so answering with the current version alone would bound an
// operator's target lower during a restart than a moment later, and clamp
// index builds to a version this very image can read past. The fallback is
// also still only about an EMPTY set: one session replaces it.
func TestTheNoSessionBoundIsWhatThisImageCanLoad(t *testing.T) {
paramtable.Init()
installForm(t)
m := newIndexEngineVersionManager()
info := segcore.GetIndexEngineInfo()
loadable := max(info.CurrentIndexVersion, info.MaxIndexVersion)
assert.Equal(t, loadable, m.GetMaximumIndexEngineVersion())
assert.GreaterOrEqual(t, loadable, info.CurrentIndexVersion,
"what an image can load is never below what it builds at")
m.Startup(map[string]*sessionutil.Session{
"qn1": {SessionRaw: sessionutil.SessionRaw{
ServerID: 1,
IndexEngineVersion: sessionutil.IndexEngineVersion{
CurrentIndexVersion: 3, MaximumIndexVersion: 4,
},
}},
})
assert.EqualValues(t, 4, m.GetMaximumIndexEngineVersion(),
"a session that exists is read the same way, and replaces the fallback")
}
// The scalar side of TestTheNoSessionBoundIsWhatThisImageCanLoad: with no
// query node session the bound is what this image can load, read the same
// way a registered node's scalar triple is. Both constants are equal today,
// so this pins the shape rather than a difference.
func TestTheNoSessionScalarBoundIsWhatThisImageCanLoad(t *testing.T) {
paramtable.Init()
installForm(t)
m := newIndexEngineVersionManager()
loadable := max(common.CurrentScalarIndexEngineVersion, common.MaximumScalarIndexEngineVersion)
assert.Equal(t, loadable, m.GetMaximumScalarIndexEngineVersion())
assert.GreaterOrEqual(t, loadable, common.CurrentScalarIndexEngineVersion,
"what an image can load is never below what it builds at")
m.Startup(map[string]*sessionutil.Session{
"qn1": {SessionRaw: sessionutil.SessionRaw{
ServerID: 1,
ScalarIndexEngineVersion: sessionutil.IndexEngineVersion{
CurrentIndexVersion: 3, MaximumIndexVersion: 4,
},
}},
})
assert.EqualValues(t, 4, m.GetMaximumScalarIndexEngineVersion(),
"a session that exists is read the same way, and replaces the fallback")
}
// The compiled-in version is read only when it is the answer: a form installed
// and no QueryNode session. Reading it is three cgo calls under the manager's
// lock, and GetMaximumIndexEngineVersion is asked for every segment index a
// compaction checks, so it must not be read when the sessions answer, or on a
// stock binary, which never uses it.
func TestTheCompiledInVersionIsReadOnlyWhenItIsTheAnswer(t *testing.T) {
paramtable.Init()
reads := 0
var origin func() segcore.IndexEngineInfo
counting := mockey.Mock(segcore.GetIndexEngineInfo).To(func() segcore.IndexEngineInfo {
reads++
return origin()
}).Origin(&origin).Build()
defer counting.UnPatch()
withSession := func(m IndexEngineVersionManager) {
m.Startup(map[string]*sessionutil.Session{
"qn1": {SessionRaw: sessionutil.SessionRaw{
ServerID: 1,
IndexEngineVersion: sessionutil.IndexEngineVersion{CurrentIndexVersion: 3, MaximumIndexVersion: 4},
}},
})
}
t.Run("stock binary, no session", func(t *testing.T) {
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
m := newIndexEngineVersionManager()
reads = 0
m.GetMaximumIndexEngineVersion()
m.GetCurrentIndexEngineVersion()
assert.Zero(t, reads, "a stock binary answers without this image's version")
})
t.Run("form installed, a session registered", func(t *testing.T) {
installForm(t)
m := newIndexEngineVersionManager()
withSession(m)
reads = 0
assert.EqualValues(t, 4, m.GetMaximumIndexEngineVersion())
assert.EqualValues(t, 3, m.GetCurrentIndexEngineVersion())
assert.Zero(t, reads, "the sessions answer, so this image's version is not read")
})
t.Run("form installed, no session", func(t *testing.T) {
installForm(t)
m := newIndexEngineVersionManager()
reads = 0
m.GetMaximumIndexEngineVersion()
assert.Equal(t, 1, reads, "with nothing registered, this image's version is the answer")
m.GetCurrentIndexEngineVersion()
assert.Equal(t, 2, reads)
})
}
internal/datacoord/index_service_drop_loaded_test.go
+109
View File
@@ -0,0 +1,109 @@
// 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"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"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/schemapb"
"github.com/milvus-io/milvus/internal/datacoord/broker"
catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// newLoadedVectorIndexServer builds a datacoord whose one collection is loaded
// and carries one vector index, and the request that drops that index.
func newLoadedVectorIndexServer(t *testing.T) (*Server, *indexpb.DropIndexRequest) {
const (
collID = UniqueID(1)
fieldID = UniqueID(10)
indexID = UniqueID(100)
indexName = "vector_idx"
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().AlterIndexes(mock.Anything, mock.Anything).Return(nil).Maybe()
b := broker.NewMockBroker(t)
b.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
DbName: "test_db",
CollectionName: "test_collection",
Schema: &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: fieldID, Name: "vec", DataType: schemapb.DataType_FloatVector},
},
},
}, nil)
s := &Server{
meta: &meta{
catalog: catalog,
indexMeta: &indexMeta{
catalog: catalog,
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "128"}},
IndexParams: []*commonpb.KeyValuePair{{Key: common.IndexTypeKey, Value: "IVF_FLAT"}},
},
},
},
segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
},
segments: NewSegmentsInfo(),
},
broker: b,
allocator: newMockAllocator(t),
notifyIndexChan: make(chan UniqueID, 1),
}
mixCoord := mocks.NewMixCoord(t)
mixCoord.EXPECT().ShowLoadCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{
Status: merr.Success(),
CollectionIDs: []int64{collID},
}, nil)
s.mixCoord = mixCoord
RegisterDDLCallbacks(s)
s.stateCode.Store(commonpb.StateCode_Healthy)
return s, &indexpb.DropIndexRequest{CollectionID: collID, IndexName: indexName}
}
func TestDropVectorIndexOnLoadedIsRefused(t *testing.T) {
initStreamingSystem(t)
s, req := newLoadedVectorIndexServer(t)
status, err := s.DropIndex(context.Background(), req)
require.NoError(t, err)
assert.False(t, merr.Ok(status))
assert.Contains(t, status.GetReason(), "vector index cannot be dropped on loaded collection")
assert.NotEmpty(t, s.meta.indexMeta.GetIndexesForCollection(req.GetCollectionID(), req.GetIndexName()),
"a refused drop must leave the index in place")
}
internal/distributed/mixcoord/extension_seam.go
+192
View File
@@ -0,0 +1,192 @@
// 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 grpcmixcoord
import (
"context"
"sync"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
// The coordinator engine is control-plane machinery a deployment hosts in the
// coordinator process. With none installed the two functions below do nothing.
func coordinatorEngine() extension.CoordinatorEngine {
return extension.InstalledCoordinatorEngine()
}
// activeNotifier is what the coordinator implements to run work once this
// replica is ACTIVE; a standby never fires it.
type activeNotifier interface {
OnActive(fn func())
}
// engineLifecycle orders the engine's Start, which runs on activation, against
// its Stop, which runs on shutdown. The lock guards nothing but the flags: it
// is never held across Start or Stop, so a shutdown that arrives while Start is
// running returns at once and the engine's own Stop is what interrupts the
// Start (see extension.CoordinatorEngine). The flags decide three things:
//
// - Stop before or without Start is a no-op. A standby is stopped without
// ever having been activated, and its engine was never started; the seam
// does not call Stop on it.
// - Stop before Start also cancels the Start: an activation that fires after
// shutdown began does not start an engine nothing will stop. "Before" here
// means before Start was entered, not before it was claimed: a Start is
// claimed on one goroutine and entered a few instructions later, and a Stop
// landing in between waits for that goroutine to enter or cancel, so the
// engine never sees a Stop followed by a Start.
// - Start runs at most once, whatever fires it.
type engineLifecycle struct {
mu sync.Mutex
started bool
stopped bool
entered bool
// settled is closed by the goroutine that claimed the Start, once it has
// either entered Start or given it up. It is nil until a Start is claimed.
settled chan struct{}
}
var lifecycle engineLifecycle
// beginStart claims the one Start, and refuses it once Stop has been asked for.
// The caller owns the claim and must settle it with enterStart.
func (l *engineLifecycle) beginStart() bool {
l.mu.Lock()
defer l.mu.Unlock()
if l.started || l.stopped {
return false
}
l.started = true
l.settled = make(chan struct{})
return true
}
// enterStart settles a claim: it reports whether the engine's Start is to be
// entered now, which it is unless a Stop arrived in the meantime, and releases
// any Stop waiting on the answer.
func (l *engineLifecycle) enterStart() bool {
l.mu.Lock()
defer l.mu.Unlock()
if !l.stopped {
l.entered = true
}
close(l.settled)
return l.entered
}
// beginStop claims the one Stop, and reports whether there is a started engine
// to stop. A claimed but not yet entered Start is waited for: the wait is the
// few instructions between the two, and it is what keeps the engine from being
// stopped before it is started.
func (l *engineLifecycle) beginStop() bool {
l.mu.Lock()
if l.stopped {
l.mu.Unlock()
return false
}
l.stopped = true
settled, started, entered := l.settled, l.started, l.entered
l.mu.Unlock()
if !started || entered {
return entered
}
<-settled
l.mu.Lock()
defer l.mu.Unlock()
return l.entered
}
// startCoordinatorEngine arranges for the installed engine to be started over
// the coordinator client, once this replica is ACTIVE and serving is the state
// it can call back into.
//
// The engine runs on a goroutine of its own, never on the activation callback
// chain: those callbacks run to completion before the coordinator's recovery
// barrier opens, and gRPC Serve waits on that barrier, so an engine started on
// that chain would block the very service its first call needs. waitServing is
// what it waits on instead. A start failure is fatal, through onFatal: a
// coordinator serving without its engine would accept work nothing accounts
// for, and by then the startup path it could have been returned to is gone.
func startCoordinatorEngine(
ctx context.Context,
coord types.MixCoordComponent,
client types.MixCoordClient,
waitServing func(context.Context) error,
onFatal func(error),
) error {
engine := coordinatorEngine()
if engine == nil {
return nil
}
notifier, ok := coord.(activeNotifier)
if !ok {
return merr.WrapErrServiceInternal(
"coordinator engine installed, but the coordinator does not report activation")
}
start := func() error {
if !lifecycle.beginStart() {
mlog.Info(ctx, "coordinator engine not started: already started, or stopped before activation")
return nil
}
if !lifecycle.enterStart() {
mlog.Info(ctx, "coordinator engine not started: shutdown began before it could start")
return nil
}
if err := engine.Start(ctx, client); err != nil {
return err
}
mlog.Info(ctx, "coordinator engine started")
return nil
}
notifier.OnActive(func() {
go func() {
if err := waitServing(ctx); err != nil {
mlog.Info(ctx, "coordinator engine not started: the coordinator stopped before it served",
mlog.Err(err))
return
}
if err := start(); err != nil {
onFatal(err)
}
}()
})
return nil
}
// stopCoordinatorEngine stops the installed engine if it was started; its
// error is logged, not returned, because the coordinator shutdown must not hang
// on it. It does not wait for a Start still in progress: the engine's Stop is
// what ends that Start.
func stopCoordinatorEngine(ctx context.Context) {
engine := coordinatorEngine()
if engine == nil {
return
}
if !lifecycle.beginStop() {
mlog.Info(ctx, "coordinator engine not stopped: it was never started")
return
}
if err := engine.Stop(); err != nil {
mlog.Warn(ctx, "coordinator engine stop failed", mlog.Err(err))
}
}
internal/distributed/mixcoord/extension_seam_test.go
+426
View File
@@ -0,0 +1,426 @@
// 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 grpcmixcoord
import (
"context"
"sync"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/pkg/v3/extension"
)
type recordingEngine struct {
startErr error
stopErr error
mu sync.Mutex
seenCoord extension.Coordinator
startCount int
stopCount int
}
func (e *recordingEngine) Start(_ context.Context, coord extension.Coordinator) error {
e.mu.Lock()
e.startCount++
e.seenCoord = coord
e.mu.Unlock()
return e.startErr
}
func (e *recordingEngine) Stop() error {
e.mu.Lock()
e.stopCount++
e.mu.Unlock()
return e.stopErr
}
func (e *recordingEngine) starts() int {
e.mu.Lock()
defer e.mu.Unlock()
return e.startCount
}
func (e *recordingEngine) stops() int {
e.mu.Lock()
defer e.mu.Unlock()
return e.stopCount
}
func (e *recordingEngine) coord() extension.Coordinator {
e.mu.Lock()
defer e.mu.Unlock()
return e.seenCoord
}
func installEngine(t *testing.T, e extension.CoordinatorEngine) {
t.Helper()
extension.ResetForTest()
resetEngineLifecycleForTest()
t.Cleanup(extension.ResetForTest)
t.Cleanup(resetEngineLifecycleForTest)
extension.SetCoordinatorEngine(e)
}
func newTestServer(t *testing.T, coord *mockMix) *Server {
t.Helper()
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
return &Server{
ctx: ctx,
cancel: cancel,
mixCoord: coord,
mixCoordClient: mocks.NewMockMixCoordClient(t),
grpcErrChan: make(chan error),
grpcServing: make(chan struct{}),
// Recovery never finishes here, so the server never tries to serve on
// a listener a test has not bound; serveNow is what says it answers.
recoveryWaiter: recoveryWaiterFunc(func(ctx context.Context) error {
<-ctx.Done()
return ctx.Err()
}),
}
}
// startedOn registers the seam on svr and returns the coordinator that fires
// activation, with the server already answering gRPC unless the caller says
// otherwise.
func serveNow(svr *Server) {
svr.markGrpcServing()
}
// fatalRecorder stands in for the process exit the server does on a failed
// engine start.
type fatalRecorder struct {
mu sync.Mutex
errs []error
seen chan struct{}
}
func newFatalRecorder() *fatalRecorder {
return &fatalRecorder{seen: make(chan struct{}, 1)}
}
func (f *fatalRecorder) record(err error) {
f.mu.Lock()
f.errs = append(f.errs, err)
f.mu.Unlock()
select {
case f.seen <- struct{}{}:
default:
}
}
func (f *fatalRecorder) wait(t *testing.T) error {
t.Helper()
select {
case <-f.seen:
case <-time.After(5 * time.Second):
t.Fatal("the engine start failure never reached the fatal path")
}
f.mu.Lock()
defer f.mu.Unlock()
return f.errs[0]
}
// eventually waits for the engine's Start, which now runs on its own goroutine.
func waitStartCount(t *testing.T, e *recordingEngine, want int) {
t.Helper()
assert.Eventually(t, func() bool { return e.starts() == want }, 5*time.Second, 5*time.Millisecond,
"the engine must have been started %d time(s), saw %d", want, e.starts())
}
func TestCoordinatorEngineSeamIsInertWithoutProvider(t *testing.T) {
extension.ResetForTest()
t.Cleanup(extension.ResetForTest)
assert.NoError(t, startCoordinatorEngine(context.Background(), &mockMix{}, nil,
func(context.Context) error { return nil }, func(error) {}))
assert.NotPanics(t, func() { stopCoordinatorEngine(context.Background()) })
}
// A coordinator that does not report activation cannot say which replica is
// the one that runs the engine, and starting one on every replica would have
// two of them accounting for the same cluster.
func TestAnEngineNeedsACoordinatorThatReportsActivation(t *testing.T) {
installEngine(t, &recordingEngine{})
err := startCoordinatorEngine(context.Background(), &mockMix{}, nil,
func(context.Context) error { return nil }, func(error) {})
assert.ErrorContains(t, err, "does not report activation")
}
func TestServerStartWithoutEngineLeavesNativeStartupUnchanged(t *testing.T) {
extension.ResetForTest()
t.Cleanup(extension.ResetForTest)
svr := newTestServer(t, &mockMix{})
assert.NoError(t, svr.start())
}
func TestServerStartHandsTheCoordinatorClientToTheEngine(t *testing.T) {
engine := &recordingEngine{}
installEngine(t, engine)
coord := &activatableCoord{}
svr := newTestServer(t, &coord.mockMix)
svr.mixCoord = coord
require.NoError(t, svr.start())
serveNow(svr)
coord.fire()
waitStartCount(t, engine, 1)
assert.Same(t, svr.mixCoordClient, engine.coord(),
"the engine must be handed the coordinator client, or it cannot load anything")
}
// A failed start cannot be returned: by the time it happens the startup path
// has returned, so it goes to the fatal path the server hands the seam.
func TestAFailedEngineStartReachesTheFatalPath(t *testing.T) {
want := errors.New("engine failed to start")
installEngine(t, &recordingEngine{startErr: want})
coord := &activatableCoord{}
fatal := newFatalRecorder()
require.NoError(t, startCoordinatorEngine(context.Background(), coord, nil,
func(context.Context) error { return nil }, fatal.record))
coord.fire()
assert.ErrorIs(t, fatal.wait(t), want)
}
func TestServerStopStopsEngine(t *testing.T) {
engine := &recordingEngine{stopErr: errors.New("engine stop failed")}
installEngine(t, engine)
coord := &activatableCoord{}
svr := newTestServer(t, &coord.mockMix)
svr.mixCoord = coord
require.NoError(t, svr.start())
serveNow(svr)
coord.fire()
waitStartCount(t, engine, 1)
assert.NoError(t, svr.Stop(),
"an engine that fails to stop must not fail the coordinator shutdown")
assert.Equal(t, 1, engine.stops(), "the coordinator must stop the engine exactly once")
}
// A standby is stopped without ever having been activated, so its engine was
// never started, and the seam must not call Stop on an engine it never
// started.
func TestStopWithoutStartIsANoOp(t *testing.T) {
engine := &recordingEngine{}
installEngine(t, engine)
coord := &activatableCoord{}
svr := newTestServer(t, &coord.mockMix)
svr.mixCoord = coord
assert.NoError(t, svr.start())
serveNow(svr)
require.Zero(t, engine.starts(), "not activated, so not started")
assert.NoError(t, svr.Stop())
assert.Zero(t, engine.stops(), "an engine that was never started is not stopped")
coord.fire()
assert.Never(t, func() bool { return engine.starts() > 0 }, 200*time.Millisecond, 10*time.Millisecond,
"an activation after shutdown began must not start an engine nothing will stop")
}
// blockingEngine is an engine whose Start does not return until Stop is
// called, which is what a slow activation looks like from the seam.
type blockingEngine struct {
recordingEngine
starting chan struct{} // closed once Start is running
release chan struct{} // closed by Stop, which is what lets Start return
once sync.Once
}
func newBlockingEngine() *blockingEngine {
return &blockingEngine{starting: make(chan struct{}), release: make(chan struct{})}
}
func (e *blockingEngine) Start(ctx context.Context, coord extension.Coordinator) error {
close(e.starting)
<-e.release
return e.recordingEngine.Start(ctx, coord)
}
func (e *blockingEngine) Stop() error {
e.once.Do(func() { close(e.release) })
return e.recordingEngine.Stop()
}
// A shutdown must not wait for a slow Start: the engine's own Stop is what
// ends it, so holding a lock across Start would deadlock the coordinator's
// shutdown against its own activation.
func TestStopDoesNotWaitForASlowStart(t *testing.T) {
engine := newBlockingEngine()
installEngine(t, engine)
coord := &activatableCoord{}
svr := newTestServer(t, &coord.mockMix)
svr.mixCoord = coord
require.NoError(t, svr.start())
serveNow(svr)
coord.fire()
select {
case <-engine.starting:
case <-time.After(5 * time.Second):
t.Fatal("the activation never reached Start")
}
stopped := make(chan struct{})
go func() {
defer close(stopped)
stopCoordinatorEngine(context.Background())
}()
select {
case <-stopped:
case <-time.After(2 * time.Second):
t.Fatal("Stop must return while Start is still running, not wait for it")
}
waitStartCount(t, &engine.recordingEngine, 1) // Stop is what let Start return
assert.Equal(t, 1, engine.stops())
}
type activatableCoord struct {
mockMix
pending []func()
}
func (c *activatableCoord) OnActive(fn func()) { c.pending = append(c.pending, fn) }
func (c *activatableCoord) fire() {
for _, fn := range c.pending {
fn()
}
c.pending = nil
}
func TestEngineStartWaitsForActivation(t *testing.T) {
engine := &recordingEngine{}
installEngine(t, engine)
coord := &activatableCoord{}
svr := newTestServer(t, &coord.mockMix)
svr.mixCoord = coord // the notifier interface must be visible to the seam
assert.NoError(t, svr.start())
serveNow(svr)
assert.Zero(t, engine.starts(),
"the engine must not start on a replica that is not ACTIVE yet")
coord.fire()
waitStartCount(t, engine, 1)
}
// The activation callbacks run to completion before the coordinator's recovery
// barrier opens, and gRPC Serve waits on that barrier. An engine started on
// that chain would block the service its own first call needs, so the seam
// must return from the callback and wait for serving on a goroutine of its
// own.
func TestActivationDoesNotWaitForTheEngine(t *testing.T) {
engine := newBlockingEngine()
installEngine(t, engine)
coord := &activatableCoord{}
svr := newTestServer(t, &coord.mockMix)
svr.mixCoord = coord
require.NoError(t, svr.start())
t.Cleanup(func() { _ = engine.Stop() })
serveNow(svr)
returned := make(chan struct{})
go func() {
defer close(returned)
coord.fire() // what enableExternalAccess does, before the barrier opens
}()
select {
case <-returned:
case <-time.After(5 * time.Second):
t.Fatal("activation must return without waiting for the engine's Start")
}
select {
case <-engine.starting:
case <-time.After(5 * time.Second):
t.Fatal("the engine must still be started, just not on the activation chain")
}
}
// The engine's first act is usually to call the coordinator back, so it must
// not be started before the coordinator answers gRPC - which happens after the
// recovery barrier, later than activation.
func TestTheEngineWaitsUntilTheCoordinatorAnswers(t *testing.T) {
engine := &recordingEngine{}
installEngine(t, engine)
coord := &activatableCoord{}
svr := newTestServer(t, &coord.mockMix)
svr.mixCoord = coord
require.NoError(t, svr.start())
coord.fire()
assert.Never(t, func() bool { return engine.starts() > 0 }, 200*time.Millisecond, 10*time.Millisecond,
"the engine must not start while gRPC is still waiting for recovery")
serveNow(svr)
waitStartCount(t, engine, 1)
}
// A coordinator that shuts down before it ever serves never starts its engine,
// and the goroutine waiting to start it goes away with the context.
func TestAShutdownBeforeServingNeverStartsTheEngine(t *testing.T) {
engine := &recordingEngine{}
installEngine(t, engine)
coord := &activatableCoord{}
svr := newTestServer(t, &coord.mockMix)
svr.mixCoord = coord
require.NoError(t, svr.start())
coord.fire()
svr.cancel()
assert.Never(t, func() bool { return engine.starts() > 0 }, 200*time.Millisecond, 10*time.Millisecond,
"a coordinator that never served must not start its engine")
stopCoordinatorEngine(context.Background())
assert.Zero(t, engine.stops(), "an engine that was never started is not stopped")
}
// A Start is claimed on one goroutine and entered a few instructions later. A
// Stop landing in that window must not reach an engine that has not started;
// the contract is that Stop follows Start, never the other way round.
func TestAStopBetweenClaimingAndEnteringStartCancelsIt(t *testing.T) {
engine := &recordingEngine{}
installEngine(t, engine)
require.True(t, lifecycle.beginStart(), "the Start is claimed")
stopped := make(chan struct{})
go func() {
defer close(stopped)
stopCoordinatorEngine(context.Background())
}()
assert.Never(t, func() bool {
select {
case <-stopped:
return true
default:
return false
}
}, 100*time.Millisecond, 10*time.Millisecond,
"Stop must wait for the claimed Start to be entered or given up")
assert.False(t, lifecycle.enterStart(), "a Start claimed before a Stop must give itself up")
select {
case <-stopped:
case <-time.After(5 * time.Second):
t.Fatal("Stop must return once the Start has given itself up")
}
assert.Zero(t, engine.stops(), "an engine that never entered Start must not be stopped")
assert.Zero(t, engine.starts())
}
internal/distributed/mixcoord/extension_seam_testonly.go
+30
View File
@@ -0,0 +1,30 @@
// 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.
//go:build test
package grpcmixcoord
// resetEngineLifecycleForTest forgets any Start or Stop the seam has seen, so a
// test begins with an engine that was never started.
func resetEngineLifecycleForTest() {
lifecycle.mu.Lock()
defer lifecycle.mu.Unlock()
lifecycle.started = false
lifecycle.stopped = false
lifecycle.entered = false
lifecycle.settled = nil
}
internal/distributed/mixcoord/service.go
+62
View File
@@ -67,6 +67,14 @@ type Server struct {
grpcWG sync.WaitGroup
grpcReady atomic.Bool
// grpcServing is closed once this server answers gRPC. Work that has to be
// able to call the coordinator back - the extension seam's engine - waits
// on this rather than on activation, because activation runs ahead of the
// recovery barrier that gates Serve, and work that blocked the activation
// callbacks would keep Serve from ever starting.
grpcServing chan struct{}
grpcServingOnce sync.Once
ctx context.Context
cancel context.CancelFunc
@@ -88,6 +96,7 @@ func NewServer(ctx context.Context, factory dependency.Factory) (*Server, error)
ctx: ctx1,
cancel: cancel,
grpcErrChan: make(chan error, 1),
grpcServing: make(chan struct{}),
}
mixCoordServer, err := mixcoord.NewMixCoordServer(ctx, factory)
@@ -244,6 +253,7 @@ func (s *Server) startGrpc() error {
case err := <-s.grpcErrChan:
if err == nil {
s.grpcReady.Store(true)
s.markGrpcServing()
}
return err
case <-s.ctx.Done():
@@ -266,6 +276,44 @@ func (s *Server) startGrpcLoop() {
}
}
// markGrpcServing releases everything waiting for this server to answer gRPC.
func (s *Server) markGrpcServing() {
s.grpcServingOnce.Do(func() {
if s.grpcServing != nil {
close(s.grpcServing)
}
})
}
// waitGrpcServing blocks until this server answers gRPC, or ctx ends. A Server
// built without the channel (a test that never serves) is not gated.
func (s *Server) waitGrpcServing(ctx context.Context) error {
if s.grpcServing == nil {
return nil
}
select {
case <-s.grpcServing:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
// engineStartFailed ends the process the way a gRPC startup failure after Run()
// does: this happens on a detached goroutine, so it cannot be returned to the
// startup path, and a coordinator serving without its engine would accept work
// nothing accounts for.
func (s *Server) engineStartFailed(err error) {
mlog.Error(s.ctx, "coordinator engine failed to start on activation", mlog.Err(err))
go func() {
if stopErr := s.Stop(); stopErr != nil {
mlog.Warn(s.ctx, "failed to cleanly stop MixCoord after engine startup failure", mlog.Err(stopErr))
}
mlog.Cleanup()
os.Exit(1)
}()
}
func (s *Server) startGrpcAfterRecovery(waiter recoveryWaiter) {
if err := waitForRecoveryAndStart(s.ctx, waiter, s.startGrpc); err != nil {
if s.ctx.Err() == nil {
@@ -308,6 +356,16 @@ func (s *Server) start() error {
return err
}
// The extension seam registers here; the engine itself starts on its own
// goroutine, once this replica is ACTIVE and gRPC is answering. An engine's
// first act is usually to read coordinator state, and the activation
// callbacks run ahead of the recovery barrier that gates Serve, so an
// engine started on that chain would deadlock its own first call.
if err := startCoordinatorEngine(s.ctx, s.mixCoord, s.mixCoordClient, s.waitGrpcServing, s.engineStartFailed); err != nil {
mlog.Error(s.ctx, "coordinator engine registration failed", mlog.Err(err))
return err
}
s.grpcWG.Add(1)
go func() {
defer s.grpcWG.Done()
@@ -333,6 +391,10 @@ func (s *Server) Stop() (err error) {
defer s.tikvCli.Close()
}
// The extension seam sits before the coordinator is torn down, so an
// engine still sees a working coordinator while it stops.
stopCoordinatorEngine(s.ctx)
if s.mixCoord != nil {
mlog.Info(s.ctx, "graceful stop rootCoord")
s.mixCoord.GracefulStop()
internal/proxy/authentication_interceptor_test.go
+48
View File
@@ -7,7 +7,10 @@ import (
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
"github.com/milvus-io/milvus/internal/proxy/privilege"
"github.com/milvus-io/milvus/internal/util/hookutil"
@@ -131,3 +134,48 @@ func TestAuthenticationInterceptor(t *testing.T) {
}
hookutil.SetTestHook(hookutil.DefaultHook{})
}
// streamForAuth is a minimal ServerStream carrying only the context the
// interceptor reads and rewrites.
type streamForAuth struct {
grpc.ServerStream
ctx context.Context
}
func (s *streamForAuth) Context() context.Context { return s.ctx }
type ctxKeyAuthedUser struct{}
// The stream interceptor is the unary one's counterpart: a failed
// authentication ends the stream before the handler exists, and a successful
// one hands the handler the authenticated context through ss.Context() - the
// only channel a stream handler has - rather than the raw incoming one.
func TestGrpcAuthStreamInterceptor(t *testing.T) {
info := &grpc.StreamServerInfo{FullMethod: "/milvus.proto.milvus.MilvusService/CreateReplicateStream"}
t.Run("a refused credential ends the stream unhandled", func(t *testing.T) {
intercept := GrpcAuthStreamInterceptor(func(ctx context.Context) (context.Context, error) {
return nil, status.Error(codes.Unauthenticated, "no")
})
handled := false
err := intercept(struct{}{}, &streamForAuth{ctx: context.Background()}, info,
func(any, grpc.ServerStream) error { handled = true; return nil })
assert.Error(t, err)
assert.False(t, handled, "an unauthenticated stream must never reach its handler")
})
t.Run("the handler reads the authenticated context", func(t *testing.T) {
intercept := GrpcAuthStreamInterceptor(func(ctx context.Context) (context.Context, error) {
return context.WithValue(ctx, ctxKeyAuthedUser{}, "alice"), nil
})
var seen any
err := intercept(struct{}{}, &streamForAuth{ctx: context.Background()}, info,
func(_ any, ss grpc.ServerStream) error {
seen = ss.Context().Value(ctxKeyAuthedUser{})
return nil
})
assert.NoError(t, err)
assert.Equal(t, "alice", seen,
"the rewritten context is the only way user identity reaches a stream handler")
})
}
internal/proxy/hook_interceptor.go
+29 -9
View File
@@ -15,6 +15,16 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// UnaryServerHookInterceptor installs the request hook on every unary RPC.
//
// This is the extension seam the "Extension seam, see internal/util/hookutil"
// comments elsewhere in this package point at: a deployment form supplies a
// hook - compiled into the binary or loaded from proxy.soPath - and the hook
// may answer an RPC itself (Mock), inspect or rewrite it before it runs
// (Before), or see its result (After). With no hook installed the default one
// does nothing and every RPC behaves exactly as it did. How a hook is installed
// and configured is internal/util/hookutil, whose package comment names the
// design doc.
func UnaryServerHookInterceptor() grpc.UnaryServerInterceptor {
return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
return HookInterceptor(ctx, req, GetCurUserFromContextOrDefault(ctx), info.FullMethod, handler)
@@ -37,11 +47,7 @@ func HookInterceptor(ctx context.Context, req any, userName, fullMethod string,
mlog.String("full method", fullMethod), mlog.Err(err))
metrics.ProxyHookFunc.WithLabelValues(metrics.HookMock, fullMethod).Inc()
updateProxyFunctionCallMetric(fullMethod, err)
if err != nil {
// NOTE: don't use the merr, because it will cause the wrong retry behavior in the sdk
err = status.Error(codes.InvalidArgument, "detail: "+err.Error())
}
return mockResp, err
return mockResp, hookError(err)
}
if newCtx, err = hoo.Before(ctx, req, fullMethod); err != nil {
@@ -49,8 +55,7 @@ func HookInterceptor(ctx context.Context, req any, userName, fullMethod string,
GetRequestFieldWithoutSensitiveInfo(req), mlog.Err(err))
metrics.ProxyHookFunc.WithLabelValues(metrics.HookBefore, fullMethod).Inc()
updateProxyFunctionCallMetric(fullMethod, err)
// NOTE: don't use the merr, because it will cause the wrong retry behavior in the sdk
return nil, status.Error(codes.InvalidArgument, "detail: "+err.Error())
return nil, hookError(err)
}
realResp, realErr = handler(newCtx, req)
if err = hoo.After(newCtx, realResp, realErr, fullMethod); err != nil {
@@ -58,12 +63,27 @@ func HookInterceptor(ctx context.Context, req any, userName, fullMethod string,
GetRequestFieldWithoutSensitiveInfo(req), mlog.Err(err))
metrics.ProxyHookFunc.WithLabelValues(metrics.HookAfter, fullMethod).Inc()
updateProxyFunctionCallMetric(fullMethod, err)
// NOTE: don't use the merr, because it will cause the wrong retry behavior in the sdk
return nil, status.Error(codes.InvalidArgument, "detail: "+err.Error())
return nil, hookError(err)
}
return realResp, realErr
}
// hookError is how a hook's refusal reaches the client.
//
// A refusal that must carry a classification the caller can act on does not
// come through here at all: the hook answers it from Mock, with the RPC's own
// response carrying merr.Status, which is how every milvus handler reports a
// refusal and what an SDK surfaces immediately. An error returned here can
// only become a gRPC status, and a bare error becomes codes.Unknown, which
// clients retry - the reason the original comment gives for not using merr.
func hookError(err error) error {
if err == nil {
return nil
}
// NOTE: don't use the merr, because it will cause the wrong retry behavior in the sdk
return status.Error(codes.InvalidArgument, "detail: "+err.Error())
}
func updateProxyFunctionCallMetric(fullMethod string, err error) {
strs := strings.Split(fullMethod, "/")
method := strs[len(strs)-1]
internal/proxy/hook_interceptor_test.go
+21
View File
@@ -9,6 +9,8 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/util/hookutil"
@@ -209,3 +211,22 @@ func TestUpdateProxyFunctionCallMetric(t *testing.T) {
updateProxyFunctionCallMetric("", nil)
})
}
// A refusal returned from Before reaches the client as InvalidArgument, not
// as codes.Unknown: a bare error would be retried by the SDK forever, which is
// what the note in hookError is about. A hook that needs the caller to see a
// classification answers from Mock instead.
func TestABeforeRefusalIsNotRetriableAtTheTransport(t *testing.T) {
// Consume the lazy plugin init first, or the GetHook inside the
// interceptor would run it and put the default hook back over this one.
hookutil.InitOnceHook()
hookutil.SetTestHook(beforeMock{err: merr.WrapErrServiceUnavailable("not ready")})
defer hookutil.SetTestHook(hookutil.DefaultHook{})
_, err := UnaryServerHookInterceptor()(context.Background(), &req{},
&grpc.UnaryServerInfo{FullMethod: "insert"},
func(ctx context.Context, req interface{}) (interface{}, error) { return nil, nil })
require.Error(t, err)
assert.Equal(t, codes.InvalidArgument, status.Code(err))
assert.False(t, merr.IsMilvusError(err))
}
internal/proxy/hook_stream_test.go
+333
View File
@@ -0,0 +1,333 @@
// 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 proxy
import (
"context"
"testing"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/hook"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/proxy/replicate"
"github.com/milvus-io/milvus/internal/util/hookutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
type fakeReplicateStream struct {
grpc.ServerStream
ctx context.Context
}
func (f fakeReplicateStream) Context() context.Context { return f.ctx }
func (f fakeReplicateStream) Send(*milvuspb.ReplicateResponse) error { return nil }
func (f fakeReplicateStream) Recv() (*milvuspb.ReplicateRequest, error) { return nil, nil }
// inspectingStreamHook is a hook that reads its request, as a plug-in built
// against hook.Hook is entitled to: it does not embed DefaultHook, so every
// method here is the one the stream actually calls, and a nil req would be
// caught by the type assertion rather than ignored.
type inspectingStreamHook struct {
t *testing.T
refuse error
mocks []string
befores []string
afters []string
results []any
errs []error
}
func (h *inspectingStreamHook) Init(map[string]string) error { return nil }
func (h *inspectingStreamHook) VerifyAPIKey(string) (string, error) { return "", nil }
func (h *inspectingStreamHook) Release() {}
func (h *inspectingStreamHook) request(req interface{}) *milvuspb.ReplicateRequest {
h.t.Helper()
require.NotNil(h.t, req, "a stream consults the hook with a request, never nil")
typed, ok := req.(*milvuspb.ReplicateRequest)
require.True(h.t, ok, "the request is the message type the stream carries, got %T", req)
return typed
}
func (h *inspectingStreamHook) Mock(_ context.Context, req interface{}, fullMethod string) (bool, interface{}, error) {
h.request(req)
h.mocks = append(h.mocks, fullMethod)
return false, nil, nil
}
func (h *inspectingStreamHook) Before(ctx context.Context, req interface{}, fullMethod string) (context.Context, error) {
h.request(req)
h.befores = append(h.befores, fullMethod)
if h.refuse != nil {
return ctx, h.refuse
}
// What a hook puts on the context must reach the stream: the cluster id
// is what the replicate server reads first - off the incoming metadata,
// as a client sends it - and the fake stream below deliberately carries
// none of its own.
return metadata.NewIncomingContext(ctx, metadata.Pairs("cluster-id", "cluster-from-the-hook")), nil
}
func (h *inspectingStreamHook) After(_ context.Context, result interface{}, err error, fullMethod string) error {
h.afters = append(h.afters, fullMethod)
h.results = append(h.results, result)
h.errs = append(h.errs, err)
return nil
}
var _ hook.Hook = (*inspectingStreamHook)(nil)
func installStreamHook(t *testing.T, h hook.Hook) {
t.Helper()
hookutil.InitOnceHook()
hookutil.SetTestHook(h)
t.Cleanup(func() { hookutil.SetTestHook(hookutil.DefaultHook{}) })
}
// CreateReplicateStream is one of the two streams that consult the hook by hand
// - DumpMessages is the other - because the interceptor that consults it for
// every other RPC is a unary one and an interceptor chain binds to one of
// gRPC's two call kinds. It must consult it
// the same way even so: Mock, Before and After in order, a typed non-nil
// request, and the handler under the context Before returned. The stream
// server here is stopped at its first read, so the test is about the seam and
// not about replication.
func TestCreateReplicateStreamConsultsTheHookLikeAUnaryRPC(t *testing.T) {
h := &inspectingStreamHook{t: t}
installStreamHook(t, h)
execute := mockey.Mock((*replicate.ReplicateStreamServer).Execute).Return(nil).Build()
defer execute.UnPatch()
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.CreateReplicateStream(fakeReplicateStream{ctx: context.Background()})
require.NoError(t, err,
"the replicate server reads the cluster id off the stream's context, and only the hook's context carries one")
want := []string{milvuspb.MilvusService_CreateReplicateStream_FullMethodName}
assert.Equal(t, want, h.mocks, "Mock is consulted, as for every other RPC")
assert.Equal(t, want, h.befores)
assert.Equal(t, want, h.afters, "After sees the stream end, as it sees every other RPC's result")
assert.Equal(t, []error{nil}, h.errs)
}
// A refusal from Before travels as a gRPC status a client does not retry, and
// After is not consulted for an RPC that never ran - the same shape the unary
// interceptor gives a refusal.
func TestCreateReplicateStreamRefusedByTheHook(t *testing.T) {
h := &inspectingStreamHook{t: t, refuse: merr.ErrServiceUnimplemented}
installStreamHook(t, h)
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.CreateReplicateStream(fakeReplicateStream{ctx: context.Background()})
require.Error(t, err)
assert.Equal(t, codes.InvalidArgument, status.Code(err),
"a stream refusal travels as a gRPC status, and must not be the codes.Unknown a client retries")
assert.Len(t, h.befores, 1)
assert.Empty(t, h.afters, "a refused RPC has no result for After to see")
}
// A hook that Mocks the stream ends it with the hook's verdict: there is no
// stream to send a mocked answer down, so what the hook returns as its error
// is what the caller gets, and the stream server never starts.
func TestCreateReplicateStreamMockedByTheHook(t *testing.T) {
installStreamHook(t, mockingStreamHook{})
started := false
execute := mockey.Mock((*replicate.ReplicateStreamServer).Execute).
To(func(*replicate.ReplicateStreamServer) error { started = true; return nil }).Build()
defer execute.UnPatch()
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.CreateReplicateStream(fakeReplicateStream{ctx: context.Background()})
require.Error(t, err)
assert.Equal(t, codes.InvalidArgument, status.Code(err))
assert.False(t, started, "a mocked stream never reaches the replicate server")
}
// mockingStreamHook answers the replicate stream from Mock, refusing it.
type mockingStreamHook struct {
hookutil.DefaultHook
}
func (mockingStreamHook) Mock(_ context.Context, _ interface{}, fullMethod string) (bool, interface{}, error) {
if fullMethod == milvuspb.MilvusService_CreateReplicateStream_FullMethodName {
return true, nil, merr.ErrServiceUnimplemented
}
return false, nil, nil
}
// dumpRequestHook records how DumpMessages consults it, and checks that the
// request it is shown is the dump request itself: this stream, unlike the
// replicate stream, carries a request message of its own.
type dumpRequestHook struct {
t *testing.T
want *milvuspb.DumpMessagesRequest
refuse error
calls []string
errs []error
}
type dumpHookContextKey struct{}
func (h *dumpRequestHook) Init(map[string]string) error { return nil }
func (h *dumpRequestHook) VerifyAPIKey(string) (string, error) { return "", nil }
func (h *dumpRequestHook) Release() {}
func (h *dumpRequestHook) Mock(_ context.Context, req interface{}, fullMethod string) (bool, interface{}, error) {
assert.Same(h.t, h.want, req, "a dump consults the hook with its own request")
h.calls = append(h.calls, "Mock "+fullMethod)
return false, nil, nil
}
func (h *dumpRequestHook) Before(ctx context.Context, req interface{}, fullMethod string) (context.Context, error) {
assert.Same(h.t, h.want, req, "a dump consults the hook with its own request")
h.calls = append(h.calls, "Before "+fullMethod)
if h.refuse != nil {
return ctx, h.refuse
}
return context.WithValue(ctx, dumpHookContextKey{}, "from-the-hook"), nil
}
func (h *dumpRequestHook) After(_ context.Context, _ interface{}, err error, fullMethod string) error {
h.calls = append(h.calls, "After "+fullMethod)
h.errs = append(h.errs, err)
return nil
}
var _ hook.Hook = (*dumpRequestHook)(nil)
// DumpMessages is the service's other stream, and it consults the hook the way
// CreateReplicateStream does: Mock, Before and After in order, the dump running
// under the context Before returned, and After seeing how the dump ended. The
// dump itself is stopped at once, so the test is about the seam and not about
// reading a WAL.
func TestDumpMessagesConsultsTheHookLikeAUnaryRPC(t *testing.T) {
req := &milvuspb.DumpMessagesRequest{}
h := &dumpRequestHook{t: t, want: req}
installStreamHook(t, h)
var seen any
dump := mockey.Mock((*Proxy).dumpMessages).
To(func(_ *Proxy, _ *milvuspb.DumpMessagesRequest, stream milvuspb.MilvusService_DumpMessagesServer) error {
seen = stream.Context().Value(dumpHookContextKey{})
return merr.WrapErrParameterMissing("pchannel")
}).Build()
defer dump.UnPatch()
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(req, &mockDumpMessagesServer{ctx: context.Background()})
assert.ErrorIs(t, err, merr.ErrParameterMissing, "the dump's own error reaches the caller unchanged")
method := milvuspb.MilvusService_DumpMessages_FullMethodName
assert.Equal(t, []string{"Mock " + method, "Before " + method, "After " + method}, h.calls)
assert.Equal(t, "from-the-hook", seen, "the dump runs under the context Before returned")
require.Len(t, h.errs, 1)
assert.ErrorIs(t, h.errs[0], merr.ErrParameterMissing, "After sees how the dump ended")
}
// A hook that refuses the dump from Before stops it before any WAL is read: the
// refusal travels as the non-retried InvalidArgument every other RPC's refusal
// does, and After is not consulted for a dump that never ran.
func TestDumpMessagesRefusedByTheHook(t *testing.T) {
req := &milvuspb.DumpMessagesRequest{}
h := &dumpRequestHook{t: t, want: req, refuse: merr.ErrServiceUnimplemented}
installStreamHook(t, h)
ran := false
dump := mockey.Mock((*Proxy).dumpMessages).
To(func(*Proxy, *milvuspb.DumpMessagesRequest, milvuspb.MilvusService_DumpMessagesServer) error {
ran = true
return nil
}).Build()
defer dump.UnPatch()
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(req, &mockDumpMessagesServer{ctx: context.Background()})
require.Error(t, err)
assert.Equal(t, codes.InvalidArgument, status.Code(err),
"a stream refusal travels as a gRPC status, and must not be the codes.Unknown a client retries")
assert.False(t, ran, "a refused dump never reads the WAL")
method := milvuspb.MilvusService_DumpMessages_FullMethodName
assert.Equal(t, []string{"Mock " + method, "Before " + method}, h.calls, "a refused RPC has no result for After to see")
}
// refusingHook refuses every RPC from Before, and records which it was asked
// about.
type refusingHook struct {
hookutil.DefaultHook
befores []string
}
func (h *refusingHook) Before(ctx context.Context, _ interface{}, fullMethod string) (context.Context, error) {
h.befores = append(h.befores, fullMethod)
return ctx, merr.ErrServiceUnimplemented
}
// A unary RPC milvus-proto adds is covered by the interceptor the moment it
// exists; a stream is consulted by hand, so a stream it adds is one a hook
// would silently never see. This drives every stream the proxy consults the
// hook for against a hook that refuses everything, and fails as soon as the
// service declares a stream the list does not know.
func TestEveryStreamTheServiceDeclaresConsultsTheHook(t *testing.T) {
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
consulted := map[string]func() error{
"CreateReplicateStream": func() error {
return node.CreateReplicateStream(fakeReplicateStream{ctx: context.Background()})
},
"DumpMessages": func() error {
return node.DumpMessages(&milvuspb.DumpMessagesRequest{}, &mockDumpMessagesServer{ctx: context.Background()})
},
}
declared := make([]string, 0, len(milvuspb.MilvusService_ServiceDesc.Streams))
for _, stream := range milvuspb.MilvusService_ServiceDesc.Streams {
declared = append(declared, stream.StreamName)
}
known := make([]string, 0, len(consulted))
for name := range consulted {
known = append(known, name)
}
assert.ElementsMatch(t, known, declared,
"every stream the service declares must consult the hook by hand, as the ones listed here do")
for name, call := range consulted {
h := &refusingHook{}
installStreamHook(t, h)
err := call()
assert.Equal(t, codes.InvalidArgument, status.Code(err), "%s: the hook's refusal ends the stream", name)
assert.Equal(t, []string{"/" + milvuspb.MilvusService_ServiceDesc.ServiceName + "/" + name}, h.befores,
"%s consults Before", name)
}
}
internal/proxy/impl.go
+71 -2
View File
@@ -3158,6 +3158,8 @@ func (node *Proxy) search(ctx context.Context, request *milvuspb.SearchRequest,
collectionName,
).Observe(float64(searchDur))
observeResourceGroupSQLatency(ctx, metrics.SearchLabel, dbName, collectionName, searchDur)
if Params.QueryNodeCfg.StorageUsageTrackingEnabled.GetAsBool() {
metrics.ProxyScannedRemoteMB.WithLabelValues(
nodeID,
@@ -3394,6 +3396,8 @@ func (node *Proxy) hybridSearch(ctx context.Context, request *milvuspb.HybridSea
collectionName,
).Observe(float64(searchDur))
observeResourceGroupSQLatency(ctx, metrics.HybridSearchLabel, dbName, collectionName, searchDur)
if Params.QueryNodeCfg.StorageUsageTrackingEnabled.GetAsBool() {
metrics.ProxyScannedRemoteMB.WithLabelValues(
nodeID,
@@ -3903,6 +3907,9 @@ func (node *Proxy) query(ctx context.Context, qt *queryTask, sp trace.Span) (*mi
request.DbName,
request.CollectionName,
).Observe(float64(tr.ElapseSpan().Milliseconds()))
observeResourceGroupSQLatency(ctx, queryLabel, request.GetDbName(), request.GetCollectionName(),
tr.ElapseSpan().Milliseconds())
}
succeeded = true
@@ -7158,9 +7165,41 @@ func (node *Proxy) GetReplicateInfo(ctx context.Context, req *milvuspb.GetReplic
}
// CreateReplicateStream establishes a replication stream on the target Milvus cluster.
func (node *Proxy) CreateReplicateStream(stream milvuspb.MilvusService_CreateReplicateStreamServer) (err error) {
func (node *Proxy) CreateReplicateStream(stream milvuspb.MilvusService_CreateReplicateStreamServer) error {
// The hook is consulted by hand here, and in DumpMessages - the service's
// two streams - because the interceptor that consults it for every other
// RPC is a UNARY one and this is a stream: an interceptor chain binds to
// one of gRPC's two call kinds.
// It is consulted the same way even so, through HookInterceptor, so a hook
// sees Mock, Before and After in the order every other RPC gives it, and
// sees a typed, non-nil request: a stream has no request message of its
// own, and hook.Hook has no "req may be nil" contract, so the request is
// the empty message of the type this stream carries. The handler runs
// under the context Before returned, which is how a hook hands the stream
// what it put there. A Mock answer has no stream to be sent down, so a
// mocked stream simply ends with the hook's verdict. The default hook
// answers nothing, so a stock binary is unchanged.
ctx := stream.Context()
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateReplicateStream")
_, err := HookInterceptor(ctx, &milvuspb.ReplicateRequest{}, GetCurUserFromContextOrDefault(ctx),
milvuspb.MilvusService_CreateReplicateStream_FullMethodName,
func(ctx context.Context, _ any) (any, error) {
return nil, node.createReplicateStream(replicateStreamWithContext{stream, ctx})
})
return err
}
// replicateStreamWithContext is the replicate stream under the context the
// hook's Before returned, so everything the stream server reads off its
// context - the cluster id, the trace - sees what the hook put there.
type replicateStreamWithContext struct {
milvuspb.MilvusService_CreateReplicateStreamServer
ctx context.Context
}
func (s replicateStreamWithContext) Context() context.Context { return s.ctx }
func (node *Proxy) createReplicateStream(stream milvuspb.MilvusService_CreateReplicateStreamServer) (err error) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(stream.Context(), "Proxy-CreateReplicateStream")
defer sp.End()
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
@@ -7313,6 +7352,36 @@ func dumpOneMessage(
// gaps or duplicates are possible if that point is filtered, inside a
// transaction, or no longer retained by the underlying WAL.
func (node *Proxy) DumpMessages(req *milvuspb.DumpMessagesRequest, stream milvuspb.MilvusService_DumpMessagesServer) error {
// The hook is consulted by hand here for the reason CreateReplicateStream
// gives: this is a stream, and the interceptor that consults it for every
// other RPC is a unary one. It is consulted the same way - Mock, Before
// and After in order, through HookInterceptor - with the request this
// stream does carry, and the dump runs under the context Before returned.
// Without it the dump would be the one RPC on the proxy a hook never sees,
// so a policy a hook enforces from Before - refusing a username and
// password on the external listener, say - would not reach a raw WAL
// read. A Mock answer has no stream to be sent down, so a mocked dump
// ends with the hook's verdict. The default hook answers nothing, so a
// stock binary is unchanged.
ctx := stream.Context()
_, err := HookInterceptor(ctx, req, GetCurUserFromContextOrDefault(ctx),
milvuspb.MilvusService_DumpMessages_FullMethodName,
func(ctx context.Context, _ any) (any, error) {
return nil, node.dumpMessages(req, dumpMessagesStreamWithContext{stream, ctx})
})
return err
}
// dumpMessagesStreamWithContext is the dump stream, answering the context the
// hook's Before returned in place of the stream's own.
type dumpMessagesStreamWithContext struct {
milvuspb.MilvusService_DumpMessagesServer
ctx context.Context
}
func (s dumpMessagesStreamWithContext) Context() context.Context { return s.ctx }
func (node *Proxy) dumpMessages(req *milvuspb.DumpMessagesRequest, stream milvuspb.MilvusService_DumpMessagesServer) error {
ctx := stream.Context()
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DumpMessages")
defer sp.End()
internal/proxy/resource_group_metrics.go
+44
View File
@@ -0,0 +1,44 @@
// 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 proxy
import (
"context"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// observeResourceGroupSQLatency records a finished search or query against the
// resource group the request was scoped to (extension.WithQueryResourceGroup).
// An unscoped request - every request on a deployment that does not pin
// queries to a resource group - emits nothing, so the series exist only where
// the scope does.
func observeResourceGroupSQLatency(ctx context.Context, queryType, dbName, collectionName string, latencyMs int64) {
resourceGroup := extension.QueryResourceGroupFromContext(ctx)
if resourceGroup == "" {
return
}
metrics.ProxyResourceGroupSQLatency.WithLabelValues(
paramtable.GetStringNodeID(),
queryType,
dbName,
collectionName,
resourceGroup,
).Observe(float64(latencyMs))
}
internal/proxy/resource_group_metrics_test.go
+151
View File
@@ -0,0 +1,151 @@
// 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 proxy
import (
"context"
"testing"
"github.com/bytedance/mockey"
"github.com/prometheus/client_golang/prometheus"
dto "github.com/prometheus/client_model/go"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"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/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/internal/util/hookutil"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
func TestObserveResourceGroupSQLatencyEmitsNoSeriesForAnUnscopedRequest(t *testing.T) {
paramtable.Init()
metrics.ProxyResourceGroupSQLatency.Reset()
observeResourceGroupSQLatency(context.Background(), metrics.SearchLabel, "db", "coll", 12)
assert.Equal(t, 0, collectSeries(t, metrics.ProxyResourceGroupSQLatency),
"a request nothing scoped to a resource group must leave no per-resource-group series")
}
func TestObserveResourceGroupSQLatencyAttributesToTheScopeRoutingUsed(t *testing.T) {
paramtable.Init()
metrics.ProxyResourceGroupSQLatency.Reset()
ctx := extension.WithQueryResourceGroup(context.Background(), "rg-a")
observeResourceGroupSQLatency(ctx, metrics.SearchLabel, "db", "coll", 17)
h, err := metrics.ProxyResourceGroupSQLatency.GetMetricWithLabelValues(
paramtable.GetStringNodeID(), metrics.SearchLabel, "db", "coll", "rg-a")
require.NoError(t, err)
assert.Equal(t, uint64(1), histogramCount(t, h),
"the completed search must land in the series of the resource group that served it")
assert.Equal(t, 1, collectSeries(t, metrics.ProxyResourceGroupSQLatency))
}
func collectSeries(t *testing.T, c prometheus.Collector) int {
t.Helper()
ch := make(chan prometheus.Metric, 64)
go func() {
c.Collect(ch)
close(ch)
}()
n := 0
for range ch {
n++
}
return n
}
func histogramCount(t *testing.T, o prometheus.Observer) uint64 {
t.Helper()
m, ok := o.(prometheus.Metric)
require.True(t, ok, "a histogram child must also be a Metric")
var pb dto.Metric
require.NoError(t, m.Write(&pb))
return pb.GetHistogram().GetSampleCount()
}
// scopingHook is a request hook that pins every request to one resource group
// from Before, the way a deployment that routes queries per resource group
// does.
type scopingHook struct {
hookutil.DefaultHook
resourceGroup string
}
func (h scopingHook) Before(ctx context.Context, _ interface{}, _ string) (context.Context, error) {
return extension.WithQueryResourceGroup(ctx, h.resourceGroup), nil
}
// The context a hook's Before returns is the context the handler runs under,
// so a resource group the hook pins reaches the search task - and from there
// the shard-leader routing - with no plumbing of its own.
func TestSearchRunsScopedToTheResourceGroupTheHookNamed(t *testing.T) {
hookutil.InitOnceHook()
hookutil.SetTestHook(scopingHook{resourceGroup: "rg-a"})
defer hookutil.SetTestHook(hookutil.DefaultHook{})
var scope string
seen := false
defer mockey.Mock((*scheduler.BaseTaskQueue).Enqueue).To(func(_ *scheduler.BaseTaskQueue, t taskmodel.Task) error {
scope = extension.QueryResourceGroupFromContext(t.TraceCtx())
seen = true
return nil
}).Build().UnPatch()
defer mockey.Mock((*taskmodel.TaskCondition).WaitToFinish).Return(nil).Build().UnPatch()
node := &Proxy{sched: &scheduler.TaskScheduler{DqQueue: &scheduler.DqTaskQueue{}}}
node.UpdateStateCode(commonpb.StateCode_Healthy)
_, err := HookInterceptor(context.Background(),
&milvuspb.SearchRequest{DbName: "db", CollectionName: "coll"}, "user",
milvuspb.MilvusService_Search_FullMethodName,
func(ctx context.Context, req any) (any, error) {
return node.Search(ctx, req.(*milvuspb.SearchRequest))
})
require.NoError(t, err)
require.True(t, seen, "the search must have reached the task queue")
assert.Equal(t, "rg-a", scope,
"the task must run scoped to the resource group the hook named, or routing ignores it")
}
func TestSearchTaskCarriesNoScopeWithTheDefaultHook(t *testing.T) {
hookutil.InitOnceHook()
hookutil.SetTestHook(hookutil.DefaultHook{})
scope := "unset"
defer mockey.Mock((*scheduler.BaseTaskQueue).Enqueue).To(func(_ *scheduler.BaseTaskQueue, t taskmodel.Task) error {
scope = extension.QueryResourceGroupFromContext(t.TraceCtx())
return nil
}).Build().UnPatch()
defer mockey.Mock((*taskmodel.TaskCondition).WaitToFinish).Return(nil).Build().UnPatch()
node := &Proxy{sched: &scheduler.TaskScheduler{DqQueue: &scheduler.DqTaskQueue{}}}
node.UpdateStateCode(commonpb.StateCode_Healthy)
_, err := HookInterceptor(context.Background(),
&milvuspb.SearchRequest{DbName: "db", CollectionName: "coll"}, "user",
milvuspb.MilvusService_Search_FullMethodName,
func(ctx context.Context, req any) (any, error) {
return node.Search(ctx, req.(*milvuspb.SearchRequest))
})
require.NoError(t, err)
assert.Equal(t, "", scope, "a stock binary must route across every replica of the collection")
}
internal/proxy/shardclient/lb_policy.go
+17 -6
View File
@@ -54,13 +54,16 @@ type ChannelWorkload struct {
// ErrCollectionNotFullyLoaded (retriable), see selectNode. Empty is the
// absence of a scope and leaves routing exactly as before.
//
// EVERY construction site must carry it, and the way to do that is
// CollectionWorkLoad.ForChannel rather than a literal. Three paths build a
// ChannelWorkload directly instead of going through Execute -- the
// namespace single-shard fast paths in task_search.go, task_query.go and
// task_delete.go -- and one that forgets does not fail: it routes that
// subset of requests to another group's leader, silently, which is the
// A construction site does NOT have to set it. Every entry point of this
// package stamps the request's scope on from the context (see
// resource_group_scope.go), because a site that forgot would not fail - it would
// route that subset of requests to another group's leader, silently, the
// same wrong-routing-with-no-signal failure as filtering the channel map.
// Three paths build a ChannelWorkload directly rather than going through
// Execute - the namespace single-shard fast paths in task_search.go,
// task_query.go and task_delete.go - and they are covered by the stamp on
// ExecuteWithRetry. Setting it explicitly still wins, for a caller that
// names a scope the context does not carry.
ResourceGroup string
}
@@ -344,6 +347,10 @@ func (lb *LBPolicyImpl) selectNode(ctx context.Context, balancer LBBalancer, wor
// ExecuteWithRetry will choose a qn to execute the workload, and retry if failed, until reach the max retryTimes.
func (lb *LBPolicyImpl) ExecuteWithRetry(ctx context.Context, workload ChannelWorkload) error {
// Extension seam, see resource_group_scope.go: the request's routing scope is
// stamped on here, at the load balancer's own entry, so that no call site
// can silently omit it. Empty with no provider installed.
workload = scopedChannelWorkload(ctx, workload)
log := mlog.With(
mlog.Int64("collectionID", workload.CollectionID),
mlog.String("channelName", workload.Channel),
@@ -505,6 +512,8 @@ func (lb *LBPolicyImpl) ExecuteWithRetry(ctx context.Context, workload ChannelWo
// Execute will execute collection workload in parallel
func (lb *LBPolicyImpl) Execute(ctx context.Context, workload CollectionWorkLoad) error {
// Extension seam, see resource_group_scope.go.
workload = scopedCollectionWorkload(ctx, workload)
log := mlog.With(
mlog.Int64("collectionID", workload.CollectionID),
)
@@ -535,6 +544,8 @@ func (lb *LBPolicyImpl) Execute(ctx context.Context, workload CollectionWorkLoad
// ExecuteOneChannel will execute at any one channel in collection
func (lb *LBPolicyImpl) ExecuteOneChannel(ctx context.Context, workload CollectionWorkLoad) error {
// Extension seam, see resource_group_scope.go.
workload = scopedCollectionWorkload(ctx, workload)
// Unscoped: any one channel will do, so the first is taken. Scoped: the
// channel list is a map's key order, so the first channel may be one the
// group has no leader on; that refusal is the retriable
internal/proxy/shardclient/resource_group_scope.go
+42
View File
@@ -0,0 +1,42 @@
// 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 shardclient
import (
"context"
"github.com/milvus-io/milvus/pkg/v3/extension"
)
// A request scoped to a resource group (extension.WithQueryResourceGroup on
// its context) is routed to the shard leaders whose replica lives in that
// group and nowhere else. A request nothing scoped keeps an empty
// ResourceGroup and is routed across every replica, as always.
func scopedCollectionWorkload(ctx context.Context, workload CollectionWorkLoad) CollectionWorkLoad {
if workload.ResourceGroup == "" {
workload.ResourceGroup = extension.QueryResourceGroupFromContext(ctx)
}
return workload
}
func scopedChannelWorkload(ctx context.Context, workload ChannelWorkload) ChannelWorkload {
if workload.ResourceGroup == "" {
workload.ResourceGroup = extension.QueryResourceGroupFromContext(ctx)
}
return workload
}
internal/proxy/shardclient/resource_group_scope_test.go
+51
View File
@@ -0,0 +1,51 @@
// 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 shardclient
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/pkg/v3/extension"
)
func TestWorkloadsStayUnscopedForAnUnscopedRequest(t *testing.T) {
ctx := context.Background()
assert.Empty(t, scopedCollectionWorkload(ctx, CollectionWorkLoad{CollectionID: 1}).ResourceGroup,
"a request that was never scoped must route across every replica")
assert.Empty(t, scopedChannelWorkload(ctx, ChannelWorkload{CollectionID: 1}).ResourceGroup)
}
func TestWorkloadsAreStampedWithTheRequestScope(t *testing.T) {
ctx := extension.WithQueryResourceGroup(context.Background(), "rg-b")
collection := scopedCollectionWorkload(ctx, CollectionWorkLoad{CollectionID: 1})
assert.Equal(t, "rg-b", collection.ResourceGroup)
channel := scopedChannelWorkload(ctx, ChannelWorkload{CollectionID: 1, Channel: "ch0"})
assert.Equal(t, "rg-b", channel.ResourceGroup)
assert.Equal(t, "rg-b", collection.ForChannel("ch0", 0).ResourceGroup,
"the scope must follow the collection workload onto every channel it fans out to")
}
func TestAnExplicitWorkloadScopeIsKept(t *testing.T) {
ctx := extension.WithQueryResourceGroup(context.Background(), "rg-b")
collection := scopedCollectionWorkload(ctx, CollectionWorkLoad{CollectionID: 1, ResourceGroup: "rg-explicit"})
assert.Equal(t, "rg-explicit", collection.ResourceGroup)
channel := scopedChannelWorkload(ctx, ChannelWorkload{CollectionID: 1, ResourceGroup: "rg-explicit"})
assert.Equal(t, "rg-explicit", channel.ResourceGroup)
}
internal/proxy/shardclient/shard_client.go
+12 -6
View File
@@ -47,12 +47,18 @@ type NodeInfo struct {
// it. It must not be read as "no resource group", and an unknown entry
// never matches a named group in FilterByResourceGroup.
//
// The one reader is FilterByResourceGroup, applied by LBPolicyImpl.selectNode
// when ChannelWorkload.ResourceGroup names a group. Nothing on the proxy's
// request path sets that field yet, so deployed traffic still builds its
// candidate sets purely from Serviceable; the filter ships with the tag so
// that the constraints on consuming it (see FilterByResourceGroup) live in
// one place rather than with each future caller.
// Written by the coordinator's shard-leader answer and matched against the
// group a workload names. That group is stamped from the request context
// at the three workload entry points - scopedCollectionWorkload and
// scopedChannelWorkload (resource_group_scope.go), called by
// LBPolicyImpl.Execute, ExecuteOneChannel and ExecuteWithRetry - and is
// read through FilterByResourceGroup in LBPolicyImpl.selectNode (the
// candidate set), in ExecuteWithRetry (whether every leader of the scope
// has been excluded, which decides the refresh-and-retry) and in
// ExecuteOneChannel (the pre-filter that picks a channel the group can
// serve). A request that names no group keeps building its candidate sets
// from Serviceable alone. The constraints on consuming the field live with
// FilterByResourceGroup rather than with each of those callers.
ResourceGroup string
}
internal/querycoordv2/assign/streaming_query_node_channel_helper.go
+1 -1
View File
@@ -59,7 +59,7 @@ func assignChannelToWALLocatedFirstForNodeInfo(
// filterSQNIfStreamingServiceEnabled filter out the non-sqn querynode.
func filterSQNIfStreamingServiceEnabled(nodes []int64) []int64 {
if streamingutil.IsStreamingServiceEnabled() {
if streamingutil.UseStreamingQueryNodeAsDelegator() {
sqns := snmanager.StaticStreamingNodeManager.GetStreamingQueryNodeIDs()
expectedSQNs := make([]int64, 0, len(nodes))
unexpectedNodes := make([]int64, 0)
internal/querycoordv2/balance/channel_level_score_balancer.go
+1 -1
View File
@@ -81,7 +81,7 @@ func (b *ChannelLevelScoreBalancer) BalanceReplica(ctx context.Context, replica
}
}()
streamingEnabled := streamingutil.IsStreamingServiceEnabled()
streamingEnabled := streamingutil.UseStreamingQueryNodeAsDelegator()
if streamingEnabled {
// Make a plan to rebalance the channel first.
// The Streaming QueryNode doesn't make the channel level score, so just fallback to the ScoreBasedBalancer.
internal/querycoordv2/balance/helper.go
+2 -2
View File
@@ -33,7 +33,7 @@ type BalanceReplicaHelper struct {
// When streaming service is enabled, it uses the compatibility helper;
// otherwise it returns the replica's RW nodes directly.
func (h *BalanceReplicaHelper) GetRWNodesForChannels(replica *meta.Replica) []int64 {
if streamingutil.IsStreamingServiceEnabled() {
if streamingutil.UseStreamingQueryNodeAsDelegator() {
rwNodes, _ := utils.GetChannelRWAndRONodesFor260(replica, h.nodeManager)
return rwNodes
}
@@ -44,7 +44,7 @@ func (h *BalanceReplicaHelper) GetRWNodesForChannels(replica *meta.Replica) []in
// When streaming service is enabled, it uses the compatibility helper;
// otherwise it returns the replica's RW and RO nodes directly.
func (h *BalanceReplicaHelper) GetRWAndRONodesForChannels(replica *meta.Replica) (rwNodes []int64, roNodes []int64) {
if streamingutil.IsStreamingServiceEnabled() {
if streamingutil.UseStreamingQueryNodeAsDelegator() {
return utils.GetChannelRWAndRONodesFor260(replica, h.nodeManager)
}
return replica.GetRWNodes(), replica.GetRONodes()
internal/querycoordv2/balance/helper_form_test.go
+75
View File
@@ -0,0 +1,75 @@
// 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 balance
import (
"testing"
"github.com/blang/semver/v4"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/internal/util/streamingutil"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
)
// The nodes a channel is balanced across follow where delegators live. On a
// stock binary with the streaming service on that is the replica's streaming
// query nodes; under an installed form it is the replica's regular query
// nodes, RO ones included, so a query node leaving its cluster is drained of
// its delegator exactly as it was before the streaming service existed.
func TestChannelBalanceNodesFollowWhereDelegatorsLive(t *testing.T) {
streaming := mockey.Mock(streamingutil.IsStreamingServiceEnabled).Return(true).Build()
defer streaming.UnPatch()
replica := meta.NewReplica(&querypb.Replica{
ID: 1, CollectionID: 1, ResourceGroup: "cluster_a",
Nodes: []int64{11}, RoNodes: []int64{12}, RwSqNodes: []int64{7},
})
// Registered at a current version: a query node the node manager does not
// know, or one older than 2.6, is listed for draining by the upgrade path
// (utils.GetChannelRWAndRONodesFor260), which is not what this is about.
nodeMgr := session.NewNodeManager()
for _, nodeID := range []int64{11, 12} {
nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: nodeID, Address: "localhost", Hostname: "localhost", Version: semver.MustParse("2.6.0"),
}))
}
helper := &BalanceReplicaHelper{nodeManager: nodeMgr}
t.Run("stock binary", func(t *testing.T) {
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
assert.Equal(t, []int64{7}, helper.GetRWNodesForChannels(replica))
rw, ro := helper.GetRWAndRONodesForChannels(replica)
assert.Equal(t, []int64{7}, rw)
assert.Equal(t, []int64{12}, ro, "the regular RO node is still one to move channels off")
})
t.Run("installed form", func(t *testing.T) {
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
ext.SetForm()
assert.Equal(t, []int64{11}, helper.GetRWNodesForChannels(replica))
rw, ro := helper.GetRWAndRONodesForChannels(replica)
assert.Equal(t, []int64{11}, rw)
assert.Equal(t, []int64{12}, ro)
})
}
internal/querycoordv2/checkers/channel_checker.go
+45 -6
View File
@@ -35,6 +35,7 @@ import (
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
"github.com/milvus-io/milvus/internal/util/streamingutil"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
@@ -109,17 +110,29 @@ func (c *ChannelChecker) Check(ctx context.Context) []task.Task {
replicas := c.meta.GetByCollection(ctx, cid)
hasTask := false
unplaced := false
for _, r := range replicas {
replicaTasks := c.checkReplica(ctx, r)
replicaTasks, replicaUnplaced := c.checkReplica(ctx, r)
if len(replicaTasks) > 0 {
hasTask = true
tasks = append(tasks, replicaTasks...)
}
if replicaUnplaced {
unplaced = true
}
}
// Only update version cache if no tasks were generated
// If tasks were generated, we need to re-check next time
if !hasTask {
// If tasks were generated, we need to re-check next time.
//
// Under a form, a replica that wanted a delegator and had no node
// to put it on is not converged either, and must not be cached as
// such: a form starts a replica's nodes on demand and spawns the
// replica before they register, and their arrival moves neither
// version the cache keys on. A stock binary keeps master's answer:
// the unplaced shape does not hold the cache open, so a replica
// waiting for a node there is cached like any other empty diff.
if versionCacheMayUpdate(hasTask, unplaced) {
c.updateVersionCache(cid, currentTargetVersion, currentDistVersion)
}
}
@@ -152,6 +165,17 @@ func (c *ChannelChecker) Check(ctx context.Context) []task.Task {
return tasks
}
// versionCacheMayUpdate reports whether a tick may cache a collection's target
// and channel-distribution versions as converged. A tick that produced tasks
// is not converged. Under a form, neither is a replica that wanted a delegator
// and had no node to put it on (unplaced): the form starts a replica's nodes
// on demand, so the replica is spawned before they register, and their arrival
// moves neither version the cache keys on. A stock binary keeps master's
// answer - the unplaced shape does not hold the cache open.
func versionCacheMayUpdate(hasTask, unplaced bool) bool {
return !hasTask && (!unplaced || !extension.FormInstalled())
}
// isCollectionSynced checks if target and dist versions are unchanged since last check
func (c *ChannelChecker) isCollectionSynced(collectionID int64, targetVersion, channelDistVersion int64) bool {
cache, ok := c.versionCache[collectionID]
@@ -186,11 +210,26 @@ func (c *ChannelChecker) cleanVersionCache(activeCollections []int64) {
}
}
func (c *ChannelChecker) checkReplica(ctx context.Context, replica *meta.Replica) []task.Task {
// checkReplica returns the tasks this replica needs, and whether it wanted a
// delegator placed but no load task came out. The case this exists for is the
// replica's read-write set being empty at this tick, so the assignment had
// nothing to place on; but the test is on the outcome, not the cause, so it
// also fires when the balancer produced no plan for some other reason, and
// that collection bypasses the version cache for as long as it lasts. That is
// accepted: the cost is one re-check per tick, and missing a replica that is
// waiting for a node would be the worse mistake.
//
// The second result exists for the version cache in Check: "no task" normally
// means converged, and caching it is what keeps the checker cheap. A replica
// waiting for a node is the opposite of converged, and the node it waits for
// joins the REPLICA, which moves neither the target version nor the channel
// distribution version - so a cached entry would never be invalidated.
func (c *ChannelChecker) checkReplica(ctx context.Context, replica *meta.Replica) ([]task.Task, bool) {
ret := make([]task.Task, 0)
lacks, redundancies := c.getDmChannelDiff(ctx, replica.GetCollectionID(), replica.GetID())
tasks := c.createChannelLoadTask(c.getTraceCtx(ctx, replica.GetCollectionID()), lacks, replica)
unplaced := len(lacks) > 0 && len(tasks) == 0
task.SetReason("lacks of channel", tasks...)
ret = append(ret, tasks...)
@@ -205,7 +244,7 @@ func (c *ChannelChecker) checkReplica(ctx context.Context, replica *meta.Replica
// All channel related tasks should be with high priority
task.SetPriority(task.TaskPriorityHigh, tasks...)
return ret
return ret, unplaced
}
// GetDmChannelDiff get channel diff between target and dist
@@ -285,7 +324,7 @@ func (c *ChannelChecker) createChannelLoadTask(ctx context.Context, channels []*
groupKeys := make([]string, 0)
for _, ch := range channels {
var rwNodes []int64
if streamingutil.IsStreamingServiceEnabled() {
if streamingutil.UseStreamingQueryNodeAsDelegator() {
rwNodes = replica.GetRWSQNodes()
} else {
if rwNodes = replica.GetChannelRWNodes(ch.GetChannelName()); len(rwNodes) == 0 {
internal/querycoordv2/checkers/channel_checker_form_test.go
+135
View File
@@ -0,0 +1,135 @@
// 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 checkers
import (
"context"
"os"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus/internal/coordinator/snmanager"
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_balancer"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/task"
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/balance"
"github.com/milvus-io/milvus/internal/util/streamingutil"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// channelOnAQueryCluster loads one channel of collection 1 into a replica that
// holds a regular query node and a streaming query node, with the WAL on a
// third node, and answers which node the delegator is grown on.
func (suite *ChannelCheckerTestSuite) channelOnAQueryCluster(regular, streamingQuery, walNode int64) int64 {
streamingutil.SetStreamingServiceEnabled()
suite.T().Cleanup(func() { os.Unsetenv(streamingutil.MilvusStreamingServiceEnabled) })
// Feed the pchannel's WAL assignment, otherwise GetWALLocated blocks forever.
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(suite.T())
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
cb(balancer.WatchChannelAssignmentsCallbackParam{
Version: typeutil.VersionInt64Pair{Global: 1, Local: 1},
CChannelAssignment: &streamingpb.CChannelAssignment{Meta: &streamingpb.CChannelMeta{Pchannel: "pchannel0"}},
Relations: []types.PChannelInfoAssigned{{
Channel: types.PChannelInfo{Name: "pchannel0", Term: 1},
Node: types.StreamingNodeInfo{ServerID: walNode, Address: "localhost:1"},
}},
})
<-ctx.Done()
return context.Cause(ctx)
}).Maybe()
streamingNodes := map[int64]*types.StreamingNodeInfoWithResourceGroup{}
for _, node := range []int64{walNode, streamingQuery} {
streamingNodes[node] = &types.StreamingNodeInfoWithResourceGroup{
StreamingNodeInfo: types.StreamingNodeInfo{ServerID: node, Address: "localhost:1"},
}
}
b.EXPECT().GetAllStreamingNodes(mock.Anything).Return(streamingNodes, nil).Maybe()
b.EXPECT().GetAvailableStreamingNodes(mock.Anything).Return(streamingNodes, nil).Maybe()
b.EXPECT().Close().Return().Maybe()
balance.Register(b)
ctx := context.Background()
checker := suite.checker
checker.scheduler.(*task.MockScheduler).EXPECT().GetChannelTaskDelta(mock.Anything, mock.Anything).Return(0).Maybe()
checker.meta.PutCollection(ctx, utils.CreateTestCollection(1, 1))
suite.meta.PutPartition(ctx, utils.CreateTestPartition(1, 1))
checker.meta.Put(ctx, meta.NewReplica(&querypb.Replica{
ID: 1, CollectionID: 1, ResourceGroup: meta.DefaultResourceGroupName,
Nodes: []int64{regular}, RwSqNodes: []int64{streamingQuery},
}))
suite.setNodeAvailable(regular, streamingQuery)
suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return(
[]*datapb.VchannelInfo{{CollectionID: 1, ChannelName: "pchannel0_1v0"}}, nil, nil)
checker.targetMgr.UpdateCollectionNextTarget(ctx, int64(1))
tasks := checker.Check(ctx)
suite.Require().Len(tasks, 1)
suite.Require().Len(tasks[0].Actions(), 1)
action := tasks[0].Actions()[0].(*task.ChannelAction)
suite.Equal(task.ActionTypeGrow, action.Type())
return action.Node()
}
// A form keeps its one streaming node for DDL and the write ahead log and
// serves queries from regular query nodes, one replica per query cluster. The
// delegator of a replica goes onto that replica's regular query node, which
// watches the channel and reads the WAL remotely - never onto the streaming
// query node the replica manager may still have handed the replica.
func (suite *ChannelCheckerTestSuite) TestUnderAFormADelegatorGoesToARegularQueryNode() {
ext.ResetForTest()
suite.T().Cleanup(ext.ResetForTest)
ext.SetForm()
suite.EqualValues(11, suite.channelOnAQueryCluster(11, 1051, 1007))
}
// A stock binary keeps milvus's placement: beside the WAL's readers, on the
// replica's streaming query node.
func (suite *ChannelCheckerTestSuite) TestOnAStockBinaryADelegatorGoesToTheStreamingQueryNode() {
ext.ResetForTest()
suite.T().Cleanup(ext.ResetForTest)
suite.EqualValues(1051, suite.channelOnAQueryCluster(11, 1051, 1007))
}
// The version-cache bypass for a replica that is unplaced is a form's: a stock
// binary caches that shape like any other tick that produced no task, as
// master did, while a form holds the cache open until a node arrives.
func (suite *ChannelCheckerTestSuite) TestTheUnplacedVersionCacheBypassIsAForms() {
ext.ResetForTest()
suite.T().Cleanup(ext.ResetForTest)
suite.True(versionCacheMayUpdate(false, false), "a converged tick is cached")
suite.True(versionCacheMayUpdate(false, true),
"a stock binary caches the unplaced shape, as master did")
suite.False(versionCacheMayUpdate(true, false), "a tick that produced tasks is never cached")
ext.SetForm()
suite.False(versionCacheMayUpdate(false, true),
"a form must not cache a replica that is waiting for a node")
suite.True(versionCacheMayUpdate(false, false), "a form still caches a converged tick")
}
internal/querycoordv2/checkers/channel_checker_unplaced_test.go
+108
View File
@@ -0,0 +1,108 @@
// 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 checkers
import (
"context"
"testing"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/task"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
)
func testReplica() *meta.Replica {
return meta.NewReplica(&querypb.Replica{ID: 1, CollectionID: 100})
}
// A replica that wants a delegator and has no node to put it on is not
// converged, and says so. The distinction matters to Check, which would
// otherwise cache it as settled: the node it waits for joins the REPLICA, and
// that moves neither version the cache keys on, so the entry would never be
// invalidated and the checker would stop looking - forever.
func TestCheckReplicaReportsADelegatorItCouldNotPlace(t *testing.T) {
c := &ChannelChecker{}
lacking := []*meta.DmChannel{{}}
diff := mockey.Mock((*ChannelChecker).getDmChannelDiff).Return(lacking, nil).Build()
defer diff.UnPatch()
loads := mockey.Mock((*ChannelChecker).createChannelLoadTask).Return([]task.Task{}).Build()
defer loads.UnPatch()
reduces := mockey.Mock((*ChannelChecker).createChannelReduceTasks).Return([]task.Task{}).Build()
defer reduces.UnPatch()
repeated := mockey.Mock((*ChannelChecker).findRepeatedChannels).Return(nil).Build()
defer repeated.UnPatch()
trace := mockey.Mock((*ChannelChecker).getTraceCtx).Return(context.Background()).Build()
defer trace.UnPatch()
tasks, unplaced := c.checkReplica(context.Background(), testReplica())
assert.Empty(t, tasks)
assert.True(t, unplaced, "a channel that was lacking and produced no task is a delegator with nowhere to go")
}
// A replica whose lack produced a task is progressing, not stuck: Check
// already refuses to cache it because a task was generated, and reporting
// "unplaced" as well would be redundant, not wrong - but the flag must mean
// what it says, so it is false here.
func TestCheckReplicaReportsNothingUnplacedWhenATaskWasMade(t *testing.T) {
c := &ChannelChecker{}
lacking := []*meta.DmChannel{{}}
diff := mockey.Mock((*ChannelChecker).getDmChannelDiff).Return(lacking, nil).Build()
defer diff.UnPatch()
loads := mockey.Mock((*ChannelChecker).createChannelLoadTask).Return([]task.Task{&task.ChannelTask{}}).Build()
defer loads.UnPatch()
reduces := mockey.Mock((*ChannelChecker).createChannelReduceTasks).Return([]task.Task{}).Build()
defer reduces.UnPatch()
repeated := mockey.Mock((*ChannelChecker).findRepeatedChannels).Return(nil).Build()
defer repeated.UnPatch()
trace := mockey.Mock((*ChannelChecker).getTraceCtx).Return(context.Background()).Build()
defer trace.UnPatch()
// The task this returns is a bare value, so the annotations milvus puts on
// a real one are stubbed out; what is under test is the flag, not them.
reason := mockey.Mock(task.SetReason).Return().Build()
defer reason.UnPatch()
priority := mockey.Mock(task.SetPriority).Return().Build()
defer priority.UnPatch()
_, unplaced := c.checkReplica(context.Background(), testReplica())
assert.False(t, unplaced)
}
// A converged replica lacks nothing, so there is nothing to fail to place.
// This is the case the version cache exists for, and it must stay cacheable.
func TestCheckReplicaReportsNothingUnplacedWhenNothingWasLacking(t *testing.T) {
c := &ChannelChecker{}
diff := mockey.Mock((*ChannelChecker).getDmChannelDiff).Return(nil, nil).Build()
defer diff.UnPatch()
loads := mockey.Mock((*ChannelChecker).createChannelLoadTask).Return([]task.Task{}).Build()
defer loads.UnPatch()
reduces := mockey.Mock((*ChannelChecker).createChannelReduceTasks).Return([]task.Task{}).Build()
defer reduces.UnPatch()
repeated := mockey.Mock((*ChannelChecker).findRepeatedChannels).Return(nil).Build()
defer repeated.UnPatch()
trace := mockey.Mock((*ChannelChecker).getTraceCtx).Return(context.Background()).Build()
defer trace.UnPatch()
tasks, unplaced := c.checkReplica(context.Background(), testReplica())
assert.Empty(t, tasks)
assert.False(t, unplaced)
}
internal/querycoordv2/checkers/leader_checker.go
+1 -1
View File
@@ -94,7 +94,7 @@ func (c *LeaderChecker) Check(ctx context.Context) []task.Task {
replicas := c.meta.GetByCollection(ctx, collectionID)
for _, replica := range replicas {
nodes := replica.GetRWNodes()
if streamingutil.IsStreamingServiceEnabled() {
if streamingutil.UseStreamingQueryNodeAsDelegator() {
nodes = replica.GetRWSQNodes()
}
for _, node := range nodes {
internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go
+69 -6
View File
@@ -23,6 +23,7 @@ import (
"github.com/milvus-io/milvus/internal/querycoordv2/job"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -46,7 +47,7 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadCollection(ctx conte
if err != nil {
return err
}
replicaNumber, resourceGroups, userSpecifiedReplicaMode, err := s.getLoadReplicaConfigForRequest(
replicaNumber, resourceGroups, scopedResourceGroups, userSpecifiedReplicaMode, err := s.getLoadReplicaConfigForRequest(
ctx,
req.GetReplicaNumber(),
req.GetResourceGroups(),
@@ -57,11 +58,51 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadCollection(ctx conte
}
currentLoadConfig := s.getCurrentLoadConfig(ctx, req.GetCollectionID())
// only check node number when the collection is not loaded
expectedReplicasNumber, err := utils.AssignReplica(ctx, s.meta, resourceGroups, replicaNumber, currentLoadConfig.Collection == nil)
// Node numbers are checked for a first load, and not for a config update
// on a loaded collection, which is master's rule and stays the stock
// binary's. With a form installed, a request that names resource groups
// on a loaded collection and asks one of them for more replicas than it
// holds is not a config update: it is a scoped expansion into those
// groups (see completePlacementForOutOfScopeResourceGroups), which
// places replicas exactly as a first load does, so it is admitted
// against the same bounds. Without the check, a group with no node would
// receive replicas that never get a delegator: the scoped task's clock
// would pause on an unknown progress forever and the group would report 0
// indefinitely.
// LoadPartitions has always passed true here for the same reason.
//
// Admission runs only for a request that ADDS replicas to a group it
// names. A scoped request that adds none is not admitted against
// anything: the same load re-sent, a shrink, or a request that changes
// only its load fields, partitions, priority or replica mode at the
// same counts - a config update, which master never admits either. None
// of these places a replica, and refusing one for a node that is
// restarting would turn an idempotent retry into a failure against a
// collection that is still serving.
//
// Whether the request is scoped is the decision getLoadReplicaConfigForRequest
// took, not a second reading of the request: under a cluster-level force
// override the groups the request named are discarded and the load states
// the whole placement, which is a config update like any other.
requestedReplicasNumber, err := utils.ReplicaNumberByResourceGroup(resourceGroups, replicaNumber)
if err != nil {
return err
}
checkNodeNum := currentLoadConfig.Collection == nil ||
(extension.FormInstalled() && len(scopedResourceGroups) > 0 && scopedLoadAddsReplicas(requestedReplicasNumber, currentLoadConfig))
expectedReplicasNumber, err := utils.AssignReplica(ctx, s.meta, resourceGroups, replicaNumber, checkNodeNum)
if err != nil {
return err
}
// With a form installed, a request that names resource groups speaks only
// for those and leaves the placement of the others alone; a request that
// names none - and every request on a stock binary - states the whole
// placement, which is the native behavior, and this returns what
// AssignReplica just produced. The scoping list comes from the same call
// that resolved the configuration, so both are decided from one reading of
// it.
expectedReplicasNumber = completePlacementForOutOfScopeResourceGroups(
ctx, req.GetCollectionID(), scopedResourceGroups, expectedReplicasNumber, currentLoadConfig)
alterLoadConfigReq := &job.AlterLoadConfigRequest{
Meta: s.meta,
CollectionInfo: coll,
@@ -89,7 +130,28 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadCollection(ctx conte
return err
}
func (s *Server) getLoadReplicaConfigForRequest(ctx context.Context, replicaNumber int32, resourceGroups []string, collectionID int64) (int32, []string, bool, error) {
// getLoadReplicaConfigForRequest resolves the replica configuration a load
// request runs with. It returns, in order: the replica number, the resource
// groups to assign replicas in (defaulted, which is what AssignReplica needs),
// the resource groups the REQUEST itself speaks for, and whether the caller
// named a replica number.
//
// The third return value is the scoping decision, and it is the raw list off
// the request, taken before the defaulting below rewrites an empty one to the
// default resource group. A request naming no group is a request about the
// collection, not a request about the default group: reading the defaulted list
// instead would turn every bare load into a request scoped to
// __default_resource_group, carry every other group's replicas through it, and
// so add a replica nobody asked for to a load_collection and refuse a
// load_partitions for "changing the replica number" of a group the caller never
// mentioned.
//
// Empty means "this request states the whole placement". A cluster-level force
// override states it by definition -- it replaces both the replica number and
// the resource groups of every load -- so it answers empty however the request
// itself was written, and it is read exactly once here, so the two answers
// cannot disagree with each other.
func (s *Server) getLoadReplicaConfigForRequest(ctx context.Context, replicaNumber int32, resourceGroups []string, collectionID int64) (int32, []string, []string, bool, error) {
// If force override is enabled with a complete cluster-level load config,
// new load requests are interpreted as cluster-managed even when the request
// carries explicit replica/RG parameters.
@@ -99,14 +161,15 @@ func (s *Server) getLoadReplicaConfigForRequest(ctx context.Context, replicaNumb
mlog.Int64("collectionID", collectionID),
mlog.Int32("replicaNumber", overrideReplicaNumber),
mlog.Strings("resourceGroups", overrideResourceGroups))
return overrideReplicaNumber, overrideResourceGroups, false, nil
return overrideReplicaNumber, overrideResourceGroups, nil, false, nil
}
scopedResourceGroups := resourceGroups
// If user specified the replica number in load request, load config changes
// won't be applied to the collection automatically.
userSpecifiedReplicaMode := replicaNumber > 0
replicaNumber, resourceGroups, err := s.getDefaultResourceGroupsAndReplicaNumber(ctx, replicaNumber, resourceGroups, collectionID)
return replicaNumber, resourceGroups, userSpecifiedReplicaMode, err
return replicaNumber, resourceGroups, scopedResourceGroups, userSpecifiedReplicaMode, err
}
func getClusterLevelLoadConfigForForceOverride() (int32, []string, bool) {
internal/querycoordv2/ddl_callbacks_alter_load_info_load_partitions.go
+31 -3
View File
@@ -21,6 +21,7 @@ import (
"github.com/milvus-io/milvus/internal/querycoordv2/job"
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
@@ -39,7 +40,7 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadPartitions(ctx conte
return err
}
replicaNumber, resourceGroups, userSpecifiedReplicaMode, err := s.getLoadReplicaConfigForRequest(
replicaNumber, resourceGroups, scopedResourceGroups, userSpecifiedReplicaMode, err := s.getLoadReplicaConfigForRequest(
ctx,
req.GetReplicaNumber(),
req.GetResourceGroups(),
@@ -49,12 +50,38 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadPartitions(ctx conte
return err
}
expectedReplicasNumber, err := utils.AssignReplica(ctx, s.meta, resourceGroups, replicaNumber, true)
currentLoadConfig := s.getCurrentLoadConfig(ctx, req.GetCollectionID())
// Node numbers are checked on every load_partitions, which is master's
// rule and stays the stock binary's. With a form, a scoped request that
// adds no replica to the groups it names - the same load re-sent while
// a node of the group restarts, or one that changes only its partitions
// at the same counts - places nothing and is not admitted against
// anything, exactly as in the LoadCollection callback: the per-group rule
// would otherwise refuse a re-send for a group whose query node is
// restarting, against a collection that is still serving.
requestedReplicasNumber, err := utils.ReplicaNumberByResourceGroup(resourceGroups, replicaNumber)
if err != nil {
return err
}
checkNodeNum := !extension.FormInstalled() || len(scopedResourceGroups) == 0 || currentLoadConfig.Collection == nil ||
scopedLoadAddsReplicas(requestedReplicasNumber, currentLoadConfig)
expectedReplicasNumber, err := utils.AssignReplica(ctx, s.meta, resourceGroups, replicaNumber, checkNodeNum)
if err != nil {
return err
}
currentLoadConfig := s.getCurrentLoadConfig(ctx, req.GetCollectionID())
// With a form installed, a request that names resource groups speaks only
// for those and leaves the placement of the others alone - same seam, same
// reason as the LoadCollection callback: without it a scoped
// LoadPartitions would be read as the whole placement and refused for
// "changing" the replica number of resource groups it never mentioned. A
// request that names none - and every request on a stock binary - states
// the whole placement and this returns what AssignReplica just produced;
// reading the defaulted list here instead would refuse a bare
// load_partitions on a collection loaded in another resource group for
// exactly that reason.
expectedReplicasNumber = completePlacementForOutOfScopeResourceGroups(
ctx, req.GetCollectionID(), scopedResourceGroups, expectedReplicasNumber, currentLoadConfig)
partitionIDsSet := typeutil.NewSet(currentLoadConfig.GetPartitionIDs()...)
// add new incoming partitionIDs.
for _, partition := range req.PartitionIDs {
@@ -72,6 +99,7 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadPartitions(ctx conte
ExpectedPriority: req.GetPriority(),
ExpectedUserSpecifiedReplicaMode: userSpecifiedReplicaMode,
},
ScopedResourceGroups: scopedResourceGroups,
}
if err := alterLoadConfigReq.CheckIfLoadPartitionsExecutable(); err != nil {
return err
internal/querycoordv2/handlers.go
+27 -1
View File
@@ -47,7 +47,33 @@ import (
// may come from different replica group. We only need these shards to form a replica that serves query
// requests.
func (s *Server) checkAnyReplicaAvailable(collectionID int64) bool {
for _, replica := range s.meta.GetByCollection(s.ctx, collectionID) {
return s.anyReplicaAvailable(s.meta.GetByCollection(s.ctx, collectionID))
}
// checkAnyReplicaAvailableInResourceGroup is the answer a ShowLoadCollections
// scoped to a resource group gives: whether THAT group can serve the
// collection, which is the group's shard-leader readiness - every shard of
// the collection has a serviceable leader in the group's replicas, on a node
// the coordinator knows (utils.ShardLeaderReadinessByResourceGroup, the same
// verdict the scoped load waits for). It is not the collection-wide rule
// restricted to one group: that rule reads a replica with no read-only node
// as available, so a replica just spawned into the group, holding nothing,
// would read true beside a progress of 0.
func (s *Server) checkAnyReplicaAvailableInResourceGroup(ctx context.Context, collectionID int64, rgName string) bool {
readiness, err := utils.ShardLeaderReadinessByResourceGroup(ctx, s.meta, s.targetMgr, s.dist, s.nodeMgr, collectionID, rgName)
if err != nil {
mlog.Warn(ctx, "the resource group's shard-leader readiness cannot be answered, reporting it as not serving",
mlog.Int64("collectionID", collectionID), mlog.String("resourceGroup", rgName), mlog.Err(err))
return false
}
return readiness.Ready
}
// anyReplicaAvailable is the collection-wide rule: a replica is available when
// every one of its read-only nodes is still known to the node manager, and
// the answer is yes as soon as one replica is.
func (s *Server) anyReplicaAvailable(replicas []*meta.Replica) bool {
for _, replica := range replicas {
isAvailable := true
for _, node := range replica.GetRONodes() {
if s.nodeMgr.Get(node) == nil {
internal/querycoordv2/job/job_load.go
+206 -9
View File
@@ -19,6 +19,7 @@ package job
import (
"context"
"fmt"
"maps"
"time"
"github.com/cockroachdb/errors"
@@ -33,6 +34,7 @@ import (
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
"github.com/milvus-io/milvus/internal/util/proxyutil"
"github.com/milvus-io/milvus/pkg/v3/eventlog"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
@@ -113,6 +115,10 @@ func (job *LoadCollectionJob) Execute() error {
mlog.Int("localReplicaCount", len(localReplicas)))
}
// Classify the request before spawn, because spawn is what mutates the
// replica meta both of the following read.
incrementalExpansion := job.isIncrementalExpansion(req, replicas)
// 2. create replica if not exist (may also remove redundant replicas)
if _, err := utils.SpawnReplicasWithReplicaConfig(job.ctx, job.meta, meta.SpawnWithReplicaConfigParams{
CollectionID: req.GetCollectionId(),
@@ -131,15 +137,83 @@ func (job *LoadCollectionJob) Execute() error {
}
// 3. put load info meta
fieldIndexIDs := make(map[int64]int64, len(req.GetLoadFields()))
fieldIDs := make([]int64, 0, len(req.GetLoadFields()))
for _, loadField := range req.GetLoadFields() {
if loadField.GetIndexId() != 0 {
fieldIndexIDs[loadField.GetFieldId()] = loadField.GetIndexId()
}
fieldIDs = append(fieldIDs, loadField.GetFieldId())
}
fieldIndexIDs, fieldIDs := requestedLoadFields(req)
replicaNumber := int32(len(replicas))
ctx, sp := otel.Tracer(typeutil.QueryCoordRole).Start(job.ctx, "LoadCollection", trace.WithNewRoot())
// 3.1 incremental resource group expansion: this request adds resource
// groups to a collection that is already loaded and changes nothing else,
// so the collection and partition meta it would write is the meta that is
// already there -- except for Status and LoadPercentage, which the write
// would reset to Loading/0 and take the resource groups that are serving
// right now down with them. Skip the overwrite; the added replicas are
// loaded by the checkers, as an added replica is on any other path.
//
// Everything on this path is idempotent, and it has to be: the spawn
// above is persisted before any of it, and if this process fails or dies
// between the two the message is replayed. The replay finds the added
// replicas already in meta, the predicate still says expansion (an
// identical replica set on a Loaded collection is exactly a replay), and
// each step below either writes the value that is already there or
// re-pulls a target that has not changed. Taking the native path on a
// replay instead would write the collection back to Loading/0 and make
// the group that is serving unreadable until the added one finishes - the
// incident this path exists to remove.
if incrementalExpansion {
// Nothing keeps this span on this path: the collection holds the
// LoadSpan of the load that created it, and that span is what
// UpdateCollectionLoadPercent ends. End this one here rather than leak
// it.
defer sp.End()
mlog.Info(job.ctx, "incremental resource group expansion, keeping loaded collection meta",
mlog.Int64("collectionID", req.GetCollectionId()),
mlog.Int32("replicaNumber", replicaNumber))
// No observer task is registered for the added groups, deliberately.
// The collection stays Loaded, so there is no status to write back and
// no aggregate to complete; how far a group has loaded and whether it
// can serve are answered from the live target and distribution
// (utils.LoadPercentageByResourceGroup and
// utils.ShardLeaderReadinessByResourceGroup), and the checkers load the
// added replicas on their own schedule. This is what upstream's
// UpdateLoadConfig does when it adds a replica to a loaded collection:
// spawn, write the count, pull the target, and leave the loading to the
// checkers. A group whose replicas never load is the caller's to
// release, exactly as an added replica that never loads is upstream.
// The replica count is the one property of the collection this request
// legitimately changes; the predicate has established that everything
// else is unchanged. UpdateReplicaNumber writes it to the collection and
// its partitions without touching their Status or LoadPercentage, so
// the already-serving resource groups stay Loaded at 100%. Leaving it
// stale instead would misreport the collection's replica count to every
// later UpdateLoadConfig, which reads it as the replica number to keep.
//
// The number written is the replicas that exist, counted inside the
// collection manager's critical section (SyncReplicaNumber), rather
// than the request's own count: they are the same unless a release of
// one of the groups runs beside this ack, and then the request's count
// written last would stand one above the replicas there are.
userSpecifiedReplicaMode := req.GetUserSpecifiedReplicaMode()
if _, err := job.meta.SyncReplicaNumber(job.ctx, req.GetCollectionId(), func() int32 {
return int32(len(job.meta.GetByCollection(job.ctx, req.GetCollectionId())))
}, &userSpecifiedReplicaMode); err != nil {
msg := "failed to update replica number"
mlog.Warn(job.ctx, msg, mlog.Err(err))
return merr.Wrapf(err, "%s", msg)
}
// The target is unchanged -- same partitions, same segments -- but the
// pull is cheap and keeps this path's ordering identical to the one
// below, where the observer only ever sees a populated next target.
if _, err = job.targetObserver.UpdateNextTarget(req.GetCollectionId()); err != nil {
return err
}
return nil
}
partitions := lo.Map(req.GetPartitionIds(), func(partID int64, _ int) *meta.Partition {
return &meta.Partition{
PartitionLoadInfo: &querypb.PartitionLoadInfo{
@@ -153,7 +227,6 @@ func (job *LoadCollectionJob) Execute() error {
}
})
ctx, sp := otel.Tracer(typeutil.QueryCoordRole).Start(job.ctx, "LoadCollection", trace.WithNewRoot())
collection := &meta.Collection{
CollectionLoadInfo: &querypb.CollectionLoadInfo{
CollectionID: req.GetCollectionId(),
@@ -223,6 +296,130 @@ func (job *LoadCollectionJob) Execute() error {
return nil
}
// requestedLoadFields splits the request's load-field configs into the
// field-to-index map and the plain list of field IDs that meta.Collection
// stores, so that the predicate below compares exactly the values the
// collection would have been overwritten with.
func requestedLoadFields(req *messagespb.AlterLoadConfigMessageHeader) (map[int64]int64, []int64) {
fieldIndexIDs := make(map[int64]int64, len(req.GetLoadFields()))
fieldIDs := make([]int64, 0, len(req.GetLoadFields()))
for _, loadField := range req.GetLoadFields() {
if loadField.GetIndexId() != 0 {
fieldIndexIDs[loadField.GetFieldId()] = loadField.GetIndexId()
}
fieldIDs = append(fieldIDs, loadField.GetFieldId())
}
return fieldIndexIDs, fieldIDs
}
// isIncrementalExpansion reports whether req asks for nothing but additional
// resource groups on a collection that is already loaded. That is the one shape
// of request whose collection/partition meta write is pure loss: the write
// stores values identical to the stored ones except for Status and
// LoadPercentage, which it resets to Loading and 0, dropping resource groups
// that are serving queries right now out of Loaded until the observer has
// walked them back up to 100%.
//
// Every other request keeps the overwrite. The predicate is therefore
// deliberately conservative -- it is false unless every part of the request
// matches what is stored -- and in particular it is false for a first load, for
// a reload, and for a plain replica-number change, so a deployment that never
// loads one collection into several resource groups never leaves the path this
// job has always taken.
//
// The legs, and what each one protects:
//
// - collection exists and is Loaded: there is serving state to protect, and
// an unfinished load must keep its Loading status and its observer task.
// - same DbID, same LoadType, same partition set, same load fields, same
// field-to-index map: everything the skipped write would have stored is
// already stored, so skipping it loses nothing. The partition leg also
// covers the partition-release branch below, which the fast path skips.
// - every existing replica appears in the new set, in the same resource
// group: no replica is being released or moved, so no resource group is
// losing state that the collection meta still claims it has.
//
// The whole path is a form's: it answers false on a stock binary
// (extension.FormInstalled), which keeps master's behavior exactly - every load
// of an already-loaded collection, a two-group LoadCollection included, writes
// the collection back to Loading/0 and the collection-wide observer walks it up
// again. The deployment shape that loads one collection into several resource
// groups independently, and cannot afford that reset, is the one a form
// builds.
//
// The new set need not be LARGER. A request whose replica set is exactly the
// stored one, on a Loaded collection, is this job's own message replayed: the
// spawn persisted the added replicas, the process failed before the rest of
// the fast path ran, and the message came back. Nothing else produces that
// shape - a request that changes nothing is dropped before it is broadcast
// (GenerateAlterLoadConfigMessage) - and taking the overwrite on it would
// write a serving collection back to Loading/0, which is the one outcome this
// predicate exists to prevent. Deciding from the requested set against what is
// stored, rather than from the count growing, is what makes the replay land on
// the same path as the original.
func (job *LoadCollectionJob) isIncrementalExpansion(req *messagespb.AlterLoadConfigMessageHeader, newReplicas []*messagespb.LoadReplicaConfig) bool {
if !extension.FormInstalled() {
return false
}
existing := job.meta.GetCollection(job.ctx, req.GetCollectionId())
if existing == nil || existing.GetStatus() != querypb.LoadStatus_Loaded {
return false
}
if existing.GetDbID() != req.GetDbId() {
return false
}
if existing.GetLoadType() != querypb.LoadType_LoadCollection {
return false
}
incomingPartitions := typeutil.NewSet(req.GetPartitionIds()...)
currentPartitions := job.meta.GetPartitionsByCollection(job.ctx, req.GetCollectionId())
if len(currentPartitions) != incomingPartitions.Len() {
return false
}
for _, partition := range currentPartitions {
if !incomingPartitions.Contain(partition.GetPartitionID()) {
return false
}
}
fieldIndexIDs, fieldIDs := requestedLoadFields(req)
if !maps.Equal(typeutil.NewSet(fieldIDs...), typeutil.NewSet(existing.GetLoadFields()...)) {
return false
}
if !maps.Equal(fieldIndexIDs, existing.GetFieldIndexID()) {
return false
}
existingReplicas := job.meta.GetByCollection(job.ctx, req.GetCollectionId())
if len(newReplicas) < len(existingReplicas) {
return false // a replica is being released
}
existingRGByReplica := make(map[int64]string, len(existingReplicas))
for _, replica := range existingReplicas {
existingRGByReplica[replica.GetID()] = replica.GetResourceGroup()
}
seen := typeutil.NewSet[int64]()
for _, replica := range newReplicas {
seen.Insert(replica.GetReplicaId())
rgName, isExisting := existingRGByReplica[replica.GetReplicaId()]
if isExisting {
if rgName != replica.GetResourceGroupName() {
return false // an existing replica is being moved to another resource group
}
continue
}
}
for _, replica := range existingReplicas {
if !seen.Contain(replica.GetID()) {
return false // an existing replica is being released
}
}
return true
}
// getLocalReplicaConfig reads the local cluster-level replica config and generates LoadReplicaConfig entries.
// It uses generateReplicas to ensure idempotency on WAL replay by reusing existing replicas from meta.
// If local config is not set, defaults to 1 replica in __default_resource_group.
internal/querycoordv2/job/job_load_test.go
+626
View File
@@ -19,14 +19,23 @@ package job
import (
"context"
"testing"
"time"
"github.com/bytedance/mockey"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/rgpb"
"github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/querycoordv2/checkers"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/observers"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
@@ -421,3 +430,620 @@ func (suite *LoadCollectionJobSuite) TestGetLocalReplicaConfig_AllocIDError() {
func TestLoadCollectionJob(t *testing.T) {
suite.Run(t, new(LoadCollectionJobSuite))
}
const (
expansionCollectionID = int64(5000)
expansionPartitionID = int64(50)
expansionChannel = "5000-dmc0"
expansionDbID = int64(7)
rgA = "rg_a"
rgB = "rg_b"
rgC = "rg_c"
)
// observedLoadTask is one registration the job handed to the CollectionObserver.
type observedLoadTask struct {
collectionID int64
partitionIDs []int64
}
// IncrementalExpansionSuite exercises LoadCollectionJob.Execute against a real
// meta store, so the assertions are about the state the job leaves behind
// rather than about which collaborator was called.
type IncrementalExpansionSuite struct {
suite.Suite
ctx context.Context
catalog *mocks.QueryCoordCatalog
meta *meta.Meta
dist *meta.DistributionManager
nodeMgr *session.NodeManager
broker *meta.MockBroker
nextID int64
}
func (suite *IncrementalExpansionSuite) SetupSuite() {
paramtable.Init()
}
// installForm makes this test's binary one with a form installed; the stock
// cases below uninstall it again.
func (suite *IncrementalExpansionSuite) installForm() {
ext.ResetForTest()
suite.T().Cleanup(ext.ResetForTest)
ext.SetForm()
}
func (suite *IncrementalExpansionSuite) SetupTest() {
suite.ctx = context.Background()
meta.GlobalFailedLoadCache = meta.NewFailedLoadCache()
suite.installForm()
suite.catalog = mocks.NewQueryCoordCatalog(suite.T())
// The collection under test always carries exactly one partition, so a
// SaveCollection carries either zero or one partition load info.
suite.catalog.On("SaveCollection", mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("SaveCollection", mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("SavePartition", mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("SaveReplica", mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("SaveReplica", mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("SaveReplica", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
// Spawning replicas batches its writes through the variadic Update. The
// action count depends on how many replicas the request reconciles, so
// every arity these cases reach is stubbed.
suite.catalog.On("Update", mock.Anything).Return(nil).Maybe()
suite.catalog.On("Update", mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("Update", mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("Update", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("SaveResourceGroup", mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("ReleaseReplica", mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("ReleasePartition", mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
suite.catalog.On("ReleaseCollection", mock.Anything, mock.Anything).Return(nil).Maybe()
suite.nextID = 9000
suite.nodeMgr = session.NewNodeManager()
suite.meta = meta.NewMeta(func() (int64, error) {
suite.nextID++
return suite.nextID, nil
}, suite.catalog, suite.nodeMgr)
suite.dist = meta.NewDistributionManager(suite.nodeMgr)
suite.broker = meta.NewMockBroker(suite.T())
suite.broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).
Return(&milvuspb.DescribeCollectionResponse{
CollectionID: expansionCollectionID,
VirtualChannelNames: []string{expansionChannel},
}, nil).Maybe()
}
// seedLoadedCollection puts a collection that has finished loading into
// resource group rg, with one replica holding one node.
func (suite *IncrementalExpansionSuite) seedLoadedCollection(replicaID int64, rg string, nodes ...int64) {
suite.seedCollection(querypb.LoadStatus_Loaded, querypb.LoadType_LoadCollection, expansionDbID)
suite.seedReplica(replicaID, rg, nodes...)
}
func (suite *IncrementalExpansionSuite) seedCollection(status querypb.LoadStatus, loadType querypb.LoadType, dbID int64) {
err := suite.meta.PutCollection(suite.ctx, &meta.Collection{
CollectionLoadInfo: &querypb.CollectionLoadInfo{
CollectionID: expansionCollectionID,
DbID: dbID,
ReplicaNumber: 1,
Status: status,
LoadType: loadType,
LoadFields: []int64{100, 101},
FieldIndexID: map[int64]int64{100: 1000},
},
LoadPercentage: 100,
CreatedAt: time.Now().Add(-3 * time.Hour),
}, &meta.Partition{
PartitionLoadInfo: &querypb.PartitionLoadInfo{
CollectionID: expansionCollectionID,
PartitionID: expansionPartitionID,
ReplicaNumber: 1,
Status: status,
FieldIndexID: map[int64]int64{100: 1000},
},
LoadPercentage: 100,
CreatedAt: time.Now().Add(-3 * time.Hour),
})
suite.Require().NoError(err)
}
func (suite *IncrementalExpansionSuite) seedReplica(replicaID int64, rg string, nodes ...int64) {
err := suite.meta.Put(suite.ctx, meta.NewReplica(&querypb.Replica{
ID: replicaID,
CollectionID: expansionCollectionID,
ResourceGroup: rg,
Nodes: nodes,
}))
suite.Require().NoError(err)
}
func replicaConfig(replicaID int64, rg string) *messagespb.LoadReplicaConfig {
return &messagespb.LoadReplicaConfig{ReplicaId: replicaID, ResourceGroupName: rg}
}
// buildExpansionRequest builds a request that matches what seedCollection
// stored, apart from the replicas handed in.
func (suite *IncrementalExpansionSuite) buildExpansionRequest(replicas ...*messagespb.LoadReplicaConfig) message.BroadcastResultAlterLoadConfigMessageV2 {
return suite.buildRequest(expansionDbID, []int64{expansionPartitionID},
[]*messagespb.LoadFieldConfig{{FieldId: 100, IndexId: 1000}, {FieldId: 101}}, replicas)
}
func (suite *IncrementalExpansionSuite) buildRequest(
dbID int64,
partitionIDs []int64,
loadFields []*messagespb.LoadFieldConfig,
replicas []*messagespb.LoadReplicaConfig,
) message.BroadcastResultAlterLoadConfigMessageV2 {
controlChannel := "_ctrl_channel"
broadcastMsg := message.NewAlterLoadConfigMessageBuilderV2().
WithHeader(&messagespb.AlterLoadConfigMessageHeader{
DbId: dbID,
CollectionId: expansionCollectionID,
PartitionIds: partitionIDs,
LoadFields: loadFields,
Replicas: replicas,
}).
WithBody(&messagespb.AlterLoadConfigMessageBody{}).
WithBroadcast([]string{controlChannel}).
MustBuildBroadcast()
return message.BroadcastResultAlterLoadConfigMessageV2{
Message: message.MustAsBroadcastAlterLoadConfigMessageV2(broadcastMsg),
Results: map[string]*message.AppendResult{controlChannel: {}},
}
}
// runJob executes the job with the target observer stubbed out (it needs a
// running goroutine to answer) and records what reached the collection
// observer, plus how many times the collection meta was overwritten. Both
// PutCollection and LoadPartitions keep running for real.
func (suite *IncrementalExpansionSuite) runJob(result message.BroadcastResultAlterLoadConfigMessageV2) (int, []observedLoadTask, error) {
putCollectionCalls := 0
var putOrigin func(*meta.CollectionManager, context.Context, *meta.Collection, ...*meta.Partition) error
putMock := mockey.Mock((*meta.CollectionManager).PutCollection).
To(func(cm *meta.CollectionManager, ctx context.Context, collection *meta.Collection, partitions ...*meta.Partition) error {
putCollectionCalls++
return putOrigin(cm, ctx, collection, partitions...)
}).Origin(&putOrigin).Build()
defer putMock.UnPatch()
tasks := make([]observedLoadTask, 0)
loadMock := mockey.Mock((*observers.CollectionObserver).LoadPartitions).
To(func(ob *observers.CollectionObserver, ctx context.Context, collectionID int64, partitionIDs []int64) {
tasks = append(tasks, observedLoadTask{collectionID, partitionIDs})
}).Build()
defer loadMock.UnPatch()
targetMock := mockey.Mock((*observers.TargetObserver).UpdateNextTarget).
To(func(ob *observers.TargetObserver, collectionID int64) (chan struct{}, error) {
return nil, nil
}).Build()
defer targetMock.UnPatch()
job := NewLoadCollectionJob(suite.ctx, result, suite.dist, suite.meta, suite.broker,
nil, &observers.TargetObserver{}, &observers.CollectionObserver{}, nil, nil, nil)
err := job.Execute()
return putCollectionCalls, tasks, err
}
// TestFirstLoadOverwritesMetaAndRegistersUnscopedTask pins the upstream path:
// a load of a collection that is not loaded yet writes the collection meta and
// hands the observer a task that names no resource group.
func (suite *IncrementalExpansionSuite) TestFirstLoadOverwritesMetaAndRegistersATask() {
putCalls, tasks, err := suite.runJob(suite.buildExpansionRequest(replicaConfig(1, rgA)))
suite.NoError(err)
suite.Equal(1, putCalls, "a first load must store the collection meta")
collection := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
suite.Require().NotNil(collection)
suite.Equal(querypb.LoadStatus_Loading, collection.GetStatus())
suite.EqualValues(1, collection.GetReplicaNumber())
suite.EqualValues(0, collection.LoadPercentage)
suite.Require().Len(tasks, 1)
suite.Equal(expansionCollectionID, tasks[0].collectionID)
suite.Equal([]int64{expansionPartitionID}, tasks[0].partitionIDs)
}
// TestAnIdenticalReplicaSetOnALoadedCollectionIsAReplay: a message whose
// replica set is exactly what meta holds, on a Loaded collection, can only be
// this job's own message coming back after the spawn persisted and the rest
// did not run - a request that changes nothing is never broadcast. It must
// leave the serving collection alone: no overwrite, no new task, and the
// replica number written to what the message says.
func (suite *IncrementalExpansionSuite) TestAnIdenticalReplicaSetOnALoadedCollectionIsAReplay() {
suite.seedLoadedCollection(1, rgA, 1)
putCalls, tasks, err := suite.runJob(suite.buildExpansionRequest(replicaConfig(1, rgA)))
suite.NoError(err)
suite.Equal(0, putCalls, "a replay must not write the collection back to Loading")
collection := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
suite.Equal(querypb.LoadStatus_Loaded, collection.GetStatus())
suite.EqualValues(100, collection.LoadPercentage)
suite.EqualValues(1, collection.GetReplicaNumber())
suite.Empty(tasks, "no resource group was added, so no task is registered")
}
// TestAReplayedExpansionKeepsTheCollectionLoaded is the replay-safety case.
// The spawn persists rgB's replica; the process then fails before the replica
// number is written (a crash between the two looks the same to the message,
// which is acknowledged only once Execute returns nil). The message is
// executed again. At no point may the serving rgA see its collection leave
// Loaded, and rgB must end up with exactly one observer task.
func (suite *IncrementalExpansionSuite) TestAReplayedExpansionKeepsTheCollectionLoaded() {
suite.seedLoadedCollection(1, rgA, 1)
message := suite.buildExpansionRequest(replicaConfig(1, rgA), replicaConfig(2, rgB))
// The first attempt dies at the replica number write, after the spawn.
crashed := errors.New("querycoord went away")
var origin func(*meta.CollectionManager, context.Context, int64, func() int32, *bool) (int32, error)
attempts := 0
crash := mockey.Mock((*meta.CollectionManager).SyncReplicaNumber).
To(func(cm *meta.CollectionManager, ctx context.Context, collectionID int64, count func() int32, userSpecified *bool) (int32, error) {
attempts++
if attempts == 1 {
return 0, crashed
}
return origin(cm, ctx, collectionID, count, userSpecified)
}).Origin(&origin).Build()
defer crash.UnPatch()
putCalls, tasks, err := suite.runJob(message)
suite.ErrorIs(err, crashed)
suite.Equal(0, putCalls)
suite.Len(suite.meta.GetByCollection(suite.ctx, expansionCollectionID), 2, "the spawn persisted before the failure")
suite.Equal(querypb.LoadStatus_Loaded, suite.meta.GetCollection(suite.ctx, expansionCollectionID).GetStatus(),
"the failed attempt must not have touched the serving state")
suite.Empty(tasks, "an expansion registers no observer task, on the first attempt or on a replay")
// The replay.
putCalls, tasks, err = suite.runJob(message)
suite.NoError(err)
suite.Equal(0, putCalls, "a replay must not write the collection back to Loading")
collection := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
suite.Equal(querypb.LoadStatus_Loaded, collection.GetStatus(), "rgA never stops serving")
suite.EqualValues(100, collection.LoadPercentage)
suite.EqualValues(2, collection.GetReplicaNumber(), "the replay completes what the first attempt could not")
partition := suite.meta.GetPartition(suite.ctx, expansionPartitionID)
suite.Equal(querypb.LoadStatus_Loaded, partition.GetStatus())
suite.EqualValues(2, partition.GetReplicaNumber())
suite.Len(suite.meta.GetByCollection(suite.ctx, expansionCollectionID), 2)
suite.Empty(tasks, "the replay registers no second task for rgB")
}
// TestReplicaNumberIncreaseInSameResourceGroupKeepsTheCollectionLoaded: an
// added replica that lands in a resource group that already holds the
// collection is an expansion like any other - the checkers load it, exactly as
// they load a replica in a group the request newly names. The fast path keeps
// the collection serving; there is no longer an overwrite leg for this shape
// to force.
func (suite *IncrementalExpansionSuite) TestReplicaNumberIncreaseInSameResourceGroupKeepsTheCollectionLoaded() {
suite.seedLoadedCollection(1, rgA, 1)
putCalls, tasks, err := suite.runJob(suite.buildExpansionRequest(
replicaConfig(1, rgA), replicaConfig(2, rgA)))
suite.NoError(err)
suite.Equal(0, putCalls, "adding a replica keeps the collection meta as it is")
collection := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
suite.Require().NotNil(collection)
suite.Equal(querypb.LoadStatus_Loaded, collection.GetStatus(), "the collection stays serving")
suite.EqualValues(100, collection.LoadPercentage)
suite.EqualValues(2, collection.GetReplicaNumber())
suite.Len(suite.meta.GetByCollection(suite.ctx, expansionCollectionID), 2)
suite.Empty(tasks, "the checkers load the added replica, as they do an added replica in a new group")
}
// TestReplicaNumberDecreaseOverwritesMeta covers the other direction: the
// request drops a replica, so it is not an expansion.
func (suite *IncrementalExpansionSuite) TestReplicaNumberDecreaseOverwritesMeta() {
suite.seedLoadedCollection(1, rgA, 1)
suite.seedReplica(2, rgB, 2)
putCalls, tasks, err := suite.runJob(suite.buildExpansionRequest(replicaConfig(1, rgA)))
suite.NoError(err)
suite.Equal(1, putCalls, "a replica-number decrease must still store the collection meta")
suite.Require().Len(tasks, 1, "a load that resets the collection is watched to completion")
}
// TestAStockBinaryResetsTheCollectionOnAnExpansion pins master's behavior on
// the very request the fast path exists for: with no form installed, a load
// that adds rgB to a collection loaded in rgA writes the collection back to
// Loading/0 and registers the collection-wide task, exactly as every load of
// an already-loaded collection always has.
func (suite *IncrementalExpansionSuite) TestAStockBinaryResetsTheCollectionOnAnExpansion() {
ext.ResetForTest()
suite.seedLoadedCollection(1, rgA, 1)
putCalls, tasks, err := suite.runJob(suite.buildExpansionRequest(
replicaConfig(1, rgA), replicaConfig(2, rgB)))
suite.NoError(err)
suite.Equal(1, putCalls, "a stock binary stores the collection meta on every load")
collection := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
suite.Require().NotNil(collection)
suite.Equal(querypb.LoadStatus_Loading, collection.GetStatus(), "master resets the collection to Loading")
suite.EqualValues(0, collection.LoadPercentage)
suite.EqualValues(2, collection.GetReplicaNumber())
suite.Require().Len(tasks, 1, "and watches it to completion")
}
// TestIncrementalExpansionKeepsLoadedResourceGroupIntact is the incident case:
// a collection loaded into rgA hours ago is loaded again to also cover rgB.
func (suite *IncrementalExpansionSuite) TestIncrementalExpansionKeepsLoadedResourceGroupIntact() {
suite.seedLoadedCollection(1, rgA, 1)
before := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
beforePartition := suite.meta.GetPartition(suite.ctx, expansionPartitionID)
putCalls, tasks, err := suite.runJob(suite.buildExpansionRequest(
replicaConfig(1, rgA), replicaConfig(2, rgB)))
suite.NoError(err)
suite.Equal(0, putCalls, "an incremental resource group expansion must not overwrite the collection meta")
collection := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
suite.Require().NotNil(collection)
suite.Equal(querypb.LoadStatus_Loaded, collection.GetStatus(), "rgA must stay loaded")
suite.EqualValues(100, collection.LoadPercentage, "rgA's load percentage must survive")
suite.Equal(before.CreatedAt, collection.CreatedAt, "the collection must not be recreated")
suite.EqualValues(2, collection.GetReplicaNumber(), "the replica count must follow the expansion")
partition := suite.meta.GetPartition(suite.ctx, expansionPartitionID)
suite.Require().NotNil(partition)
suite.Equal(querypb.LoadStatus_Loaded, partition.GetStatus())
suite.EqualValues(100, partition.LoadPercentage)
suite.Equal(beforePartition.CreatedAt, partition.CreatedAt)
suite.EqualValues(2, partition.GetReplicaNumber())
suite.Empty(tasks, "the collection never left Loaded, so there is nothing for an observer task to complete")
}
// TestIncrementalExpansionRegistersNoTask: adding two resource groups to a
// loaded collection adds their replicas and nothing else. Upstream's
// UpdateLoadConfig does the same when it adds a replica to a loaded
// collection, and the checkers are what load them.
func (suite *IncrementalExpansionSuite) TestIncrementalExpansionRegistersNoTask() {
suite.seedLoadedCollection(1, rgA, 1)
putCalls, tasks, err := suite.runJob(suite.buildExpansionRequest(
replicaConfig(1, rgA), replicaConfig(2, rgB), replicaConfig(3, rgC)))
suite.NoError(err)
suite.Equal(0, putCalls)
suite.Empty(tasks,
"an expansion leaves the collection Loaded, so there is no load status for an observer task to complete; "+
"the checkers load the added replicas and the per-resource-group figures report them")
collection := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
suite.Require().NotNil(collection)
suite.Equal(querypb.LoadStatus_Loaded, collection.GetStatus())
suite.EqualValues(3, collection.GetReplicaNumber(), "the added replicas are counted")
suite.Len(suite.meta.GetByCollection(suite.ctx, expansionCollectionID), 3)
}
// newPredicateJob builds a job over the seeded meta without executing it, so
// each leg of the predicate can be exercised on its own.
func (suite *IncrementalExpansionSuite) newPredicateJob(result message.BroadcastResultAlterLoadConfigMessageV2) *LoadCollectionJob {
return NewLoadCollectionJob(suite.ctx, result, suite.dist, suite.meta, suite.broker,
nil, &observers.TargetObserver{}, &observers.CollectionObserver{}, nil, nil, nil)
}
func (suite *IncrementalExpansionSuite) TestIsIncrementalExpansionLegs() {
sameFields := []*messagespb.LoadFieldConfig{{FieldId: 100, IndexId: 1000}, {FieldId: 101}}
expansionReplicas := []*messagespb.LoadReplicaConfig{replicaConfig(1, rgA), replicaConfig(2, rgB)}
cases := []struct {
name string
seed func()
request message.BroadcastResultAlterLoadConfigMessageV2
expected bool
}{
{
name: "expansion into a new resource group",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields, expansionReplicas),
expected: true,
},
{
name: "the same expansion on a stock binary",
seed: func() {
ext.ResetForTest()
suite.seedLoadedCollection(1, rgA, 1)
},
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields, expansionReplicas),
expected: false,
},
{
name: "collection not loaded at all",
seed: func() { suite.seedReplica(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields, expansionReplicas),
expected: false,
},
{
name: "collection still loading",
seed: func() {
suite.seedCollection(querypb.LoadStatus_Loading, querypb.LoadType_LoadCollection, expansionDbID)
suite.seedReplica(1, rgA, 1)
},
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields, expansionReplicas),
expected: false,
},
{
name: "different db",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID+1, []int64{expansionPartitionID}, sameFields, expansionReplicas),
expected: false,
},
{
name: "collection was loaded by partition",
seed: func() {
suite.seedCollection(querypb.LoadStatus_Loaded, querypb.LoadType_LoadPartition, expansionDbID)
suite.seedReplica(1, rgA, 1)
},
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields, expansionReplicas),
expected: false,
},
{
name: "different partition set",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID, expansionPartitionID + 1}, sameFields, expansionReplicas),
expected: false,
},
{
name: "different load fields",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID},
[]*messagespb.LoadFieldConfig{{FieldId: 100, IndexId: 1000}}, expansionReplicas),
expected: false,
},
{
name: "different field index",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID},
[]*messagespb.LoadFieldConfig{{FieldId: 100, IndexId: 2000}, {FieldId: 101}}, expansionReplicas),
expected: false,
},
{
name: "an existing replica is dropped",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields,
[]*messagespb.LoadReplicaConfig{replicaConfig(2, rgB), replicaConfig(3, rgC)}),
expected: false,
},
{
name: "an existing replica moves to another resource group",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields,
[]*messagespb.LoadReplicaConfig{replicaConfig(1, rgB), replicaConfig(2, rgC)}),
expected: false,
},
{
name: "an added replica lands in a loaded resource group",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields,
[]*messagespb.LoadReplicaConfig{replicaConfig(1, rgA), replicaConfig(2, rgA)}),
expected: true,
},
{
name: "an identical replica set on a loaded collection is a replayed expansion",
seed: func() { suite.seedLoadedCollection(1, rgA, 1) },
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields,
[]*messagespb.LoadReplicaConfig{replicaConfig(1, rgA)}),
expected: true,
},
{
name: "an identical replica set on a loading collection keeps the overwrite",
seed: func() {
suite.seedCollection(querypb.LoadStatus_Loading, querypb.LoadType_LoadCollection, expansionDbID)
suite.seedReplica(1, rgA, 1)
},
request: suite.buildRequest(expansionDbID, []int64{expansionPartitionID}, sameFields,
[]*messagespb.LoadReplicaConfig{replicaConfig(1, rgA)}),
expected: false,
},
}
for _, tc := range cases {
suite.Run(tc.name, func() {
suite.SetupTest()
tc.seed()
job := suite.newPredicateJob(tc.request)
req := tc.request.Message.Header()
suite.Equal(tc.expected, job.isIncrementalExpansion(req, req.GetReplicas()))
})
}
}
// TestExpandedCollectionKeepsServingWhileNewResourceGroupLoads runs a real
// CollectionObserver tick after the expansion. The loaded resource group must
// still report a full load while the added one is still at zero: that is what
// "the resource groups load independently" means to everything downstream.
func (suite *IncrementalExpansionSuite) TestExpandedCollectionKeepsServingWhileNewResourceGroupLoads() {
suite.seedLoadedCollection(1, rgA, 1)
targetMgr := meta.NewMockTargetManager(suite.T())
channels := map[string]*meta.DmChannel{
expansionChannel: meta.DmChannelFromVChannel(&datapb.VchannelInfo{
CollectionID: expansionCollectionID,
ChannelName: expansionChannel,
}),
}
targetMgr.EXPECT().GetDmChannelsByCollection(mock.Anything, mock.Anything, mock.Anything).Return(channels).Maybe()
targetMgr.EXPECT().GetSealedSegmentsByCollection(mock.Anything, mock.Anything, mock.Anything).
Return(map[int64]*datapb.SegmentInfo{}).Maybe()
targetMgr.EXPECT().GetSealedSegmentsByPartition(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(map[int64]*datapb.SegmentInfo{}).Maybe()
targetMgr.EXPECT().IsNextTargetExist(mock.Anything, mock.Anything).Return(true).Maybe()
targetMgr.EXPECT().IsCurrentTargetExist(mock.Anything, mock.Anything, mock.Anything).Return(true).Maybe()
// The observer reads the next target's version to decide whether a
// figure of 100 may be reused; a constant is enough for a mock target.
targetMgr.EXPECT().GetCollectionTargetVersion(mock.Anything, mock.Anything, mock.Anything).Return(int64(1)).Maybe()
// Node 1 belongs to rgA's replica and already serves the only channel
// target; rgB's replica has no node carrying anything.
suite.dist.ChannelDistManager.Update(1, &meta.DmChannel{
VchannelInfo: &datapb.VchannelInfo{
CollectionID: expansionCollectionID,
ChannelName: expansionChannel,
},
Node: 1,
View: &meta.LeaderView{
ID: 1,
CollectionID: expansionCollectionID,
Channel: expansionChannel,
Segments: make(map[int64]*querypb.SegmentDist),
},
})
targetMock := mockey.Mock((*observers.TargetObserver).UpdateNextTarget).
To(func(ob *observers.TargetObserver, collectionID int64) (chan struct{}, error) {
return nil, nil
}).Build()
defer targetMock.UnPatch()
targetObserver := &observers.TargetObserver{}
collectionObserver := observers.NewCollectionObserver(suite.dist, suite.meta, targetMgr,
targetObserver, &checkers.CheckerController{}, nil)
result := suite.buildExpansionRequest(replicaConfig(1, rgA), replicaConfig(2, rgB))
job := NewLoadCollectionJob(suite.ctx, result, suite.dist, suite.meta, suite.broker,
targetMgr, targetObserver, collectionObserver, nil, nil, nil)
suite.Require().NoError(job.Execute())
// The per-resource-group percentage refuses a group that does not exist,
// so both have to be registered before it is asked about them.
rgCfg := &rgpb.ResourceGroupConfig{
Requests: &rgpb.ResourceGroupLimit{NodeNum: 0},
Limits: &rgpb.ResourceGroupLimit{NodeNum: 1},
}
_, err := suite.meta.AddResourceGroup(suite.ctx, rgA, rgCfg)
suite.Require().NoError(err)
_, err = suite.meta.AddResourceGroup(suite.ctx, rgB, rgCfg)
suite.Require().NoError(err)
// The added replica has no nodes yet, so rgB is at zero while rgA is full.
percentA, err := utils.LoadPercentageByResourceGroup(suite.ctx, suite.meta, targetMgr, suite.dist, expansionCollectionID, rgA)
suite.NoError(err)
suite.EqualValues(100, percentA)
percentB, err := utils.LoadPercentageByResourceGroup(suite.ctx, suite.meta, targetMgr, suite.dist, expansionCollectionID, rgB)
suite.NoError(err)
suite.EqualValues(0, percentB)
collectionObserver.Observe(suite.ctx)
collection := suite.meta.GetCollection(suite.ctx, expansionCollectionID)
suite.Require().NotNil(collection, "the observer must not release the collection")
suite.Equal(querypb.LoadStatus_Loaded, collection.GetStatus())
suite.EqualValues(100, suite.meta.CalculateLoadPercentage(suite.ctx, expansionCollectionID),
"the serving resource group must not be averaged down by the loading one")
suite.EqualValues(100, suite.meta.GetPartitionLoadPercentage(suite.ctx, expansionPartitionID))
suite.Len(suite.meta.GetByCollection(suite.ctx, expansionCollectionID), 2,
"neither resource group's replica may be torn down")
}
func TestIncrementalExpansion(t *testing.T) {
suite.Run(t, new(IncrementalExpansionSuite))
}
internal/querycoordv2/job/load_config.go
+34 -1
View File
@@ -27,6 +27,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/pkg/v3/extension"
"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/merr"
@@ -37,13 +38,45 @@ type AlterLoadConfigRequest struct {
CollectionInfo *milvuspb.DescribeCollectionResponse
Expected ExpectedLoadConfig
Current CurrentLoadConfig
// ScopedResourceGroups is the list of resource groups the REQUEST named, as
// the caller wrote it; empty when it named none. Only a form reads it (see
// CheckIfLoadPartitionsExecutable): with one installed, a request naming
// groups speaks only for those, and Expected.ExpectedReplicaNumber carries
// the other groups' counts through unchanged.
ScopedResourceGroups []string
}
// CheckIfLoadPartitionsExecutable checks if the load partitions is executable.
// CheckIfLoadPartitionsExecutable checks if the load partitions is executable:
// loading more partitions of a loaded collection may not change its replica
// number.
//
// On a stock binary that is the total, as it always was. With a form installed
// and a request that names resource groups, the request speaks only for those
// groups (the completed placement carries the others through unchanged, see
// completePlacementForOutOfScopeResourceGroups), so the rule is applied to the
// named groups alone: a named group that already holds replicas must keep its
// count, while a named group that holds none is being added, which is the
// expansion a scoped load exists for and not a replica-number change.
// Comparing the total there would refuse every such expansion.
func (req *AlterLoadConfigRequest) CheckIfLoadPartitionsExecutable() error {
if req.Current.Collection == nil {
return nil
}
if extension.FormInstalled() && len(req.ScopedResourceGroups) > 0 {
current := req.Current.GetReplicaNumber()
for _, rgName := range req.ScopedResourceGroups {
held := current[rgName]
if held == 0 {
continue // the request adds this group
}
if expected := req.Expected.ExpectedReplicaNumber[rgName]; held != expected {
return merr.WrapErrParameterInvalid(held, expected,
"can't change the replica number for loaded partitions in resource group "+rgName)
}
}
return nil
}
expectedReplicaNumber := 0
for _, num := range req.Expected.ExpectedReplicaNumber {
expectedReplicaNumber += num
internal/querycoordv2/job/load_config_test.go
+154
View File
@@ -0,0 +1,154 @@
//go:build test
package job
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
// These two tests pin the reconciliation the load-placement seam upstream of
// them depends on. The seam's whole job is to hand this function the second
// shape rather than the first, so if the reconciliation ever stopped
// distinguishing them the seam would go quietly ineffective.
func replicaIn(id int64, collectionID int64, rgName string) *meta.Replica {
return meta.NewReplica(&querypb.Replica{
ID: id,
CollectionID: collectionID,
ResourceGroup: rgName,
})
}
// A request that asks for a replica in rg_1 and leaves rg_0 at zero does not
// merely drop rg_0's replica: it reuses that very replica for rg_1. The
// resulting record says the replica lives in rg_1 and never mentions rg_0,
// which is why the placement has to be completed before the record is built
// and cannot be repaired afterwards.
func TestGenerateReplicasTransfersWhenAResourceGroupIsLeftAtZero(t *testing.T) {
req := &AlterLoadConfigRequest{
Current: CurrentLoadConfig{
Replicas: map[int64]*meta.Replica{1: replicaIn(1, 7, "rg_0")},
},
Expected: ExpectedLoadConfig{
ExpectedReplicaNumber: map[string]int{"rg_1": 1},
},
}
got, err := req.generateReplicas(context.Background())
require.NoError(t, err)
require.Len(t, got, 1)
assert.Equal(t, int64(1), got[0].GetReplicaId(),
"the existing replica is reused, not left alone")
assert.Equal(t, "rg_1", got[0].GetResourceGroupName(),
"and it is moved out of the resource group the request left at zero")
}
// The cumulative shape the seam produces: rg_0 keeps the replica it holds and
// rg_1 gets a newly allocated one, so the record describes a placement that
// grew rather than one that moved.
func TestGenerateReplicasAddsWhenEveryResourceGroupIsNamed(t *testing.T) {
m := meta.NewMeta(func() (int64, error) { return 99, nil }, nil, session.NewNodeManager())
req := &AlterLoadConfigRequest{
Meta: m,
Current: CurrentLoadConfig{
Replicas: map[int64]*meta.Replica{1: replicaIn(1, 7, "rg_0")},
},
Expected: ExpectedLoadConfig{
ExpectedReplicaNumber: map[string]int{"rg_0": 1, "rg_1": 1},
},
}
got, err := req.generateReplicas(context.Background())
require.NoError(t, err)
require.Len(t, got, 2)
placement := make(map[int64]string, len(got))
for _, replica := range got {
placement[replica.GetReplicaId()] = replica.GetResourceGroupName()
}
assert.Equal(t, "rg_0", placement[1],
"the replica that is already serving stays where it is")
assert.Equal(t, "rg_1", placement[99],
"and the newly named resource group gets a freshly allocated replica")
}
// loadedCollectionIn is a CurrentLoadConfig for a collection registered as
// loaded with one replica in each of the given resource groups, so that
// CheckIfLoadPartitionsExecutable has a loaded collection to judge.
func loadedCollectionIn(rgs ...string) CurrentLoadConfig {
replicas := make(map[int64]*meta.Replica, len(rgs))
for i, rg := range rgs {
replicas[int64(i+1)] = replicaIn(int64(i+1), 7, rg)
}
return CurrentLoadConfig{
Collection: &meta.Collection{CollectionLoadInfo: &querypb.CollectionLoadInfo{CollectionID: 7, ReplicaNumber: int32(len(rgs))}},
Replicas: replicas,
}
}
// setForm makes this test's binary one with a form installed, or a stock one,
// and restores a stock binary when the test ends.
func setForm(t *testing.T, installed bool) {
t.Helper()
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
if installed {
ext.SetForm()
}
}
// With a form installed, a scoped LoadPartitions that adds a resource group to
// a loaded collection is not a replica-number change: the groups it did not
// name are carried through unchanged by the completed placement, and the group
// it names holds nothing yet. Comparing the total, as the stock rule does,
// would refuse every such expansion for "changing the replica number".
func TestAScopedLoadPartitionsAddingAResourceGroupIsExecutableForAForm(t *testing.T) {
setForm(t, true)
req := &AlterLoadConfigRequest{
Current: loadedCollectionIn("rg_a"),
Expected: ExpectedLoadConfig{ExpectedReplicaNumber: map[string]int{"rg_a": 1, "rg_b": 1}},
ScopedResourceGroups: []string{"rg_b"},
}
assert.NoError(t, req.CheckIfLoadPartitionsExecutable())
}
// A scoped request that would change the count of a group that already holds
// replicas is still a replica-number change, and is refused as before.
func TestAScopedLoadPartitionsChangingALoadedGroupIsRefusedForAForm(t *testing.T) {
setForm(t, true)
req := &AlterLoadConfigRequest{
Current: loadedCollectionIn("rg_a"),
Expected: ExpectedLoadConfig{ExpectedReplicaNumber: map[string]int{"rg_a": 2}},
ScopedResourceGroups: []string{"rg_a"},
}
assert.ErrorIs(t, req.CheckIfLoadPartitionsExecutable(), merr.ErrParameterInvalid)
}
// A stock binary keeps master's rule - the total replica count of loaded
// partitions cannot change - whatever the request named.
func TestAStockBinaryComparesTheTotalReplicaCount(t *testing.T) {
setForm(t, false)
req := &AlterLoadConfigRequest{
Current: loadedCollectionIn("rg_a"),
Expected: ExpectedLoadConfig{ExpectedReplicaNumber: map[string]int{"rg_a": 1, "rg_b": 1}},
ScopedResourceGroups: []string{"rg_b"},
}
assert.ErrorIs(t, req.CheckIfLoadPartitionsExecutable(), merr.ErrParameterInvalid)
same := &AlterLoadConfigRequest{
Current: loadedCollectionIn("rg_a"),
Expected: ExpectedLoadConfig{ExpectedReplicaNumber: map[string]int{"rg_b": 1}},
}
assert.NoError(t, same.CheckIfLoadPartitionsExecutable(), "one replica for one replica is not a change")
}
internal/querycoordv2/load_admission_test.go
+438
View File
@@ -0,0 +1,438 @@
// 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 querycoordv2
import (
"context"
"testing"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/rgpb"
"github.com/milvus-io/milvus/internal/coordinator/snmanager"
"github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/querycoordv2/job"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/params"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
"github.com/milvus-io/milvus/internal/util/streamingutil"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// admissionFixture is a meta whose resource groups hold real nodes, so that
// the LoadCollection callback's admission - AssignReplica - runs for real
// against them instead of being replaced by a mock. Nothing here touches etcd:
// the catalog is a mock that accepts every write.
type admissionFixture struct {
meta *meta.Meta
nodeMgr *session.NodeManager
broker *meta.MockBroker
}
func newAdmissionFixture(t *testing.T) *admissionFixture {
t.Helper()
paramtable.Init()
catalog := mocks.NewQueryCoordCatalog(t)
catalog.On("SaveReplica", mock.Anything, mock.Anything).Return(nil).Maybe()
catalog.On("SaveReplica", mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
catalog.On("SaveResourceGroup", mock.Anything, mock.Anything).Return(nil).Maybe()
nodeMgr := session.NewNodeManager()
m := &meta.Meta{
CollectionManager: meta.NewCollectionManager(catalog),
ReplicaManager: meta.NewReplicaManager(params.RandomIncrementIDAllocator(), catalog),
ResourceManager: meta.NewResourceManager(catalog, nodeMgr),
}
broker := meta.NewMockBroker(t)
broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).
RunAndReturn(func(_ context.Context, collectionID int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{CollectionID: collectionID}, nil
}).Maybe()
broker.EXPECT().GetPartitions(mock.Anything, mock.Anything).Return([]int64{1}, nil).Maybe()
return &admissionFixture{meta: m, nodeMgr: nodeMgr, broker: broker}
}
// putResourceGroup registers rgName holding exactly the given regular query
// nodes; with none, the group has no compute at all.
func (f *admissionFixture) putResourceGroup(t *testing.T, rgName string, nodeIDs ...int64) {
t.Helper()
ctx := context.Background()
wanted := int32(len(nodeIDs))
_, err := f.meta.AddResourceGroup(ctx, rgName, &rgpb.ResourceGroupConfig{
Requests: &rgpb.ResourceGroupLimit{NodeNum: wanted},
Limits: &rgpb.ResourceGroupLimit{NodeNum: wanted},
})
require.NoError(t, err)
for _, nodeID := range nodeIDs {
f.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: nodeID,
Address: "localhost",
Hostname: "localhost",
}))
f.meta.HandleNodeUp(ctx, nodeID)
}
nodes, err := f.meta.GetNodes(ctx, rgName)
require.NoError(t, err)
require.Len(t, nodes, len(nodeIDs), "the resource manager must hold exactly the query nodes given")
}
// loadedIn registers collectionID as loaded with one replica in each of the
// given resource groups, which is what makes the next LoadCollection on it a
// request on a loaded collection rather than a first load.
func (f *admissionFixture) loadedIn(t *testing.T, collectionID int64, rgNames ...string) {
t.Helper()
ctx := context.Background()
require.NoError(t, f.meta.PutCollectionWithoutSave(ctx, &meta.Collection{
CollectionLoadInfo: &querypb.CollectionLoadInfo{
CollectionID: collectionID,
ReplicaNumber: int32(len(rgNames)),
Status: querypb.LoadStatus_Loaded,
},
}))
for i, rgName := range rgNames {
if !f.meta.ContainResourceGroup(ctx, rgName) {
f.putResourceGroup(t, rgName)
}
require.NoError(t, f.meta.Put(ctx, meta.NewReplica(&querypb.Replica{
ID: collectionID*100 + int64(i+1),
CollectionID: collectionID,
ResourceGroup: rgName,
})))
}
}
// withStreamingQueryNodes turns the streaming service on and makes the
// streaming node manager answer from byRG, so a test can state which resource
// group holds which streaming query nodes without one running. The
// cluster-wide set is derived from byRG: the two are the same nodes.
func withStreamingQueryNodes(byRG map[string]typeutil.UniqueSet) func() {
all := typeutil.NewUniqueSet()
for _, nodes := range byRG {
all.Insert(nodes.Collect()...)
}
enabled := mockey.Mock(streamingutil.IsStreamingServiceEnabled).Return(true).Build()
byGroup := mockey.Mock((*snmanager.StreamingNodeManager).GetStreamingQueryNodeIDsByResourceGroup).
Return(byRG).Build()
cluster := mockey.Mock((*snmanager.StreamingNodeManager).GetStreamingQueryNodeIDs).
Return(all).Build()
return func() {
cluster.UnPatch()
byGroup.UnPatch()
enabled.UnPatch()
}
}
// loadCollection runs the LoadCollection callback against the fixture with the
// real admission, and reports whether the callback asked the admission to check
// node numbers. The broadcast itself is stubbed: what is under test is the
// decision taken before it.
func (f *admissionFixture) loadCollection(t *testing.T, req *querypb.LoadCollectionRequest) (checkedNodeNum bool, err error) {
t.Helper()
s := &Server{meta: f.meta, broker: f.broker}
var origin func(context.Context, *meta.Meta, []string, int32, bool) (map[string]int, error)
broadcast := mockey.Mock((*Server).startBroadcastWithCollectionIDLock).Return(stubBroadcaster{}, nil).Build()
defer broadcast.UnPatch()
admission := mockey.Mock(utils.AssignReplica).To(
func(ctx context.Context, m *meta.Meta, resourceGroups []string, replicaNumber int32, checkNodeNum bool) (map[string]int, error) {
checkedNodeNum = checkNodeNum
return origin(ctx, m, resourceGroups, replicaNumber, checkNodeNum)
}).Origin(&origin).Build()
defer admission.UnPatch()
generate := mockey.Mock(job.GenerateAlterLoadConfigMessage).Return(nil, nil).Build()
defer generate.UnPatch()
err = s.broadcastAlterLoadConfigCollectionV2ForLoadCollection(context.Background(), req)
return checkedNodeNum, err
}
// loadPartitions is loadCollection for the LoadPartitions callback, which
// checks node numbers on every request.
func (f *admissionFixture) loadPartitions(t *testing.T, req *querypb.LoadPartitionsRequest) (checkedNodeNum bool, err error) {
t.Helper()
s := &Server{meta: f.meta, broker: f.broker}
var origin func(context.Context, *meta.Meta, []string, int32, bool) (map[string]int, error)
broadcast := mockey.Mock((*Server).startBroadcastWithCollectionIDLock).Return(stubBroadcaster{}, nil).Build()
defer broadcast.UnPatch()
admission := mockey.Mock(utils.AssignReplica).To(
func(ctx context.Context, m *meta.Meta, resourceGroups []string, replicaNumber int32, checkNodeNum bool) (map[string]int, error) {
checkedNodeNum = checkNodeNum
return origin(ctx, m, resourceGroups, replicaNumber, checkNodeNum)
}).Origin(&origin).Build()
defer admission.UnPatch()
generate := mockey.Mock(job.GenerateAlterLoadConfigMessage).Return(nil, nil).Build()
defer generate.UnPatch()
err = s.broadcastAlterLoadConfigCollectionV2ForLoadPartitions(context.Background(), req)
return checkedNodeNum, err
}
// theFormsOnlyStreamingNode is the streaming side of a form's instance: one
// streaming node, for DDL and the write ahead log, in a group of its own that
// holds no replica. Every query cluster is a group of regular query nodes.
func theFormsOnlyStreamingNode() map[string]typeutil.UniqueSet {
return map[string]typeutil.UniqueSet{"internal": typeutil.NewUniqueSet(7)}
}
// LoadPartitions checks node numbers on every request. A collection loaded
// into rg_b with two replicas on its two query nodes, one of them restarting,
// and the identical LoadPartitions re-sent: the request adds no replica and is
// not admitted against anything, while the same request asking a third
// replica is.
func TestAnIdenticalScopedLoadPartitionsReSendIsNotRefused(t *testing.T) {
setForm(t, true)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_b", 1)
f.loadedIn(t, 17, "rg_b", "rg_b")
defer withStreamingQueryNodes(theFormsOnlyStreamingNode())()
checked, err := f.loadPartitions(t, &querypb.LoadPartitionsRequest{
CollectionID: 17,
PartitionIDs: []int64{1},
ReplicaNumber: 2,
ResourceGroups: []string{"rg_b"},
})
assert.False(t, checked, "a re-send adds no replica, so there is nothing to admit")
assert.NoError(t, err)
checked, err = f.loadPartitions(t, &querypb.LoadPartitionsRequest{
CollectionID: 17,
PartitionIDs: []int64{1},
ReplicaNumber: 3,
ResourceGroups: []string{"rg_b"},
})
assert.True(t, checked, "a third replica is an expansion")
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrResourceGroupNodeNotEnough, "three replicas on rg_b's one query node")
}
// A stock binary checks node numbers on every load_partitions, as it always
// has, whatever the request repeats.
func TestAStockLoadPartitionsAlwaysChecksNodeNumbers(t *testing.T) {
setForm(t, false)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_b", 1, 2)
f.loadedIn(t, 18, "rg_b", "rg_b")
defer withStreamingQueryNodes(map[string]typeutil.UniqueSet{
"rg_elsewhere": typeutil.NewUniqueSet(301, 302),
})()
checked, err := f.loadPartitions(t, &querypb.LoadPartitionsRequest{
CollectionID: 18,
PartitionIDs: []int64{1},
ReplicaNumber: 2,
ResourceGroups: []string{"rg_b"},
})
assert.True(t, checked, "master's rule: every load_partitions is checked")
assert.NoError(t, err, "two replicas on two regular nodes, as master admits")
}
// The reviewer's finding: a client re-sending the load that placed a
// collection's two replicas in rg_b, while one of rg_b's two query nodes
// restarts, names a group on a loaded collection and so was admitted like an
// expansion: two replicas against one node, refused, thirty lines before the
// callback would have recognized a request that changes nothing. Admission
// runs only for what a request adds; an identical re-send adds nothing and
// reaches the no-op, and the same request asking a third replica is refused.
func TestAnIdenticalScopedReSendIsANoOpNotARefusal(t *testing.T) {
setForm(t, true)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_b", 1)
f.loadedIn(t, 15, "rg_b", "rg_b")
defer withStreamingQueryNodes(theFormsOnlyStreamingNode())()
checked, err := f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 15,
ReplicaNumber: 2,
ResourceGroups: []string{"rg_b"},
})
assert.False(t, checked, "a re-send adds no replica, so there is nothing to admit")
assert.NoError(t, err, "the collection is still serving; the re-send is the no-op it always was")
checked, err = f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 15,
ReplicaNumber: 3,
ResourceGroups: []string{"rg_b"},
})
assert.True(t, checked, "a third replica is an expansion")
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrResourceGroupNodeNotEnough)
}
// With BOTH of rg_b's query nodes away the regular bound would refuse two
// replicas on none; the re-send still adds nothing and is still the no-op.
func TestAnIdenticalScopedReSendIsANoOpWhileAllTheGroupsQueryNodesRestart(t *testing.T) {
setForm(t, true)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_b")
f.loadedIn(t, 16, "rg_b", "rg_b")
defer withStreamingQueryNodes(theFormsOnlyStreamingNode())()
checked, err := f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 16,
ReplicaNumber: 2,
ResourceGroups: []string{"rg_b"},
})
assert.False(t, checked)
assert.NoError(t, err)
}
// The reviewer's finding: master checks node numbers only for a first load,
// and a config update on a loaded collection skips the check. With a form, a
// LoadCollection naming resource groups on a loaded collection is a scoped
// expansion into those groups - a placement decision like a first load - and
// it was admitted with no bound at all. A group with no node would receive a
// replica that never gets a delegator, the scoped task's clock would pause
// forever, and ShowLoadCollections scoped to the group would report 0 for as
// long as anybody polled it.
func TestAScopedExpansionIntoAGroupWithNoNodeIsRefused(t *testing.T) {
setForm(t, true)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_a", 1)
f.putResourceGroup(t, "rg_b")
f.loadedIn(t, 7, "rg_a")
defer withStreamingQueryNodes(theFormsOnlyStreamingNode())()
checked, err := f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 7,
ReplicaNumber: 1,
ResourceGroups: []string{"rg_b"},
})
assert.True(t, checked, "a scoped expansion is a placement decision and is checked like a first load")
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrResourceGroupNodeNotEnough)
}
// One query node in the group serves one replica; asking two is refused up
// front rather than leaving the second waiting forever.
func TestAScopedExpansionAskingMoreReplicasThanTheGroupHasQueryNodesIsRefused(t *testing.T) {
setForm(t, true)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_a", 1)
f.putResourceGroup(t, "rg_b", 21)
f.loadedIn(t, 8, "rg_a")
defer withStreamingQueryNodes(theFormsOnlyStreamingNode())()
checked, err := f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 8,
ReplicaNumber: 2,
ResourceGroups: []string{"rg_b"},
})
assert.True(t, checked)
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrResourceGroupNodeNotEnough)
}
// A group with enough query nodes admits the expansion, exactly as a first
// load into it would be admitted - and with a form installed it is admitted
// although the whole instance runs one streaming node: a form's delegators
// go onto the replicas' own query nodes (UseStreamingQueryNodeAsDelegator).
func TestAScopedExpansionIntoAGroupWithEnoughNodesIsAdmitted(t *testing.T) {
setForm(t, true)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_a", 1)
f.putResourceGroup(t, "rg_b", 21, 22)
f.loadedIn(t, 9, "rg_a")
defer withStreamingQueryNodes(theFormsOnlyStreamingNode())()
checked, err := f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 9,
ReplicaNumber: 2,
ResourceGroups: []string{"rg_b"},
})
assert.True(t, checked)
assert.NoError(t, err)
}
// A form's request that names no resource group on a loaded collection is a
// config update, not a placement into a new group, and keeps master's rule:
// no node check.
func TestAFormsConfigUpdateOnALoadedCollectionSkipsTheNodeCheck(t *testing.T) {
setForm(t, true)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_a", 1)
f.loadedIn(t, 10, "rg_a")
defer withStreamingQueryNodes(map[string]typeutil.UniqueSet{
"rg_a": typeutil.NewUniqueSet(101),
})()
checked, err := f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 10,
ReplicaNumber: 1,
})
assert.False(t, checked, "a config update on a loaded collection skips the node check, as master does")
assert.NoError(t, err)
}
// The stock binary is unchanged: a LoadCollection naming resource groups on a
// loaded collection is a config update there, and master never checks node
// numbers for one, whatever the named group holds.
func TestAStockLoadOnALoadedCollectionSkipsTheNodeCheckWhateverGroupsItNames(t *testing.T) {
setForm(t, false)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_a", 1)
f.putResourceGroup(t, "rg_b")
f.loadedIn(t, 11, "rg_a")
defer withStreamingQueryNodes(map[string]typeutil.UniqueSet{
"rg_a": typeutil.NewUniqueSet(101),
})()
checked, err := f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 11,
ReplicaNumber: 1,
ResourceGroups: []string{"rg_b"},
})
assert.False(t, checked, "on a stock binary a load on a loaded collection never checks node numbers")
assert.NoError(t, err, "the empty group is not checked, exactly as before")
}
// A cluster-level force override discards the groups the request named and
// makes the load state the whole placement; on a loaded collection that is a
// config update, and the scoping decision the placement follows says so. The
// node check follows the same decision, not a second reading of the request.
func TestAFormsForceOverriddenLoadOnALoadedCollectionSkipsTheNodeCheck(t *testing.T) {
setForm(t, true)
f := newAdmissionFixture(t)
f.putResourceGroup(t, "rg_a", 1)
f.putResourceGroup(t, "rg_b")
f.loadedIn(t, 12, "rg_a")
defer withStreamingQueryNodes(map[string]typeutil.UniqueSet{
"rg_a": typeutil.NewUniqueSet(101),
})()
p := paramtable.Get()
require.NoError(t, p.Save(p.QueryCoordCfg.ClusterLevelLoadForceOverrideUserReplicaMode.Key, "true"))
require.NoError(t, p.Save(p.QueryCoordCfg.ClusterLevelLoadReplicaNumber.Key, "1"))
require.NoError(t, p.Save(p.QueryCoordCfg.ClusterLevelLoadResourceGroups.Key, "rg_a"))
t.Cleanup(func() {
p.Reset(p.QueryCoordCfg.ClusterLevelLoadForceOverrideUserReplicaMode.Key)
p.Reset(p.QueryCoordCfg.ClusterLevelLoadReplicaNumber.Key)
p.Reset(p.QueryCoordCfg.ClusterLevelLoadResourceGroups.Key)
})
checked, err := f.loadCollection(t, &querypb.LoadCollectionRequest{
CollectionID: 12,
ReplicaNumber: 1,
ResourceGroups: []string{"rg_b"},
})
assert.False(t, checked, "the override made this a whole-placement config update")
assert.NoError(t, err)
}
internal/querycoordv2/load_placement.go
+107
View File
@@ -0,0 +1,107 @@
// 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 querycoordv2
import (
"context"
"maps"
"github.com/milvus-io/milvus/internal/querycoordv2/job"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// completePlacementForOutOfScopeResourceGroups returns the per-resource-group
// replica counts to record for this load request.
//
// On a stock binary the answer is always what utils.AssignReplica produced: a
// load request states the whole placement, and a second LoadCollection naming
// rg_1 on a collection loaded in rg_0 moves the replica to rg_1, as it always
// has (job/load_config_test.go pins the reconciliation that does the moving).
//
// With a form installed (extension.FormInstalled) a load request that NAMES
// resource groups speaks only for those: the counts the collection already has
// in the groups it did not name are carried through alongside the requested
// ones. A request that names none still speaks for the whole placement. The
// result is the cumulative placement the collection ends up with, so the
// reconciliation leaves those replicas where they are, and the record it
// builds describes a placement that grew rather than one that moved - which is
// what lets the load job recognize a request that only adds resource groups
// as a pure expansion and keep the collection's serving state intact. A
// distribution that loads one collection into several resource groups
// independently needs that reading; a stock deployment, whose LoadCollection
// has always meant "this is the placement", keeps its contract, including the
// request's replica_number being the total.
//
// requestedResourceGroups is the list the REQUEST named, not the one the load
// path defaulted: see getLoadReplicaConfigForRequest, which returns the two
// separately, for why the difference decides the answer.
//
// expected is the non-nil map utils.AssignReplica just returned. It is never
// mutated, because the caller may still log it.
func completePlacementForOutOfScopeResourceGroups(
ctx context.Context,
collectionID int64,
requestedResourceGroups []string,
expected map[string]int,
current job.CurrentLoadConfig,
) map[string]int {
if !extension.FormInstalled() || len(requestedResourceGroups) == 0 {
// A stock binary states the whole placement on every request, and a
// form does so for a request naming no group: what AssignReplica
// produced is the whole answer, and carrying anything over would add
// replicas the request did not ask for.
return expected
}
named := typeutil.NewSet(requestedResourceGroups...)
completed := maps.Clone(expected)
carried := make([]string, 0)
for rgName, replicaNumber := range current.GetReplicaNumber() {
if named.Contain(rgName) {
continue
}
completed[rgName] = replicaNumber
carried = append(carried, rgName)
}
if len(carried) > 0 {
mlog.Info(ctx, "load request is scoped to the resource groups it names, keeping the placement of the others",
mlog.Int64("collectionID", collectionID),
mlog.Strings("requestedResourceGroups", requestedResourceGroups),
mlog.Strings("keptResourceGroups", carried),
)
}
return completed
}
// scopedLoadAddsReplicas reports whether a scoped request asks any of the
// groups it names for more replicas than the collection already holds there.
// Only such a request places anything and is admitted against node numbers;
// a request that changes nothing - a client re-sending the load that placed
// the collection, while a node of the group restarts - or that only shrinks
// a group adds no replica, and reaches the no-op or the shrink it asks for
// without being refused for compute it does not need.
func scopedLoadAddsReplicas(requested map[string]int, current job.CurrentLoadConfig) bool {
held := current.GetReplicaNumber()
for rgName, num := range requested {
if num > held[rgName] {
return true
}
}
return false
}
internal/querycoordv2/load_placement_test.go
+459
View File
@@ -0,0 +1,459 @@
// 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 querycoordv2
import (
"context"
"maps"
"slices"
"testing"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/querycoordv2/job"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// setForm makes this test's binary one with a form installed, or a stock one,
// and restores a stock binary when the test ends.
func setForm(t *testing.T, installed bool) {
t.Helper()
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
if installed {
ext.SetForm()
}
}
func loadedIn(collectionID int64, resourceGroups ...string) job.CurrentLoadConfig {
replicas := make(map[int64]*meta.Replica, len(resourceGroups))
for i, rgName := range resourceGroups {
id := int64(i + 1)
replicas[id] = meta.NewReplica(&querypb.Replica{
ID: id,
CollectionID: collectionID,
ResourceGroup: rgName,
})
}
return job.CurrentLoadConfig{Replicas: replicas}
}
func TestCompletePlacementKeepsOutOfScopeResourceGroups(t *testing.T) {
paramtable.Init()
setForm(t, true)
expected := map[string]int{"rg_1": 1}
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, []string{"rg_1"}, expected, loadedIn(7, "rg_0"))
assert.Equal(t, map[string]int{"rg_0": 1, "rg_1": 1}, got,
"the resource group the request did not name keeps the replica it holds")
assert.Equal(t, map[string]int{"rg_1": 1}, expected,
"the map AssignReplica returned must not be mutated")
}
func TestCompletePlacementDoesNotOverrideNamedResourceGroups(t *testing.T) {
paramtable.Init()
setForm(t, true)
current := loadedIn(7, "rg_0", "rg_1")
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, []string{"rg_1"}, map[string]int{"rg_1": 2}, current)
assert.Equal(t, map[string]int{"rg_0": 1, "rg_1": 2}, got)
}
func TestCompletePlacementOnFirstLoadIsUnchanged(t *testing.T) {
paramtable.Init()
setForm(t, true)
expected := map[string]int{"rg_0": 1}
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, []string{"rg_0"}, expected, job.CurrentLoadConfig{})
assert.Equal(t, map[string]int{"rg_0": 1}, got)
}
func TestCompletePlacementWhenRequestNamesEveryLoadedResourceGroup(t *testing.T) {
paramtable.Init()
setForm(t, true)
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, []string{"rg_0"}, map[string]int{"rg_0": 1}, loadedIn(7, "rg_0"))
assert.Equal(t, map[string]int{"rg_0": 1}, got)
}
func TestCompletePlacementKeepsEverySibling(t *testing.T) {
paramtable.Init()
setForm(t, true)
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, []string{"rg_2"}, map[string]int{"rg_2": 1},
loadedIn(7, "rg_0", "rg_1"))
assert.Equal(t, map[string]int{"rg_0": 1, "rg_1": 1, "rg_2": 1}, got)
}
func TestCompletePlacementKeepsSiblingReplicaCount(t *testing.T) {
paramtable.Init()
setForm(t, true)
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, []string{"rg_1"}, map[string]int{"rg_1": 1},
loadedIn(7, "rg_0", "rg_0"))
assert.Equal(t, map[string]int{"rg_0": 2, "rg_1": 1}, got)
}
type stubBroadcaster struct{}
func (stubBroadcaster) Broadcast(context.Context, message.BroadcastMutableMessage) (*types.BroadcastAppendResult, error) {
return nil, nil
}
func (stubBroadcaster) Close() {}
// capturedPlacement holds the parts of an AlterLoadConfigRequest these tests
// assert on. The request itself is a local of the callback under test, and the
// compiler is free to keep it on that callback's frame: escape analysis ran
// against the real GenerateAlterLoadConfigMessage, which does not let the
// argument outlive the call. A mock that stored the pointer would therefore be
// reading a dead frame once the callback returned, so the mock copies out what
// the test needs while the request is still alive.
type capturedPlacement struct {
ExpectedReplicaNumber map[string]int
ScopedResourceGroups []string
}
func capturePlacement(req *job.AlterLoadConfigRequest) *capturedPlacement {
return &capturedPlacement{
ExpectedReplicaNumber: maps.Clone(req.Expected.ExpectedReplicaNumber),
ScopedResourceGroups: slices.Clone(req.ScopedResourceGroups),
}
}
func TestLoadCollectionBroadcastAppliesTheCompletedPlacement(t *testing.T) {
paramtable.Init()
setForm(t, true)
const collectionID = int64(7)
broker := meta.NewMockBroker(t)
broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
Return(&milvuspb.DescribeCollectionResponse{CollectionID: collectionID}, nil).Maybe()
broker.EXPECT().GetPartitions(mock.Anything, collectionID).Return([]int64{1}, nil).Maybe()
s := &Server{broker: broker}
var captured *capturedPlacement
mockey.PatchConvey("a load naming one resource group carries the other's placement", t, func() {
mockey.Mock((*Server).startBroadcastWithCollectionIDLock).
Return(stubBroadcaster{}, nil).Build()
mockey.Mock(utils.AssignReplica).Return(map[string]int{"rg_1": 1}, nil).Build()
mockey.Mock((*Server).getCurrentLoadConfig).Return(loadedIn(collectionID, "rg_0")).Build()
mockey.Mock(job.GenerateAlterLoadConfigMessage).To(
func(_ context.Context, req *job.AlterLoadConfigRequest) (message.BroadcastMutableMessage, error) {
captured = capturePlacement(req)
return nil, nil
}).Build()
err := s.broadcastAlterLoadConfigCollectionV2ForLoadCollection(context.Background(),
&querypb.LoadCollectionRequest{
CollectionID: collectionID,
ReplicaNumber: 1,
ResourceGroups: []string{"rg_1"},
})
assert.NoError(t, err)
})
require.NotNil(t, captured, "the broadcast request must have been built")
assert.Equal(t, map[string]int{"rg_0": 1, "rg_1": 1},
captured.ExpectedReplicaNumber,
"the load path must broadcast the completed placement, not the one it asked for")
}
func TestCompletePlacementForARequestNamingNoResourceGroup(t *testing.T) {
paramtable.Init()
setForm(t, true)
expected := map[string]int{meta.DefaultResourceGroupName: 1}
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, nil, expected, loadedIn(7, "rg_0"))
assert.Equal(t, map[string]int{meta.DefaultResourceGroupName: 1}, got,
"a request naming no resource group states the whole placement, so nothing is carried over")
}
func TestCompletePlacementWhenTheDefaultResourceGroupIsNamedExplicitly(t *testing.T) {
paramtable.Init()
setForm(t, true)
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, []string{meta.DefaultResourceGroupName},
map[string]int{meta.DefaultResourceGroupName: 1}, loadedIn(7, "rg_0"))
assert.Equal(t, map[string]int{meta.DefaultResourceGroupName: 1, "rg_0": 1}, got,
"naming the default resource group explicitly is a scoped request like any other")
}
// The scoping list and the list AssignReplica works from come out of one call,
// and they differ exactly where the defaulting rewrote the request.
func TestLoadReplicaConfigSeparatesTheRequestedGroupsFromTheDefaultedOnes(t *testing.T) {
paramtable.Init()
s := &Server{}
replicaNumber, resourceGroups, scoped, userSpecified, err := s.getLoadReplicaConfigForRequest(
context.Background(), 1, nil, 7)
require.NoError(t, err)
assert.EqualValues(t, 1, replicaNumber)
assert.Equal(t, []string{meta.DefaultResourceGroupName}, resourceGroups,
"the defaulted list is what AssignReplica is given")
assert.Empty(t, scoped, "a request naming no group states the whole placement")
assert.True(t, userSpecified)
_, resourceGroups, scoped, _, err = s.getLoadReplicaConfigForRequest(
context.Background(), 1, []string{"rg_0"}, 7)
require.NoError(t, err)
assert.Equal(t, []string{"rg_0"}, resourceGroups)
assert.Equal(t, []string{"rg_0"}, scoped, "a request naming a group is scoped to it")
}
// A cluster-level force override states the whole placement, whatever the
// request named, and the flag behind it is read once so the two answers cannot
// disagree.
func TestLoadReplicaConfigUnderClusterLevelForceOverride(t *testing.T) {
paramtable.Init()
p := paramtable.Get()
require.NoError(t, p.Save(p.QueryCoordCfg.ClusterLevelLoadForceOverrideUserReplicaMode.Key, "true"))
require.NoError(t, p.Save(p.QueryCoordCfg.ClusterLevelLoadReplicaNumber.Key, "2"))
require.NoError(t, p.Save(p.QueryCoordCfg.ClusterLevelLoadResourceGroups.Key, "rg_0"))
t.Cleanup(func() {
p.Reset(p.QueryCoordCfg.ClusterLevelLoadForceOverrideUserReplicaMode.Key)
p.Reset(p.QueryCoordCfg.ClusterLevelLoadReplicaNumber.Key)
p.Reset(p.QueryCoordCfg.ClusterLevelLoadResourceGroups.Key)
})
s := &Server{}
replicaNumber, resourceGroups, scoped, userSpecified, err := s.getLoadReplicaConfigForRequest(
context.Background(), 1, []string{"rg_1"}, 7)
require.NoError(t, err)
assert.EqualValues(t, 2, replicaNumber)
assert.Equal(t, []string{"rg_0"}, resourceGroups)
assert.Empty(t, scoped)
assert.False(t, userSpecified)
}
// TestLoadCollectionBroadcastOfABareRequestKeepsThePlacementItAsksFor is the
// LoadCollection half of the reviewer's scenario: a load naming no resource
// group must not be read as a load scoped to the default resource group, or
// the replicas the collection already has elsewhere are carried through and
// the collection ends up with one more replica than anybody asked for.
func TestLoadCollectionBroadcastOfABareRequestKeepsThePlacementItAsksFor(t *testing.T) {
paramtable.Init()
setForm(t, true)
const collectionID = int64(8)
broker := meta.NewMockBroker(t)
broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
Return(&milvuspb.DescribeCollectionResponse{CollectionID: collectionID}, nil).Maybe()
broker.EXPECT().GetPartitions(mock.Anything, collectionID).Return([]int64{1}, nil).Maybe()
s := &Server{broker: broker}
var captured *capturedPlacement
mockey.PatchConvey("a load naming no resource group states the whole placement", t, func() {
mockey.Mock((*Server).startBroadcastWithCollectionIDLock).
Return(stubBroadcaster{}, nil).Build()
mockey.Mock(utils.AssignReplica).
Return(map[string]int{meta.DefaultResourceGroupName: 1}, nil).Build()
mockey.Mock((*Server).getCurrentLoadConfig).Return(loadedIn(collectionID, "rg_a")).Build()
mockey.Mock(job.GenerateAlterLoadConfigMessage).To(
func(_ context.Context, req *job.AlterLoadConfigRequest) (message.BroadcastMutableMessage, error) {
captured = capturePlacement(req)
return nil, nil
}).Build()
err := s.broadcastAlterLoadConfigCollectionV2ForLoadCollection(context.Background(),
&querypb.LoadCollectionRequest{
CollectionID: collectionID,
ReplicaNumber: 1,
})
assert.NoError(t, err)
})
require.NotNil(t, captured, "the broadcast request must have been built")
assert.Equal(t, map[string]int{meta.DefaultResourceGroupName: 1},
captured.ExpectedReplicaNumber,
"a bare load must not carry the replicas of a resource group it never named")
}
// TestLoadPartitionsBroadcastOfABareRequestIsNotRefused is the LoadPartitions
// half: reading a bare request as scoped to the default resource group makes
// the expected placement one replica larger than the collection has, and
// CheckIfLoadPartitionsExecutable refuses the load for "changing the replica
// number" of a resource group the caller never mentioned.
func TestLoadPartitionsBroadcastOfABareRequestIsNotRefused(t *testing.T) {
paramtable.Init()
setForm(t, true)
const collectionID = int64(9)
broker := meta.NewMockBroker(t)
broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
Return(&milvuspb.DescribeCollectionResponse{CollectionID: collectionID}, nil).Maybe()
s := &Server{broker: broker}
current := loadedIn(collectionID, "rg_a")
current.Collection = &meta.Collection{
CollectionLoadInfo: &querypb.CollectionLoadInfo{CollectionID: collectionID, ReplicaNumber: 1},
}
var captured *capturedPlacement
mockey.PatchConvey("a bare load_partitions on a collection loaded elsewhere", t, func() {
mockey.Mock((*Server).startBroadcastWithCollectionIDLock).
Return(stubBroadcaster{}, nil).Build()
mockey.Mock(utils.AssignReplica).
Return(map[string]int{meta.DefaultResourceGroupName: 1}, nil).Build()
mockey.Mock((*Server).getCurrentLoadConfig).Return(current).Build()
mockey.Mock(job.GenerateAlterLoadConfigMessage).To(
func(_ context.Context, req *job.AlterLoadConfigRequest) (message.BroadcastMutableMessage, error) {
captured = capturePlacement(req)
return nil, nil
}).Build()
err := s.broadcastAlterLoadConfigCollectionV2ForLoadPartitions(context.Background(),
&querypb.LoadPartitionsRequest{
CollectionID: collectionID,
PartitionIDs: []int64{1},
ReplicaNumber: 1,
})
assert.NoError(t, err, "a bare load_partitions must not be refused for a replica number it never asked to change")
})
require.NotNil(t, captured, "the broadcast request must have been built")
assert.Equal(t, map[string]int{meta.DefaultResourceGroupName: 1},
captured.ExpectedReplicaNumber)
}
// TestLoadPartitionsBroadcastOfAScopedRequestCarriesTheSiblings pins the other
// direction on the same callback: a request that does name a resource group is
// scoped, and the groups it did not name keep their replicas.
func TestLoadPartitionsBroadcastOfAScopedRequestCarriesTheSiblings(t *testing.T) {
paramtable.Init()
setForm(t, true)
const collectionID = int64(10)
broker := meta.NewMockBroker(t)
broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
Return(&milvuspb.DescribeCollectionResponse{CollectionID: collectionID}, nil).Maybe()
s := &Server{broker: broker}
// The collection is registered as loaded, so CheckIfLoadPartitionsExecutable
// runs for real: the scoped expansion must pass it.
current := loadedIn(collectionID, "rg_a")
current.Collection = &meta.Collection{
CollectionLoadInfo: &querypb.CollectionLoadInfo{CollectionID: collectionID, ReplicaNumber: 1},
}
var captured *capturedPlacement
mockey.PatchConvey("a load_partitions naming one resource group carries the other's placement", t, func() {
mockey.Mock((*Server).startBroadcastWithCollectionIDLock).
Return(stubBroadcaster{}, nil).Build()
mockey.Mock(utils.AssignReplica).Return(map[string]int{"rg_b": 1}, nil).Build()
mockey.Mock((*Server).getCurrentLoadConfig).Return(current).Build()
mockey.Mock(job.GenerateAlterLoadConfigMessage).To(
func(_ context.Context, req *job.AlterLoadConfigRequest) (message.BroadcastMutableMessage, error) {
captured = capturePlacement(req)
return nil, nil
}).Build()
err := s.broadcastAlterLoadConfigCollectionV2ForLoadPartitions(context.Background(),
&querypb.LoadPartitionsRequest{
CollectionID: collectionID,
PartitionIDs: []int64{1},
ReplicaNumber: 1,
ResourceGroups: []string{"rg_b"},
})
assert.NoError(t, err, "adding a group to a loaded collection is not a replica-number change for a form")
})
require.NotNil(t, captured)
assert.Equal(t, map[string]int{"rg_a": 1, "rg_b": 1}, captured.ExpectedReplicaNumber)
assert.Equal(t, []string{"rg_b"}, captured.ScopedResourceGroups, "the scoping list travels with the request")
}
// A stock binary never completes a placement: the request states it whole,
// whatever groups it names, and the function hands back exactly what
// AssignReplica produced.
func TestAStockBinaryStatesTheWholePlacementOnEveryRequest(t *testing.T) {
paramtable.Init()
setForm(t, false)
expected := map[string]int{"rg_1": 1}
got := completePlacementForOutOfScopeResourceGroups(
context.Background(), 7, []string{"rg_1"}, expected, loadedIn(7, "rg_0"))
assert.Equal(t, map[string]int{"rg_1": 1}, got,
"on a stock binary a load naming rg_1 moves the replica there, as it always has")
}
// The LoadCollection callback on a stock binary: a second load naming rg_1 on
// a collection loaded in rg_0 broadcasts a placement of rg_1 alone, so the
// reconciliation moves the replica - master's contract, in which the request's
// replica_number is the total.
func TestAStockLoadCollectionBroadcastStatesTheWholePlacement(t *testing.T) {
paramtable.Init()
setForm(t, false)
const collectionID = int64(11)
broker := meta.NewMockBroker(t)
broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
Return(&milvuspb.DescribeCollectionResponse{CollectionID: collectionID}, nil).Maybe()
broker.EXPECT().GetPartitions(mock.Anything, collectionID).Return([]int64{1}, nil).Maybe()
s := &Server{broker: broker}
var captured *capturedPlacement
mockey.PatchConvey("a stock load naming one resource group states the whole placement", t, func() {
mockey.Mock((*Server).startBroadcastWithCollectionIDLock).
Return(stubBroadcaster{}, nil).Build()
mockey.Mock(utils.AssignReplica).Return(map[string]int{"rg_1": 1}, nil).Build()
mockey.Mock((*Server).getCurrentLoadConfig).Return(loadedIn(collectionID, "rg_0")).Build()
mockey.Mock(job.GenerateAlterLoadConfigMessage).To(
func(_ context.Context, req *job.AlterLoadConfigRequest) (message.BroadcastMutableMessage, error) {
captured = capturePlacement(req)
return nil, nil
}).Build()
err := s.broadcastAlterLoadConfigCollectionV2ForLoadCollection(context.Background(),
&querypb.LoadCollectionRequest{
CollectionID: collectionID,
ReplicaNumber: 1,
ResourceGroups: []string{"rg_1"},
})
assert.NoError(t, err)
})
require.NotNil(t, captured, "the broadcast request must have been built")
assert.Equal(t, map[string]int{"rg_1": 1}, captured.ExpectedReplicaNumber,
"a stock binary must not carry the replicas of a resource group the request left out")
}
// The LoadPartitions callback on a stock binary, same shape: the placement
// broadcast is the one the request asked for, and the load is not refused,
// since one replica is replaced by one replica.
func TestAStockLoadPartitionsBroadcastStatesTheWholePlacement(t *testing.T) {
paramtable.Init()
setForm(t, false)
const collectionID = int64(12)
broker := meta.NewMockBroker(t)
broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
Return(&milvuspb.DescribeCollectionResponse{CollectionID: collectionID}, nil).Maybe()
s := &Server{broker: broker}
current := loadedIn(collectionID, "rg_a")
current.Collection = &meta.Collection{
CollectionLoadInfo: &querypb.CollectionLoadInfo{CollectionID: collectionID, ReplicaNumber: 1},
}
var captured *capturedPlacement
mockey.PatchConvey("a stock load_partitions naming one resource group states the whole placement", t, func() {
mockey.Mock((*Server).startBroadcastWithCollectionIDLock).
Return(stubBroadcaster{}, nil).Build()
mockey.Mock(utils.AssignReplica).Return(map[string]int{"rg_b": 1}, nil).Build()
mockey.Mock((*Server).getCurrentLoadConfig).Return(current).Build()
mockey.Mock(job.GenerateAlterLoadConfigMessage).To(
func(_ context.Context, req *job.AlterLoadConfigRequest) (message.BroadcastMutableMessage, error) {
captured = capturePlacement(req)
return nil, nil
}).Build()
err := s.broadcastAlterLoadConfigCollectionV2ForLoadPartitions(context.Background(),
&querypb.LoadPartitionsRequest{
CollectionID: collectionID,
PartitionIDs: []int64{1},
ReplicaNumber: 1,
ResourceGroups: []string{"rg_b"},
})
assert.NoError(t, err, "one replica for one replica is not a replica-number change")
})
require.NotNil(t, captured)
assert.Equal(t, map[string]int{"rg_b": 1}, captured.ExpectedReplicaNumber)
}
internal/querycoordv2/meta/collection_manager.go
+43
View File
@@ -672,6 +672,49 @@ func (m *CollectionManager) UpdateReplicaNumber(ctx context.Context, collectionI
if !ok {
return merr.WrapErrCollectionNotFound(collectionID)
}
return m.setReplicaNumberLocked(ctx, collection, replicaNumber, userSpecifiedReplicaMode)
}
// SyncReplicaNumber sets the collection's ReplicaNumber, and its partitions',
// to what count returns, and answers the number the collection then holds.
//
// count is evaluated INSIDE the manager's critical section, so that reading
// how many replicas there are and writing their number is one step. A writer
// that reads the replicas first and calls UpdateReplicaNumber afterwards can
// be overtaken in between: a resource group's teardown reads two replicas, a
// concurrent expansion spawns a third and writes three, and the teardown's
// two lands last - a count one short of the replicas that exist, with nothing
// left to correct it. Every writer whose number IS the number of replicas
// goes through here; a writer spawns or removes its replicas before it calls,
// so whichever of two writers lands last has counted the other's replicas.
// count must not call back into this manager.
//
// userSpecifiedReplicaMode is written with the number when it is not nil, and
// nil keeps the stored one: a writer that only follows the replicas has no
// say in how they were asked for. Nothing is written when neither changes.
func (m *CollectionManager) SyncReplicaNumber(ctx context.Context, collectionID typeutil.UniqueID, count func() int32, userSpecifiedReplicaMode *bool) (int32, error) {
m.rwmutex.Lock()
defer m.rwmutex.Unlock()
collection, ok := m.collections[collectionID]
if !ok {
return 0, merr.WrapErrCollectionNotFound(collectionID)
}
replicaNumber := count()
mode := collection.GetUserSpecifiedReplicaMode()
if userSpecifiedReplicaMode != nil {
mode = *userSpecifiedReplicaMode
}
if collection.GetReplicaNumber() == replicaNumber && collection.GetUserSpecifiedReplicaMode() == mode {
return replicaNumber, nil
}
return replicaNumber, m.setReplicaNumberLocked(ctx, collection, replicaNumber, mode)
}
// setReplicaNumberLocked writes the replica number and mode to the collection
// and its partitions. The caller holds rwmutex.
func (m *CollectionManager) setReplicaNumberLocked(ctx context.Context, collection *Collection, replicaNumber int32, userSpecifiedReplicaMode bool) error {
collectionID := collection.GetCollectionID()
newCollection := collection.Clone()
newCollection.ReplicaNumber = replicaNumber
newCollection.UserSpecifiedReplicaMode = userSpecifiedReplicaMode
internal/querycoordv2/meta/collection_manager_replica_number_test.go
+172
View File
@@ -0,0 +1,172 @@
// 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 meta
import (
"context"
"testing"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
const (
syncCollection = int64(4200)
syncPartition = int64(4201)
)
// newSyncFixture is a manager holding one collection with one partition, both
// at replicaNumber, whose catalog accepts exactly saves writes - the one that
// puts the collection there included.
func newSyncFixture(t *testing.T, saves int, replicaNumber int32, userSpecified bool) *CollectionManager {
t.Helper()
catalog := catalogmocks.NewQueryCoordCatalog(t)
catalog.EXPECT().SaveCollection(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(saves)
mgr := NewCollectionManager(catalog)
require.NoError(t, mgr.PutCollection(context.Background(),
&Collection{CollectionLoadInfo: &querypb.CollectionLoadInfo{
CollectionID: syncCollection, ReplicaNumber: replicaNumber, UserSpecifiedReplicaMode: userSpecified,
}},
&Partition{PartitionLoadInfo: &querypb.PartitionLoadInfo{
CollectionID: syncCollection, PartitionID: syncPartition, ReplicaNumber: replicaNumber,
}}))
return mgr
}
// The point of SyncReplicaNumber: the replicas are counted while the manager
// holds its lock, so no other writer of the number can land between the count
// and the write. A reader that cannot get in while count runs is the proof.
func TestSyncReplicaNumberCountsInsideTheCriticalSection(t *testing.T) {
ctx := context.Background()
mgr := newSyncFixture(t, 2, 3, true)
counted := false
got, err := mgr.SyncReplicaNumber(ctx, syncCollection, func() int32 {
counted = true
if mgr.rwmutex.TryRLock() {
mgr.rwmutex.RUnlock()
t.Error("the manager's lock is free while the replicas are counted")
}
return 2
}, nil)
require.NoError(t, err)
require.True(t, counted)
assert.EqualValues(t, 2, got)
assert.EqualValues(t, 2, mgr.GetReplicaNumber(ctx, syncCollection))
partitions := mgr.GetPartitionsByCollection(ctx, syncCollection)
require.Len(t, partitions, 1)
assert.EqualValues(t, 2, partitions[0].GetReplicaNumber(), "the partitions follow the collection")
assert.True(t, mgr.GetCollection(ctx, syncCollection).GetUserSpecifiedReplicaMode(),
"nil keeps the stored mode: a writer that only follows the replicas has no say in how they were asked for")
}
func TestSyncReplicaNumberWritesTheModeItIsGiven(t *testing.T) {
ctx := context.Background()
mgr := newSyncFixture(t, 2, 1, false)
userSpecified := true
got, err := mgr.SyncReplicaNumber(ctx, syncCollection, func() int32 { return 1 }, &userSpecified)
require.NoError(t, err)
assert.EqualValues(t, 1, got)
assert.True(t, mgr.GetCollection(ctx, syncCollection).GetUserSpecifiedReplicaMode(),
"a changed mode is written even when the number is not")
}
// The catalog accepts one write, the one that put the collection there: a
// sync that changes nothing must not make another.
func TestSyncReplicaNumberWritesNothingWhenNothingChanges(t *testing.T) {
ctx := context.Background()
mgr := newSyncFixture(t, 1, 2, true)
got, err := mgr.SyncReplicaNumber(ctx, syncCollection, func() int32 { return 2 }, nil)
require.NoError(t, err)
assert.EqualValues(t, 2, got)
userSpecified := true
_, err = mgr.SyncReplicaNumber(ctx, syncCollection, func() int32 { return 2 }, &userSpecified)
require.NoError(t, err)
}
func TestSyncReplicaNumberOfAMissingCollection(t *testing.T) {
mgr := NewCollectionManager(catalogmocks.NewQueryCoordCatalog(t))
_, err := mgr.SyncReplicaNumber(context.Background(), syncCollection, func() int32 {
t.Error("there is nothing to count replicas for")
return 0
}, nil)
assert.ErrorIs(t, err, merr.ErrCollectionNotFound)
}
func TestSyncReplicaNumberLeavesTheNumberAloneWhenTheCatalogRefuses(t *testing.T) {
ctx := context.Background()
catalog := catalogmocks.NewQueryCoordCatalog(t)
catalog.EXPECT().SaveCollection(mock.Anything, mock.Anything).Return(nil).Once()
mgr := NewCollectionManager(catalog)
require.NoError(t, mgr.PutCollection(ctx, &Collection{CollectionLoadInfo: &querypb.CollectionLoadInfo{
CollectionID: syncCollection, ReplicaNumber: 3,
}}))
refused := errors.New("catalog is away")
catalog.EXPECT().SaveCollection(mock.Anything, mock.Anything).Return(refused).Once()
_, err := mgr.SyncReplicaNumber(ctx, syncCollection, func() int32 { return 2 }, nil)
assert.ErrorIs(t, err, refused)
assert.EqualValues(t, 3, mgr.GetReplicaNumber(ctx, syncCollection), "so the caller's retry still has something to correct")
}
// The race this exists for, in both orders. A resource group's teardown has
// removed a replica and an expansion adds one: each changes the replicas
// before it syncs, so whichever syncs last has counted the other's, and the
// number ends at the replicas that exist. With the count read outside the
// lock and written afterwards, the teardown's stale two could land last.
func TestWhicheverWriterSyncsLastHasCountedTheOthersReplicas(t *testing.T) {
ctx := context.Background()
t.Run("the expansion syncs last", func(t *testing.T) {
// Three writes: the put, the teardown's three to two, the expansion's two to three.
mgr := newSyncFixture(t, 3, 3, false)
replicas := int32(2) // the teardown has removed one of three
count := func() int32 { return replicas }
_, err := mgr.SyncReplicaNumber(ctx, syncCollection, count, nil)
require.NoError(t, err)
replicas++ // the expansion spawns its replica, then syncs
_, err = mgr.SyncReplicaNumber(ctx, syncCollection, count, nil)
require.NoError(t, err)
assert.EqualValues(t, 3, mgr.GetReplicaNumber(ctx, syncCollection))
})
t.Run("the teardown syncs last", func(t *testing.T) {
// One write, the put: three were stored, and three is what both count.
mgr := newSyncFixture(t, 1, 3, false)
replicas := int32(2)
count := func() int32 { return replicas }
replicas++ // the expansion spawns and syncs first
_, err := mgr.SyncReplicaNumber(ctx, syncCollection, count, nil)
require.NoError(t, err)
_, err = mgr.SyncReplicaNumber(ctx, syncCollection, count, nil)
require.NoError(t, err)
assert.EqualValues(t, 3, mgr.GetReplicaNumber(ctx, syncCollection))
})
}
internal/querycoordv2/ops_services.go
+1 -1
View File
@@ -476,7 +476,7 @@ func (s *Server) TransferChannel(ctx context.Context, req *querypb.TransferChann
// when no dst node specified, default to use all other nodes in same
dstNodeSet := typeutil.NewUniqueSet()
if req.GetToAllNodes() {
if streamingutil.IsStreamingServiceEnabled() {
if streamingutil.UseStreamingQueryNodeAsDelegator() {
dstNodeSet.Insert(replica.GetRWSQNodes()...)
} else {
dstNodeSet.Insert(replica.GetRWNodes()...)
internal/querycoordv2/resource_group_load_percentage_test.go
+97 -1
View File
@@ -48,6 +48,7 @@ type rgLoadPercentageFixture struct {
targetMgr *meta.TargetManager
dist *meta.DistributionManager
broker *meta.MockBroker
nodeMgr *session.NodeManager
}
func newRGLoadPercentageFixture(t *testing.T) *rgLoadPercentageFixture {
@@ -74,6 +75,7 @@ func newRGLoadPercentageFixture(t *testing.T) *rgLoadPercentageFixture {
targetMgr: meta.NewTargetManager(broker, m),
dist: meta.NewDistributionManager(nodeMgr),
broker: broker,
nodeMgr: nodeMgr,
}
}
@@ -84,7 +86,7 @@ func newRGLoadPercentageFixture(t *testing.T) *rgLoadPercentageFixture {
// merr.CheckHealthy(s.State()), and a zero Server reports
// StateCode_Initializing.
func (f *rgLoadPercentageFixture) server() *Server {
s := &Server{meta: f.meta, targetMgr: f.targetMgr, dist: f.dist}
s := &Server{meta: f.meta, targetMgr: f.targetMgr, dist: f.dist, nodeMgr: f.nodeMgr}
s.status.Store(int32(commonpb.StateCode_Healthy))
return s
}
@@ -187,6 +189,30 @@ func (f *rgLoadPercentageFixture) putDelegator(collectionID, nodeID int64, chann
f.dist.ChannelDistManager.Update(nodeID, channel)
}
// promoteTarget makes the collection's next target its current one, which
// is the target the shard-leader readiness reads. The percentage is
// unchanged by it: that figure reads the next target first and falls back
// to the current one.
func (f *rgLoadPercentageFixture) promoteTarget(t *testing.T, collectionID int64) {
t.Helper()
require.True(t, f.targetMgr.UpdateCollectionCurrentTarget(context.Background(), collectionID))
}
// registerNode makes nodeID one the coordinator knows, which a leader must
// sit on to count as serving.
func (f *rgLoadPercentageFixture) registerNode(nodeID int64) {
f.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: nodeID, Address: "localhost", Hostname: "localhost"}))
}
// putServiceableDelegator is putDelegator with the leader view reporting
// itself serviceable: the delegator can answer a query.
func (f *rgLoadPercentageFixture) putServiceableDelegator(collectionID, nodeID int64, channelName string, segmentIDs ...int64) {
f.putDelegator(collectionID, nodeID, channelName, segmentIDs...)
channel := f.dist.ChannelDistManager.GetByFilter(meta.WithNodeID2Channel(nodeID))[0]
channel.View.Status = &querypb.LeaderViewStatus{Serviceable: true}
f.dist.ChannelDistManager.Update(nodeID, channel)
}
// TestGetLoadPercentageByResourceGroup_NoReplicaOnRG asserts that a
// collection with no replica in the requested resource group reports -1,
// even though the collection is fully loaded elsewhere. Deleting the
@@ -337,6 +363,76 @@ func TestGetLoadPercentageByResourceGroup_MultipleReplicasSameRG(t *testing.T) {
assert.EqualValues(t, 0, percentage)
}
// TestReplicaLoadPercentagesByResourceGroup_OneFigurePerReplica pins the
// per-replica reading the group figure is folded from. The observer's load
// timeout releases the replicas that stalled rather than the group, so it
// has to know which replica holds the group at 0 here: the map names it,
// and LoadPercentageByResourceGroup is exactly its minimum.
func TestReplicaLoadPercentagesByResourceGroup_OneFigurePerReplica(t *testing.T) {
f := newRGLoadPercentageFixture(t)
f.putTarget(t, 610, 6100, "610-dmc0", 1, 2)
f.putReplica(t, 610, 62, "rg-shared")
f.putDelegator(610, 62, "610-dmc0", 1, 2) // this replica: fully loaded
f.putReplica(t, 610, 63, "rg-shared")
f.putDelegator(610, 63, "610-dmc0", 1) // this replica: the channel and one of two segments
f.putReplica(t, 610, 64, "rg-shared")
// this replica: nothing loaded
f.putReplica(t, 610, 65, "rg-other")
// a replica of another group is not in the answer at all
ctx := context.Background()
figures, err := utils.ReplicaLoadPercentagesByResourceGroup(ctx, f.meta, f.targetMgr, f.dist, 610, "rg-shared")
require.NoError(t, err)
require.Len(t, figures, 3)
byNode := make(map[int64]int32, len(figures))
for _, replica := range f.meta.GetByCollection(ctx, 610) {
if figure, ok := figures[replica.GetID()]; ok {
byNode[replica.GetNodes()[0]] = figure
}
}
assert.Equal(t, map[int64]int32{62: 100, 63: 66, 64: 0}, byNode)
percentage, err := f.freeFn(ctx, 610, "rg-shared")
require.NoError(t, err)
assert.EqualValues(t, 0, percentage, "the group's figure is the minimum over its replicas")
}
// TestReplicaLoadPercentagesByResourceGroup_OutcomesMatchTheGroupFigure pins
// that the per-replica reading keeps the group figure's contract for the
// outcomes that carry no figure: a group holding no replica is an empty map
// with no error, where the group figure is -1; an unknown group and an
// unloaded collection with a recorded failure are the same errors.
func TestReplicaLoadPercentagesByResourceGroup_OutcomesMatchTheGroupFigure(t *testing.T) {
f := newRGLoadPercentageFixture(t)
f.putTarget(t, 620, 6200, "620-dmc0", 1)
f.putReplica(t, 620, 66, "rg-here")
f.putResourceGroup(t, "rg-empty")
ctx := context.Background()
figures, err := utils.ReplicaLoadPercentagesByResourceGroup(ctx, f.meta, f.targetMgr, f.dist, 620, "rg-empty")
assert.NoError(t, err)
assert.Empty(t, figures, "a group holding no replica has nothing to report a figure for")
_, err = utils.ReplicaLoadPercentagesByResourceGroup(ctx, f.meta, f.targetMgr, f.dist, 620, "rg-typo")
assert.ErrorIs(t, err, merr.ErrResourceGroupNotFound)
_, err = utils.ReplicaLoadPercentagesByResourceGroup(ctx, nil, f.targetMgr, f.dist, 620, "rg-here")
assert.ErrorIs(t, err, merr.ErrServiceNotReady)
figures, err = utils.ReplicaLoadPercentagesByResourceGroup(ctx, f.meta, f.targetMgr, f.dist, 621, "rg-here")
assert.NoError(t, err)
assert.Empty(t, figures, "a collection that is not registered as loaded has no figure either")
seedFailedLoadCache(t, 621, merr.WrapErrServiceUnavailable("query node restarting"))
_, err = utils.ReplicaLoadPercentagesByResourceGroup(ctx, f.meta, f.targetMgr, f.dist, 621, "rg-here")
assert.ErrorIs(t, err, merr.ErrCollectionNotLoaded)
}
// seedFailedLoadCache installs a fresh GlobalFailedLoadCache holding err for
// collectionID, and restores the previous cache when the test ends so the
// package-global state does not leak into other tests in this package.
internal/querycoordv2/services.go
+51 -1
View File
@@ -66,6 +66,19 @@ func (s *Server) ShowLoadCollections(ctx context.Context, req *querypb.ShowColle
}
defer meta.GlobalFailedLoadCache.TryExpire()
// A scope names a resource group; one the resource manager does not know
// is refused before any collection is looked at. The loaded path below
// validates it too, but a collection that is not loaded answers -1 for
// the group without reaching that path, and a client polling a misspelled
// group after a refused load would wait on -1 forever.
if rgName := req.GetResourceGroup(); rgName != "" && !s.meta.ContainResourceGroup(ctx, rgName) {
err := merr.WrapErrResourceGroupNotFound(rgName)
mlog.Warn(ctx, "show collection failed on an unknown resource group", mlog.String("resourceGroup", rgName), mlog.Err(err))
return &querypb.ShowCollectionsResponse{
Status: merr.Status(err),
}, nil
}
isGetAll := false
collectionSet := typeutil.NewUniqueSet(req.GetCollectionIDs()...)
if len(req.GetCollectionIDs()) == 0 {
@@ -102,6 +115,21 @@ func (s *Server) ShowLoadCollections(ctx context.Context, req *querypb.ShowColle
}, nil
}
if req.GetResourceGroup() != "" {
// Scoped to a resource group, a collection that is not loaded
// at all and has no failure recorded is a collection with no
// replica in that group: -1, the same answer a loaded
// collection gives for a group it has no replica in, so a
// caller polling one group's progress right after issuing a
// load is not told the load failed.
resp.CollectionIDs = append(resp.CollectionIDs, collectionID)
resp.InMemoryPercentages = append(resp.InMemoryPercentages, -1)
resp.QueryServiceAvailable = append(resp.QueryServiceAvailable, false)
resp.RefreshProgress = append(resp.RefreshProgress, 0)
resp.LoadFields = append(resp.LoadFields, &schemapb.LongArray{})
continue
}
err = merr.WrapErrCollectionNotLoaded(collectionID)
mlog.Warn(context.TODO(), "show collection failed", mlog.Err(err))
return &querypb.ShowCollectionsResponse{
@@ -113,9 +141,31 @@ func (s *Server) ShowLoadCollections(ctx context.Context, req *querypb.ShowColle
refreshProgress = 100
}
// A request that names a resource group asks for the progress of the
// replicas living there, not the collection-wide figure; -1 says the
// group holds no replica of this collection. Whether the service is
// available follows the same scope: the caller asks whether THAT
// group can serve, and a group that cannot must not be reported as
// serving because some other group can.
var queryServiceAvailable bool
if rgName := req.GetResourceGroup(); rgName != "" {
scoped, err := utils.LoadPercentageByResourceGroup(ctx, s.meta, s.targetMgr, s.dist, collectionID, rgName)
if err != nil {
mlog.Warn(ctx, "show collection failed on the resource-group scoped progress",
mlog.Int64("collectionID", collectionID), mlog.String("resourceGroup", rgName), mlog.Err(err))
return &querypb.ShowCollectionsResponse{
Status: merr.Status(err),
}, nil
}
percentage = scoped
queryServiceAvailable = s.checkAnyReplicaAvailableInResourceGroup(ctx, collectionID, rgName)
} else {
queryServiceAvailable = s.checkAnyReplicaAvailable(collectionID)
}
resp.CollectionIDs = append(resp.CollectionIDs, collectionID)
resp.InMemoryPercentages = append(resp.InMemoryPercentages, int64(percentage))
resp.QueryServiceAvailable = append(resp.QueryServiceAvailable, s.checkAnyReplicaAvailable(collectionID))
resp.QueryServiceAvailable = append(resp.QueryServiceAvailable, queryServiceAvailable)
resp.RefreshProgress = append(resp.RefreshProgress, refreshProgress)
resp.LoadFields = append(resp.LoadFields, &schemapb.LongArray{
Data: loadFields,
internal/querycoordv2/show_load_collections_rg_test.go
+277
View File
@@ -0,0 +1,277 @@
// 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 querycoordv2
import (
"context"
"testing"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
// withFailedLoadCache makes sure the global failed-load cache ShowLoadCollections
// expires on every call exists for the test.
func withFailedLoadCache(t *testing.T) {
t.Helper()
if meta.GlobalFailedLoadCache == nil {
meta.GlobalFailedLoadCache = meta.NewFailedLoadCache()
}
}
func TestShowLoadCollectionsWithoutAResourceGroupKeepsTheCollectionWideFigure(t *testing.T) {
withFailedLoadCache(t)
f := newRGLoadPercentageFixture(t)
f.putTarget(t, 100, 1000, "100-dmc0", 1, 2)
f.putReplica(t, 100, 10, "rg-a")
f.putReplica(t, 100, 20, "rg-b")
f.putDelegator(100, 10, "100-dmc0", 1, 2) // rg-a fully loaded, rg-b has nothing yet
resp, err := f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{100},
})
require.NoError(t, err)
require.NoError(t, merr.Error(resp.GetStatus()))
require.Equal(t, []int64{100}, resp.GetCollectionIDs())
assert.Equal(t, f.meta.CalculateLoadPercentage(context.Background(), 100), int32(resp.GetInMemoryPercentages()[0]),
"without a resource group the answer is the collection-wide percentage, as before")
}
func TestShowLoadCollectionsScopedToAResourceGroup(t *testing.T) {
withFailedLoadCache(t)
f := newRGLoadPercentageFixture(t)
f.putTarget(t, 100, 1000, "100-dmc0", 1, 2)
f.putReplica(t, 100, 10, "rg-a")
f.putReplica(t, 100, 20, "rg-b")
f.putDelegator(100, 10, "100-dmc0", 1, 2) // rg-a fully loaded, rg-b has nothing yet
f.putResourceGroup(t, "rg-empty")
show := func(rg string) *querypb.ShowCollectionsResponse {
resp, err := f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{100},
ResourceGroup: rg,
})
require.NoError(t, err)
require.NoError(t, merr.Error(resp.GetStatus()))
require.Equal(t, []int64{100}, resp.GetCollectionIDs())
return resp
}
assert.EqualValues(t, 100, show("rg-a").GetInMemoryPercentages()[0],
"the group whose replica serves every target is at 100 whatever its sibling does")
assert.EqualValues(t, 0, show("rg-b").GetInMemoryPercentages()[0],
"the group whose replica serves nothing yet is at 0, not at the collection-wide figure")
assert.EqualValues(t, -1, show("rg-empty").GetInMemoryPercentages()[0],
"a group that holds no replica of the collection answers -1, which is not 0")
}
func TestShowLoadCollectionsScopedToAnUnknownResourceGroupIsRefused(t *testing.T) {
withFailedLoadCache(t)
f := newRGLoadPercentageFixture(t)
f.putTarget(t, 100, 1000, "100-dmc0", 1, 2)
f.putReplica(t, 100, 10, "rg-a")
f.putDelegator(100, 10, "100-dmc0", 1, 2)
resp, err := f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{100},
ResourceGroup: "rg-missing",
})
require.NoError(t, err)
assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrResourceGroupNotFound)
// The reviewer's finding: a collection that is not loaded took the -1
// branch before anything looked at the group, so a client polling
// ShowLoadCollections{C, "rg_typo"} after a refused load waited on -1
// forever. The group is validated before the collections are.
const neverLoaded = int64(777)
resp, err = f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{neverLoaded},
ResourceGroup: "rg-missing",
})
require.NoError(t, err)
assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrResourceGroupNotFound,
"an unknown group is refused whether or not the collection is loaded")
resp, err = f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
ResourceGroup: "rg-missing",
})
require.NoError(t, err)
assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrResourceGroupNotFound,
"and so is a request that names no collection at all")
}
func TestShowLoadCollectionsScopedAnswersMinusOneForAnUnloadedCollection(t *testing.T) {
withFailedLoadCache(t)
f := newRGLoadPercentageFixture(t)
f.putResourceGroup(t, "rg-a")
const neverLoaded = int64(777)
resp, err := f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{neverLoaded},
})
require.NoError(t, err)
assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrCollectionNotLoaded,
"unscoped, an unloaded collection is refused exactly as before")
resp, err = f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{neverLoaded},
ResourceGroup: "rg-a",
})
require.NoError(t, err)
require.NoError(t, merr.Error(resp.GetStatus()))
require.Equal(t, []int64{neverLoaded}, resp.GetCollectionIDs())
assert.EqualValues(t, -1, resp.GetInMemoryPercentages()[0],
"scoped, an unloaded collection with no recorded failure holds no replica in the group: -1, not a refusal")
assert.False(t, resp.GetQueryServiceAvailable()[0])
}
// putReplicaWithRONode registers a replica of collectionID in rgName homed on
// nodeID, with roNodeID as a read-only node the node manager does not know:
// under the availability rule ShowLoadCollections answers with, which is
// master's - a replica is available when every one of its read-only nodes is
// known to the node manager - that is a replica that is not available.
func (f *rgLoadPercentageFixture) putReplicaWithRONode(t *testing.T, collectionID, nodeID, roNodeID int64, rgName string) {
t.Helper()
f.putResourceGroup(t, rgName)
require.Nil(t, f.nodeMgr.Get(roNodeID), "the read-only node must be unknown to the node manager")
require.NoError(t, f.meta.Put(context.Background(), meta.NewReplica(&querypb.Replica{
ID: nodeID,
CollectionID: collectionID,
ResourceGroup: rgName,
Nodes: []int64{nodeID},
RoNodes: []int64{roNodeID},
})))
}
// A scoped row's query_service_available answers from the group's
// shard-leader readiness - every shard of the collection has a serviceable
// leader in the group's replicas, on a node the coordinator knows - not
// from the collection-wide rule (every read-only node alive, a replica with
// none counting as available), which read true for a replica that had
// loaded nothing. A collection loaded in rg-a and just expanded into rg-b,
// whose fresh replica holds no delegator, answers 0 and false for rg-b,
// 100 and true for rg-a, and the collection-wide 100 and true - by the
// unchanged rule - without a scope.
func TestShowLoadCollectionsScopedAnswersQueryServiceAvailableFromTheGroupsShardLeaders(t *testing.T) {
withFailedLoadCache(t)
f := newRGLoadPercentageFixture(t)
f.putTarget(t, 100, 1000, "100-dmc0", 1, 2)
f.putReplica(t, 100, 10, "rg-a")
f.registerNode(10)
f.putServiceableDelegator(100, 10, "100-dmc0", 1, 2)
f.promoteTarget(t, 100)
// The collection-wide figure is the mean of its partitions' own: rg-a
// carries everything, so the collection reads as loaded.
require.NoError(t, f.meta.PutPartitionWithoutSave(context.Background(), &meta.Partition{
PartitionLoadInfo: &querypb.PartitionLoadInfo{CollectionID: 100, PartitionID: 1000},
LoadPercentage: 100,
}))
// rg-b's replica was just spawned: no delegator, nothing loaded, and no
// read-only node either, which is what made the old rule call it available.
f.putReplica(t, 100, 20, "rg-b")
show := func(rg string) *querypb.ShowCollectionsResponse {
resp, err := f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{100},
ResourceGroup: rg,
})
require.NoError(t, err)
require.NoError(t, merr.Error(resp.GetStatus()))
require.Equal(t, []int64{100}, resp.GetCollectionIDs())
return resp
}
fresh := show("rg-b")
assert.EqualValues(t, 0, fresh.GetInMemoryPercentages()[0], "sanity: the group's own progress")
assert.False(t, fresh.GetQueryServiceAvailable()[0],
"a replica that serves no shard is not available, whatever its read-only nodes say")
serving := show("rg-a")
assert.EqualValues(t, 100, serving.GetInMemoryPercentages()[0])
assert.True(t, serving.GetQueryServiceAvailable()[0], "every shard has a serviceable leader in rg-a")
unscoped := show("")
assert.EqualValues(t, 100, unscoped.GetInMemoryPercentages()[0])
assert.True(t, unscoped.GetQueryServiceAvailable()[0],
"unscoped, the collection-wide answer is unchanged: rg-a serves")
}
// A group whose leader is not serviceable, or sits on a node the
// coordinator does not know, is not serving, however far its load is.
func TestShowLoadCollectionsScopedIsNotAvailableWithoutAServiceableLeader(t *testing.T) {
withFailedLoadCache(t)
f := newRGLoadPercentageFixture(t)
f.putTarget(t, 100, 1000, "100-dmc0", 1, 2)
f.putReplica(t, 100, 10, "rg-a")
f.putDelegator(100, 10, "100-dmc0", 1, 2) // every segment there, but the view does not report serviceable
f.promoteTarget(t, 100)
resp, err := f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{100},
ResourceGroup: "rg-a",
})
require.NoError(t, err)
require.NoError(t, merr.Error(resp.GetStatus()))
assert.EqualValues(t, 100, resp.GetInMemoryPercentages()[0], "sanity: fully loaded by the percentage")
assert.False(t, resp.GetQueryServiceAvailable()[0], "loaded is not serving: the leader is not serviceable")
}
// Each request answers availability in its own scope and computes only that
// answer: a request naming a resource group asks the group's shard leaders and
// never the collection-wide check, which it would only overwrite.
func TestShowLoadCollectionsComputesOnlyTheAvailabilityItsScopeAsksFor(t *testing.T) {
withFailedLoadCache(t)
f := newRGLoadPercentageFixture(t)
f.putTarget(t, 100, 1000, "100-dmc0", 1, 2)
f.putReplica(t, 100, 10, "rg-a")
f.registerNode(10)
f.putServiceableDelegator(100, 10, "100-dmc0", 1, 2)
f.promoteTarget(t, 100)
require.NoError(t, f.meta.PutPartitionWithoutSave(context.Background(), &meta.Partition{
PartitionLoadInfo: &querypb.PartitionLoadInfo{CollectionID: 100, PartitionID: 1000},
LoadPercentage: 100,
}))
collectionWide := 0
var origin func(*Server, int64) bool
counting := mockey.Mock((*Server).checkAnyReplicaAvailable).To(func(s *Server, collectionID int64) bool {
collectionWide++
return origin(s, collectionID)
}).Origin(&origin).Build()
defer counting.UnPatch()
show := func(rg string) bool {
resp, err := f.server().ShowLoadCollections(context.Background(), &querypb.ShowCollectionsRequest{
CollectionIDs: []int64{100},
ResourceGroup: rg,
})
require.NoError(t, err)
require.NoError(t, merr.Error(resp.GetStatus()))
return resp.GetQueryServiceAvailable()[0]
}
assert.True(t, show("rg-a"))
assert.Zero(t, collectionWide, "a scoped request does not compute the collection-wide answer")
assert.True(t, show(""))
assert.Equal(t, 1, collectionWide, "an unscoped request does")
}
internal/querycoordv2/utils/delegator_placement_form_test.go
+164
View File
@@ -0,0 +1,164 @@
// 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 utils
import (
"context"
"testing"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/rgpb"
"github.com/milvus-io/milvus/internal/coordinator/snmanager"
etcdKV "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
. "github.com/milvus-io/milvus/internal/querycoordv2/params"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/internal/util/streamingutil"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/util/etcd"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// installForm turns this test's binary into one with a form installed, and
// turns it back into a stock binary when the test ends.
func installForm(t *testing.T) {
t.Helper()
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
ext.SetForm()
}
// stockBinary makes sure nothing is installed, so the test speaks for a stock
// deployment whatever an earlier test left behind.
func stockBinary(t *testing.T) {
t.Helper()
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
}
// withStreamingQueryNodes turns the streaming service on and makes the
// streaming node manager answer from byRG, so a test can state which resource
// group holds which streaming query nodes without one running.
//
// The cluster-wide set is derived from byRG rather than given separately: the
// two are the same nodes, and letting them disagree would test a state that
// cannot happen. It is needed as well as the per-group map, because
// AssignReplica's first check - the one milvus has always had - refuses a
// replica count above the number of streaming nodes in the whole cluster.
func withStreamingQueryNodes(byRG map[string]typeutil.UniqueSet) func() {
all := typeutil.NewUniqueSet()
for _, nodes := range byRG {
all.Insert(nodes.Collect()...)
}
enabled := mockey.Mock(streamingutil.IsStreamingServiceEnabled).Return(true).Build()
byGroup := mockey.Mock((*snmanager.StreamingNodeManager).GetStreamingQueryNodeIDsByResourceGroup).
Return(byRG).Build()
cluster := mockey.Mock((*snmanager.StreamingNodeManager).GetStreamingQueryNodeIDs).
Return(all).Build()
return func() {
cluster.UnPatch()
byGroup.UnPatch()
enabled.UnPatch()
}
}
// metaWithQueryClusters builds a meta holding one resource group per name,
// each with one regular query node: the shape of a form's query clusters.
func metaWithQueryClusters(t *testing.T, rgNames ...string) (context.Context, *meta.Meta) {
t.Helper()
paramtable.Init()
config := GenerateEtcdConfig()
cli, err := etcd.GetEtcdClient(
config.UseEmbedEtcd.GetAsBool(),
config.EtcdUseSSL.GetAsBool(),
config.Endpoints.GetAsStrings(),
config.EtcdTLSCert.GetValue(),
config.EtcdTLSKey.GetValue(),
config.EtcdTLSCACert.GetValue(),
config.EtcdTLSMinVersion.GetValue())
require.NoError(t, err)
t.Cleanup(func() { cli.Close() })
ctx := context.Background()
nodeMgr := session.NewNodeManager()
m := meta.NewMeta(RandomIncrementIDAllocator(),
querycoord.NewCatalog(etcdKV.NewEtcdKV(cli, config.MetaRootPath.GetValue())),
nodeMgr)
for i, rgName := range rgNames {
_, err = m.AddResourceGroup(ctx, rgName, &rgpb.ResourceGroupConfig{
Requests: &rgpb.ResourceGroupLimit{NodeNum: 1},
Limits: &rgpb.ResourceGroupLimit{NodeNum: 1},
})
require.NoError(t, err)
nodeID := int64(100 + i)
nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: nodeID, Address: "localhost", Hostname: "localhost"}))
m.HandleNodeUp(ctx, nodeID)
nodes, err := m.GetNodes(ctx, rgName)
require.NoError(t, err)
require.Len(t, nodes, 1, "each query cluster holds its one regular query node")
}
return ctx, m
}
// A form's instance runs ONE streaming node, for DDL and the write ahead log,
// and serves queries from resource groups of regular query nodes, several of
// which load the same collection. Milvus's own bound - no more replicas than
// streaming query nodes, because that is where a delegator goes - would refuse
// the second one. Under a form the delegator goes onto the replica's regular
// query nodes, so that bound does not apply and the regular one decides.
func TestAFormLoadsACollectionIntoMoreQueryClustersThanItHasStreamingNodes(t *testing.T) {
ctx, m := metaWithQueryClusters(t, "cluster_a", "cluster_b")
defer withStreamingQueryNodes(map[string]typeutil.UniqueSet{
"internal": typeutil.NewUniqueSet(7), // the instance's only streaming node, in neither cluster
})()
installForm(t)
assignment, err := AssignReplica(ctx, m, []string{"cluster_a", "cluster_b"}, 2, true)
require.NoError(t, err)
assert.Equal(t, map[string]int{"cluster_a": 1, "cluster_b": 1}, assignment)
}
// The same request on a stock binary keeps the refusal it has always had.
func TestAStockBinaryStillBoundsReplicasByStreamingNodes(t *testing.T) {
ctx, m := metaWithQueryClusters(t, "cluster_a", "cluster_b")
defer withStreamingQueryNodes(map[string]typeutil.UniqueSet{
"internal": typeutil.NewUniqueSet(7),
})()
stockBinary(t)
_, err := AssignReplica(ctx, m, []string{"cluster_a", "cluster_b"}, 2, true)
assert.ErrorIs(t, err, merr.ErrServiceResourceInsufficient, "two replicas, one streaming node")
}
// The regular bound still stands under a form: a query cluster with one
// regular query node takes one replica, not two.
func TestAFormStillBoundsAQueryClusterByItsRegularNodes(t *testing.T) {
ctx, m := metaWithQueryClusters(t, "cluster_a")
defer withStreamingQueryNodes(map[string]typeutil.UniqueSet{
"internal": typeutil.NewUniqueSet(7),
})()
installForm(t)
_, err := AssignReplica(ctx, m, []string{"cluster_a"}, 2, true)
assert.ErrorIs(t, err, merr.ErrResourceGroupNodeNotEnough)
}
internal/querycoordv2/utils/meta.go
+24 -9
View File
@@ -106,18 +106,15 @@ func RecoverAllCollection(m *meta.Meta) {
}
}
func AssignReplica(ctx context.Context, m *meta.Meta, resourceGroups []string, replicaNumber int32, checkNodeNum bool) (map[string]int, error) {
// ReplicaNumberByResourceGroup is the layout a load request asks for: how
// many replicas in which resource group. It is the first thing AssignReplica
// decides, and a caller that has to compare the request with what the
// collection already holds before deciding whether to admit it reads the
// same answer from here.
func ReplicaNumberByResourceGroup(resourceGroups []string, replicaNumber int32) (map[string]int, error) {
if len(resourceGroups) != 0 && len(resourceGroups) != 1 && len(resourceGroups) != int(replicaNumber) {
return nil, merr.WrapErrParameterInvalidMsg("replica=[%d] resource group=[%s], resource group num can only be 0, 1 or same as replica number", replicaNumber, strings.Join(resourceGroups, ","))
}
if streamingutil.IsStreamingServiceEnabled() && checkNodeNum {
streamingNodeCount := snmanager.StaticStreamingNodeManager.GetStreamingQueryNodeIDs().Len()
if replicaNumber > int32(streamingNodeCount) {
return nil, merr.WrapErrStreamingNodeNotEnough(streamingNodeCount, int(replicaNumber), fmt.Sprintf("when load %d replica count", replicaNumber))
}
}
replicaNumInRG := make(map[string]int)
if len(resourceGroups) == 0 {
// All replicas should be spawned in default resource group.
@@ -131,6 +128,24 @@ func AssignReplica(ctx context.Context, m *meta.Meta, resourceGroups []string, r
replicaNumInRG[rgName] += 1
}
}
return replicaNumInRG, nil
}
func AssignReplica(ctx context.Context, m *meta.Meta, resourceGroups []string, replicaNumber int32, checkNodeNum bool) (map[string]int, error) {
replicaNumInRG, err := ReplicaNumberByResourceGroup(resourceGroups, replicaNumber)
if err != nil {
return nil, err
}
// A replica needs a streaming query node of its own only when that is
// where its delegator goes; under an installed form it goes onto the
// replica's regular query nodes, and the regular bound below is the bound.
if streamingutil.UseStreamingQueryNodeAsDelegator() && checkNodeNum {
streamingNodeCount := snmanager.StaticStreamingNodeManager.GetStreamingQueryNodeIDs().Len()
if replicaNumber > int32(streamingNodeCount) {
return nil, merr.WrapErrStreamingNodeNotEnough(streamingNodeCount, int(replicaNumber), fmt.Sprintf("when load %d replica count", replicaNumber))
}
}
// TODO: !!!Warning, ResourceManager and ReplicaManager doesn't protected with each other in concurrent operation.
// 1. replica1 got rg1's node snapshot but doesn't spawn finished.
internal/querycoordv2/utils/resource_group_load_percentage.go
+52 -11
View File
@@ -209,6 +209,10 @@ import (
//
// This is a read-only composition over existing querycoord state. It adds no
// proto field, no persisted state, and no change to the replica model.
//
// The figure is the minimum of ReplicaLoadPercentagesByResourceGroup, which
// is where the walk lives; the per-replica split exists so the min is taken
// across one consistent snapshot of the targets and the distribution.
func LoadPercentageByResourceGroup(
ctx context.Context,
m *meta.Meta,
@@ -217,8 +221,47 @@ func LoadPercentageByResourceGroup(
collectionID int64,
rgName string,
) (int32, error) {
figures, err := ReplicaLoadPercentagesByResourceGroup(ctx, m, targetMgr, dist, collectionID, rgName)
if err != nil {
return -1, err
}
return MinReplicaLoadPercentage(figures), nil
}
// MinReplicaLoadPercentage folds per-replica figures into the group's figure:
// the laggard, for the reason given under "Multiple replicas" on
// LoadPercentageByResourceGroup. No figures at all is the "no replica here"
// outcome, -1.
func MinReplicaLoadPercentage(figures map[int64]int32) int32 {
if len(figures) == 0 {
return -1
}
percentage := int32(100)
for _, figure := range figures {
if figure < percentage {
percentage = figure
}
}
return percentage
}
// ReplicaLoadPercentagesByResourceGroup is LoadPercentageByResourceGroup
// before the fold: one figure per selected replica, keyed by replica ID,
// every one measured against the same snapshot of the targets and the
// distribution. The outcomes and the error contract are those of
// LoadPercentageByResourceGroup, with "-1, nil error" spelled as an empty map:
// a collection that is not registered as loaded, or a group holding no
// replica of it, has nothing to report a figure for.
func ReplicaLoadPercentagesByResourceGroup(
ctx context.Context,
m *meta.Meta,
targetMgr meta.TargetManagerInterface,
dist *meta.DistributionManager,
collectionID int64,
rgName string,
) (map[int64]int32, error) {
if m == nil || targetMgr == nil || dist == nil {
return -1, merr.WrapErrServiceNotReadyMsg("querycoord read stores are not wired up yet")
return nil, merr.WrapErrServiceNotReadyMsg("querycoord read stores are not wired up yet")
}
// Only the scoped form consults the resource manager, so only the scoped
@@ -233,10 +276,10 @@ func LoadPercentageByResourceGroup(
// whatever the collection's state.
if rgName != "" {
if m.ResourceManager == nil {
return -1, merr.WrapErrServiceNotReadyMsg("querycoord resource manager is not wired up yet")
return nil, merr.WrapErrServiceNotReadyMsg("querycoord resource manager is not wired up yet")
}
if !m.ContainResourceGroup(ctx, rgName) {
return -1, merr.WrapErrResourceGroupNotFound(rgName)
return nil, merr.WrapErrResourceGroupNotFound(rgName)
}
}
@@ -284,10 +327,10 @@ func LoadPercentageByResourceGroup(
// indistinguishable from one that is, and the caller would
// retry until the cache entry expires 24h later. The recorded
// cause is kept in the message.
return -1, merr.WrapErrCollectionNotLoaded(collectionID, err.Error())
return nil, merr.WrapErrCollectionNotLoaded(collectionID, err.Error())
}
}
return -1, nil
return nil, nil
}
// Query-invisible replicas (load-config updates spawn replicas invisible
@@ -305,7 +348,7 @@ func LoadPercentageByResourceGroup(
}
}
if len(replicas) == 0 {
return -1, nil
return nil, nil
}
// The targets are read ONCE for every selected replica, not per replica.
@@ -341,13 +384,11 @@ func LoadPercentageByResourceGroup(
// channels.
delegators := prefetchDelegatorsByChannel(dist, channelTargets, segmentTargets)
percentage := int32(100)
figures := make(map[int64]int32, len(replicas))
for _, replica := range replicas {
if p := replicaLoadPercentage(replica, channelTargets, segmentTargets, delegators); p < percentage {
percentage = p
}
figures[replica.GetID()] = replicaLoadPercentage(replica, channelTargets, segmentTargets, delegators)
}
return percentage, nil
return figures, nil
}
// prefetchDelegatorsByChannel resolves every channel named by the target set
internal/querycoordv2/utils/resource_group_load_percentage_internal_test.go
+11
View File
@@ -119,3 +119,14 @@ func TestReplicaLoadPercentageUsesTheSuppliedSnapshots(t *testing.T) {
assert.EqualValues(t, 0, replicaLoadPercentage(replica, channelTargets, segmentTargets,
map[string][]*meta.DmChannel{}))
}
// TestMinReplicaLoadPercentageIsTheLaggard pins the fold from per-replica
// figures to the group's: the slowest replica, and -1 -- the "no replica
// here" outcome -- when there is nothing to fold.
func TestMinReplicaLoadPercentageIsTheLaggard(t *testing.T) {
assert.EqualValues(t, -1, MinReplicaLoadPercentage(nil))
assert.EqualValues(t, -1, MinReplicaLoadPercentage(map[int64]int32{}))
assert.EqualValues(t, 100, MinReplicaLoadPercentage(map[int64]int32{1: 100}))
assert.EqualValues(t, 33, MinReplicaLoadPercentage(map[int64]int32{1: 100, 2: 33, 3: 66}))
assert.EqualValues(t, 0, MinReplicaLoadPercentage(map[int64]int32{1: 0, 2: 100}))
}
internal/querynodev2/compiled_query_hook_test.go
+177
View File
@@ -0,0 +1,177 @@
// 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 querynodev2
import (
"context"
"testing"
"github.com/bytedance/mockey"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus/internal/util/hookutil"
"github.com/milvus-io/milvus/internal/util/searchutil/optimizers"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// recordingQueryHook is a compiled-in query hook that remembers how it was
// initialized and can refuse either step.
type recordingQueryHook struct {
initErr, tuningErr error
initWith []string
tuningWith []map[string]string
}
func (h *recordingQueryHook) Run(map[string]any) error { return nil }
func (h *recordingQueryHook) Init(s string) error {
h.initWith = append(h.initWith, s)
return h.initErr
}
func (h *recordingQueryHook) InitTuningConfig(m map[string]string) error {
h.tuningWith = append(h.tuningWith, m)
return h.tuningErr
}
func (h *recordingQueryHook) DeleteTuningConfig(string) error { return nil }
func (h *recordingQueryHook) CalculateEffectiveSegmentNum([]int64, int64) int { return 0 }
func installQueryHook(t *testing.T, h extension.QueryHook) *QueryNode {
t.Helper()
paramtable.Init()
extension.ResetForTest()
t.Cleanup(extension.ResetForTest)
extension.SetQueryHook(h)
return &QueryNode{ctx: context.Background()}
}
// saveQueryNodeKey writes one configuration key and resets it when the test
// ends. Save and Reset both dispatch to the config watchers synchronously, on
// this goroutine, so a watcher initHook registered sees the reset only after
// the test's assertions have run.
func saveQueryNodeKey(t *testing.T, key, value string) {
t.Helper()
p := paramtable.Get()
require.NoError(t, p.Save(key, value))
t.Cleanup(func() { p.Reset(key) })
}
func TestInitHookUsesTheCompiledInQueryHook(t *testing.T) {
h := &recordingQueryHook{}
node := installQueryHook(t, h)
saveQueryNodeKey(t, paramtable.Get().AutoIndexConfig.AutoIndexSearchConfig.Key, `{"level": 2}`)
require.NoError(t, node.initHook())
assert.Same(t, h, node.queryHook)
assert.Equal(t, []string{`{"level": 2}`}, h.initWith,
"the compiled-in hook is initialized with autoIndex.params.search, as a plug-in is")
require.Len(t, h.tuningWith, 1, "and its tuning configuration, as a plug-in is")
assert.Equal(t, paramtable.Get().AutoIndexConfig.AutoIndexTuningConfig.GetValue(), h.tuningWith[0])
}
// Two tuners for the same search is a deployment mistake, and it is reported
// rather than silently resolved by start-up order. It is also the one hook
// failure that is fatal whatever autoIndex.enable says, so it carries a
// distinguishable error.
func TestInitHookRefusesACompiledInQueryHookBesideAPlugin(t *testing.T) {
node := installQueryHook(t, &recordingQueryHook{})
saveQueryNodeKey(t, paramtable.Get().QueryNodeCfg.SoPath.Key, "/tmp/some-tuner.so")
err := node.initHook()
require.Error(t, err)
assert.Contains(t, err.Error(), "only one can")
assert.True(t, errors.Is(err, errQueryHookConflict),
"the conflict must be recognizable, or Init cannot make it fatal on its own")
assert.Nil(t, node.queryHook)
}
// The conflict stops the QueryNode whatever autoIndex.enable is; every other
// hook failure is fatal only when auto index is on.
func TestHookInitIsFatal(t *testing.T) {
saveQueryNodeKey(t, paramtable.Get().AutoIndexConfig.Enable.Key, "false")
assert.True(t, hookInitIsFatal(errQueryHookConflict),
"a configured-both-ways hook is fatal even with auto index off")
assert.False(t, hookInitIsFatal(errors.New("fail to set the plugin path")),
"an absent hook is not fatal when auto index is off")
saveQueryNodeKey(t, paramtable.Get().AutoIndexConfig.Enable.Key, "true")
assert.True(t, hookInitIsFatal(errors.New("fail to set the plugin path")),
"with auto index on, any hook failure is fatal")
}
// With nothing compiled in - a stock binary - nothing changes: an empty
// queryNode.soPath is still the error it always was.
func TestInitHookWithoutACompiledInQueryHookIsUnchanged(t *testing.T) {
node := installQueryHook(t, nil)
err := node.initHook()
require.Error(t, err)
assert.Contains(t, err.Error(), "fail to set the plugin path")
assert.Nil(t, node.queryHook)
}
// A queryNode.soPath that does not load is still the load error it always
// was: the plug-in branch is master's, untouched.
func TestInitHookWithoutACompiledInQueryHookStillLoadsThePlugin(t *testing.T) {
node := installQueryHook(t, nil)
saveQueryNodeKey(t, paramtable.Get().QueryNodeCfg.SoPath.Key, "/nonexistent/some-tuner.so")
err := node.initHook()
require.Error(t, err)
assert.Contains(t, err.Error(), "/nonexistent/some-tuner.so")
assert.Nil(t, node.queryHook)
}
// A plug-in that loads is initialized and installed through the same path a
// compiled-in hook takes, which is what makes the two interchangeable.
func TestInitHookInitializesALoadedPluginAsItAlwaysDid(t *testing.T) {
h := &recordingQueryHook{}
node := installQueryHook(t, nil)
saveQueryNodeKey(t, paramtable.Get().QueryNodeCfg.SoPath.Key, "/tmp/some-tuner.so")
loaded := mockey.Mock(hookutil.LoadPlugin[optimizers.QueryHook], mockey.OptGeneric).
Return(h, nil).Build()
defer loaded.UnPatch()
require.NoError(t, node.initHook())
assert.Same(t, h, node.queryHook)
assert.Equal(t, []string{paramtable.Get().AutoIndexConfig.AutoIndexSearchConfig.GetValue()}, h.initWith,
"a plug-in is initialized with the same autoIndex.params.search the compiled-in hook receives")
require.Len(t, h.tuningWith, 1)
assert.Equal(t, paramtable.Get().AutoIndexConfig.AutoIndexTuningConfig.GetValue(), h.tuningWith[0])
}
func TestInitHookReportsACompiledInQueryHookThatCannotInitialize(t *testing.T) {
t.Run("search config", func(t *testing.T) {
node := installQueryHook(t, &recordingQueryHook{initErr: errors.New("bad search config")})
err := node.initHook()
require.Error(t, err)
assert.Contains(t, err.Error(), "bad search config")
assert.Nil(t, node.queryHook, "a hook that cannot initialize is not installed")
})
t.Run("tuning config", func(t *testing.T) {
node := installQueryHook(t, &recordingQueryHook{tuningErr: errors.New("bad tuning config")})
err := node.initHook()
require.Error(t, err)
assert.Contains(t, err.Error(), "bad tuning config")
assert.Nil(t, node.queryHook)
})
}
internal/querynodev2/server.go
+52 -6
View File
@@ -38,6 +38,7 @@ import (
"sync"
"time"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"github.com/tidwall/gjson"
clientv3 "go.etcd.io/etcd/client/v3"
@@ -67,6 +68,7 @@ import (
"github.com/milvus-io/milvus/internal/util/streamingutil/util"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/config"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/mq/msgdispatcher"
@@ -341,8 +343,12 @@ func (node *QueryNode) Init() error {
err = node.initHook()
if err != nil {
// auto index cannot work if hook init failed
if paramtable.Get().AutoIndexConfig.Enable.GetAsBool() {
// A hook configured both ways is a contradiction in the
// deployment, not an optional plug-in that is absent: refuse to
// start whatever autoIndex.enable says. Any other failure keeps
// the old rule - auto index cannot work without a hook, so it is
// fatal only when auto index is on.
if hookInitIsFatal(err) {
mlog.Error(node.ctx, "QueryNode init hook failed", mlog.Err(err))
initError = err
return
@@ -634,9 +640,43 @@ func (node *QueryNode) SetAddress(address string) {
node.address = address
}
// initHook initializes parameter tuning hook.
// errQueryHookConflict is returned when a compiled-in query hook and
// queryNode.soPath are both configured. A hook that fails to load is an
// optional plug-in that may be absent, but this is a contradiction in the
// deployment, so Init refuses to start on it whatever autoIndex.enable says.
var errQueryHookConflict = errors.New("queryNode.soPath and a compiled-in query hook are both configured")
// hookInitIsFatal reports whether an initHook error must stop the QueryNode.
// A queryNode.soPath configured beside a compiled-in hook is a contradiction
// and is always fatal; any other failure - no path at all, a plug-in that
// does not load - is fatal only when auto index is on, because that is the
// only feature that needs a hook.
func hookInitIsFatal(err error) bool {
return errors.Is(err, errQueryHookConflict) || paramtable.Get().AutoIndexConfig.Enable.GetAsBool()
}
// initHook installs the parameter tuning hook: the one compiled into this
// binary if a distribution installed one, otherwise the plug-in at
// queryNode.soPath.
//
// A compiled-in hook is used in preference to the plug-in, and a deployment
// that configures both is refused (errQueryHookConflict): both would tune
// every search, and picking silently would make which one wins depend on
// start-up order rather than on the deployment. It is otherwise treated
// exactly as the plug-in is: the same two Init calls with the autoIndex
// configuration, the same watchers.
func (node *QueryNode) initHook() error {
path := paramtable.Get().QueryNodeCfg.SoPath.GetValue()
if compiled := extension.InstalledQueryHook(); compiled != nil {
if path != "" {
return merr.Wrapf(errQueryHookConflict,
"queryNode.soPath is set to %q and a query hook is also compiled in; "+
"both tune every search, and only one can", path)
}
return node.useQueryHook(compiled)
}
if path == "" {
return merr.WrapErrServiceInternalMsg("fail to set the plugin path")
}
@@ -645,17 +685,23 @@ func (node *QueryNode) initHook() error {
if err != nil {
return err
}
return node.useQueryHook(hoo)
}
if err = hoo.Init(paramtable.Get().AutoIndexConfig.AutoIndexSearchConfig.GetValue()); err != nil {
// useQueryHook initializes a query hook with the autoIndex configuration,
// installs it and watches that configuration on its behalf. It is the same
// for a plug-in and a compiled-in hook, which is the point: a hook a
// distribution compiled in must not be a second-class one.
func (node *QueryNode) useQueryHook(hoo optimizers.QueryHook) error {
if err := hoo.Init(paramtable.Get().AutoIndexConfig.AutoIndexSearchConfig.GetValue()); err != nil {
return merr.Wrap(err, "fail to init configs for the hook")
}
if err = hoo.InitTuningConfig(paramtable.Get().AutoIndexConfig.AutoIndexTuningConfig.GetValue()); err != nil {
if err := hoo.InitTuningConfig(paramtable.Get().AutoIndexConfig.AutoIndexTuningConfig.GetValue()); err != nil {
return merr.Wrap(err, "fail to init tuning configs for the hook")
}
node.queryHook = hoo
node.handleQueryHookEvent()
return nil
}
internal/util/hookutil/cipher.go
+35 -5
View File
@@ -33,6 +33,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/hook"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/config"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/indexcgopb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
@@ -422,6 +423,28 @@ func initCipher() error {
pathGo := paramtable.GetCipherParams().SoPathGo.GetValue()
pathCpp := paramtable.GetCipherParams().SoPathCpp.GetValue()
// A cipher compiled into this binary replaces the Go half of the plug-in
// pair, cipherPlugin.soPathGo, and only that half: the C++ half is still
// loaded by the core from cipherPlugin.soPathCpp, so that path stays the
// deployment's declaration that encryption is on. A compiled-in cipher
// with no soPathCpp is left idle, exactly as no plug-in would be loaded.
// Configuring soPathGo beside a compiled-in cipher is refused: both would
// answer for the same keys, and picking silently would make which one
// wins depend on start-up order rather than on the deployment.
if compiled := ext.InstalledCipher(); compiled != nil {
if pathGo != "" {
return merr.WrapErrServiceInternalMsg(
"hookutil: cipherPlugin.soPathGo is set to %q and a cipher is also compiled in; "+
"both answer for the encryption keys, and only one can", pathGo)
}
if pathCpp == "" {
mlog.Info(context.TODO(), "empty cipherPlugin.soPathCpp, encryption is off and the compiled-in cipher stays idle")
return nil
}
return useCipher(compiled, "compiled-in cipher")
}
if pathGo == "" || pathCpp == "" {
mlog.Info(context.TODO(), "empty so path for cipher plugin, skip to load plugin")
return nil
@@ -431,14 +454,21 @@ func initCipher() error {
if err != nil {
return err
}
return useCipher(cipherVal, "cipher plugin")
}
initConfigs := buildCipherInitConfig()
if err = cipherVal.Init(initConfigs); err != nil {
return merr.Wrap(err, "fail to init configs for the cipher plugin")
// useCipher initializes a cipher with the cipherPlugin.* configuration,
// registers the reload callbacks and installs it. It is the same for a
// plug-in and a compiled-in cipher, which is the point: a cipher a
// distribution compiled in must not be a second-class one. origin names it
// in the error and the log.
func useCipher(cipherVal hook.Cipher, origin string) error {
if err := cipherVal.Init(buildCipherInitConfig()); err != nil {
return merr.Wrap(err, "fail to init configs for the "+origin)
}
registerCallback()
storeCipher((cipherVal))
storeCipher(cipherVal)
mlog.Info(context.TODO(), "using the "+origin)
return nil
}
internal/util/hookutil/compiled_cipher_test.go
+179
View File
@@ -0,0 +1,179 @@
/*
* 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 hookutil
import (
"context"
"sync"
"testing"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/hook"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// recordingCipher is a compiled-in cipher that remembers how often it was
// initialized and can refuse. The reload callback re-initializes it from a
// configuration event, so the count is taken under a lock.
type recordingCipher struct {
testCipher
initErr error
mu sync.Mutex
inits int
}
func (c *recordingCipher) Init(map[string]string) error {
c.mu.Lock()
defer c.mu.Unlock()
c.inits++
return c.initErr
}
func (c *recordingCipher) initCount() int {
c.mu.Lock()
defer c.mu.Unlock()
return c.inits
}
var _ hook.Cipher = (*recordingCipher)(nil)
// installCipher makes this test's binary one with a cipher compiled in, on a
// cipher state no earlier test has touched.
func installCipher(t *testing.T, c hook.Cipher) {
t.Helper()
paramtable.Init()
initCipherOnce = sync.Once{}
storeCipher(nil)
ext.ResetForTest()
t.Cleanup(func() {
ext.ResetForTest()
initCipherOnce = sync.Once{}
storeCipher(nil)
})
ext.SetCipher(c)
}
// saveCipherKey writes one of the two cipherPlugin.so* path keys and clears
// it when the test ends. Both default to empty, which is what clearing
// restores; a key with a non-empty default would need its value read first.
func saveCipherKey(t *testing.T, key, value string) {
t.Helper()
cp := paramtable.GetCipherParams()
require.NoError(t, cp.Save(key, value))
t.Cleanup(func() { _ = cp.Save(key, "") })
}
// installedCipher reads the stored cipher without going through GetCipher,
// whose InitOnceCipher would run initCipher a second time after a test that
// called it directly - and would panic on the error those tests are about.
func installedCipher() hook.Cipher {
return Cipher.Load().(cipherContainer).cipher
}
func TestInitCipherUsesTheCompiledInCipher(t *testing.T) {
c := &recordingCipher{}
installCipher(t, c)
saveCipherKey(t, paramtable.GetCipherParams().SoPathCpp.Key, "/tmp/some-cipher.so")
InitOnceCipher()
assert.Same(t, c, GetCipher())
assert.True(t, IsClusterEncryptionEnabled(), "a compiled-in cipher with the C++ half configured is encryption on")
assert.Equal(t, 1, c.initCount(), "initialized once, with the cipherPlugin.* configuration, as a plug-in is")
}
// Two answers for the same keys is a deployment mistake, and it is reported
// rather than silently resolved by start-up order.
func TestInitCipherRefusesACompiledInCipherBesideAPlugin(t *testing.T) {
installCipher(t, &recordingCipher{})
saveCipherKey(t, paramtable.GetCipherParams().SoPathGo.Key, "/tmp/some-cipher-go.so")
saveCipherKey(t, paramtable.GetCipherParams().SoPathCpp.Key, "/tmp/some-cipher.so")
err := initCipher()
require.Error(t, err)
assert.Contains(t, err.Error(), "only one can")
assert.Nil(t, installedCipher())
}
// The C++ half is still a file the core loads from cipherPlugin.soPathCpp, so
// that path stays the declaration that encryption is on. Without it the
// compiled-in cipher is left idle, as no plug-in would be loaded.
func TestInitCipherLeavesTheCompiledInCipherIdleWithoutTheCppHalf(t *testing.T) {
c := &recordingCipher{}
installCipher(t, c)
InitOnceCipher()
assert.Nil(t, GetCipher())
assert.False(t, IsClusterEncryptionEnabled())
assert.Equal(t, 0, c.initCount(), "an idle cipher is not initialized: it has no KMS configuration to initialize with")
}
func TestInitCipherReportsACompiledInCipherThatCannotInitialize(t *testing.T) {
installCipher(t, &recordingCipher{initErr: errors.New("no kms key")})
saveCipherKey(t, paramtable.GetCipherParams().SoPathCpp.Key, "/tmp/some-cipher.so")
err := initCipher()
require.Error(t, err)
assert.Contains(t, err.Error(), "no kms key")
assert.Contains(t, err.Error(), "compiled-in cipher")
assert.Nil(t, installedCipher(), "a cipher that cannot initialize is not installed")
}
// With nothing compiled in - a stock binary - nothing changes: no so paths,
// no cipher.
func TestInitCipherWithoutACompiledInCipherIsUnchanged(t *testing.T) {
installCipher(t, nil)
InitOnceCipher()
assert.Nil(t, GetCipher())
}
// A compiled-in cipher is reconfigured exactly as a plug-in is: a
// cipherPlugin.* edit re-initializes it.
func TestReloadCipherConfigReachesTheCompiledInCipher(t *testing.T) {
c := &recordingCipher{}
installCipher(t, c)
saveCipherKey(t, paramtable.GetCipherParams().SoPathCpp.Key, "/tmp/some-cipher.so")
InitOnceCipher()
require.Equal(t, 1, c.initCount())
require.NoError(t, reloadCipherConfig(context.Background(), paramtable.GetCipherParams().DefaultRootKey.Key, "", "aws-kms://new"))
assert.Equal(t, 2, c.initCount())
}
// And with nothing compiled in, cipherPlugin.soPathGo is still loaded as a
// plug-in: the compiled-in branch must not swallow the path a stock binary
// takes. A file that is not there is the nearest thing to a plug-in a test
// can load.
func TestInitCipherWithoutACompiledInCipherStillLoadsThePlugin(t *testing.T) {
installCipher(t, nil)
saveCipherKey(t, paramtable.GetCipherParams().SoPathGo.Key, "/tmp/no-such-cipher-go.so")
saveCipherKey(t, paramtable.GetCipherParams().SoPathCpp.Key, "/tmp/no-such-cipher.so")
err := initCipher()
require.Error(t, err)
assert.Contains(t, err.Error(), "fail to open plugin")
assert.Nil(t, installedCipher())
}
internal/util/hookutil/compiled_hook_test.go
+353
View File
@@ -0,0 +1,353 @@
/*
* 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 hookutil
import (
"context"
"sync"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/hook"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// saveHookKey writes one hook.* configuration key and clears it again when the
// test ends, so that no test inherits the configuration another one left
// behind and the order they run in does not matter.
func saveHookKey(t *testing.T, key, value string) {
t.Helper()
hp := paramtable.GetHookParams()
require.NoError(t, hp.Save(key, value))
t.Cleanup(func() { _ = hp.Save(key, "") })
}
// hookProvider is a form that supplies nothing but a hook, which is the
// smallest thing a distribution can do to take over the request path.
func installHook(t *testing.T, h hook.Hook) {
t.Helper()
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
ext.SetHook(h)
}
func TestInitHookUsesTheCompiledInHook(t *testing.T) {
paramtable.Init()
installHook(t, MockAPIHook{User: "root"})
require.NoError(t, initHook())
got, err := GetHook().VerifyAPIKey("whatever")
assert.NoError(t, err)
assert.Equal(t, "root", got)
}
// With no provider - a stock binary - nothing changes: the default hook, and
// then whatever proxy.soPath says.
func TestInitHookWithoutACompiledInHookIsUnchanged(t *testing.T) {
paramtable.Init()
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
require.NoError(t, initHook())
_, ok := GetHook().(DefaultHook)
assert.True(t, ok, "a stock binary keeps the default hook")
}
// A provider that fills in no hook is a form that does not want the request
// path, and it must not displace a plug-in or the default.
func TestInitHookIgnoresANilHook(t *testing.T) {
paramtable.Init()
installHook(t, nil)
require.NoError(t, initHook())
_, ok := GetHook().(DefaultHook)
assert.True(t, ok)
}
// Two authorities for the same question is a deployment mistake, and it is
// reported rather than silently resolved by start-up order.
func TestInitHookRefusesACompiledInHookBesideAPlugin(t *testing.T) {
paramtable.Init()
installHook(t, MockAPIHook{User: "root"})
p := paramtable.Get()
require.NoError(t, p.Save(p.ProxyCfg.SoPath.Key, "/tmp/some-hook.so"))
t.Cleanup(func() { p.Reset(p.ProxyCfg.SoPath.Key) })
err := initHook()
require.Error(t, err)
assert.Contains(t, err.Error(), "only one can")
}
// initRecordingHook is a compiled-in hook that remembers how it was
// initialized, and can refuse. The config-reload watcher re-initializes the
// hook from its own goroutine, so the record is taken under a lock.
type initRecordingHook struct {
MockAPIHook
initErr error
mu sync.Mutex
params map[string]string
inits int
}
func (h *initRecordingHook) Init(params map[string]string) error {
h.mu.Lock()
defer h.mu.Unlock()
h.inits++
h.params = params
return h.initErr
}
// failInitsWith makes every later Init refuse, which is what an operator
// editing a hook.* key to a value the hook cannot accept looks like from here.
func (h *initRecordingHook) failInitsWith(err error) {
h.mu.Lock()
defer h.mu.Unlock()
h.initErr = err
}
// initCount and initParams read what Init recorded.
func (h *initRecordingHook) initCount() int {
h.mu.Lock()
defer h.mu.Unlock()
return h.inits
}
func (h *initRecordingHook) initParams() map[string]string {
h.mu.Lock()
defer h.mu.Unlock()
return h.params
}
// storeAwareHook is an initRecordingHook that also counts the Init calls made
// before it was stored. Saving a hook.* key re-initializes whichever hook is
// installed, from the config watcher's own goroutine, and the watcher is
// registered once for the whole package: a key a test saves before initHook
// can therefore reach this hook through a watcher an earlier test registered,
// after initHook has stored it. Only the calls made before it is stored are
// initHook's.
type storeAwareHook struct {
initRecordingHook
beforeStore int
}
func (h *storeAwareHook) Init(params map[string]string) error {
// Read the stored hook directly. GetHook would run InitOnceHook, and on a
// package state no earlier test initialized that re-enters the same
// sync.Once that is calling this Init, which never returns.
container, _ := hoo.Load().(hookContainer)
stored, _ := container.hook.(*storeAwareHook)
h.mu.Lock()
if stored != h {
h.beforeStore++
}
h.mu.Unlock()
return h.initRecordingHook.Init(params)
}
func (h *storeAwareHook) initsBeforeStore() int {
h.mu.Lock()
defer h.mu.Unlock()
return h.beforeStore
}
// A compiled-in hook is initialized the way a plug-in is: once, with the hook
// configuration, before it is stored. It is the one call that tells the hook
// it runs in the proxy process.
func TestInitHookInitializesTheCompiledInHookWithTheHookConfig(t *testing.T) {
paramtable.Init()
saveHookKey(t, "somekey", "someValue")
h := &storeAwareHook{initRecordingHook: initRecordingHook{MockAPIHook: MockAPIHook{User: "root"}}}
installHook(t, h)
require.NoError(t, initHook())
assert.Equal(t, 1, h.initsBeforeStore(), "initialized exactly once, before it is stored")
assert.Equal(t, "someValue", h.initParams()["somekey"], "the hook sees the hook.* configuration, as a plug-in does")
assert.Same(t, h, GetHook(), "the initialized hook is the one stored")
}
// A hook that cannot initialize is a proxy that does not start, exactly as
// for a plug-in.
func TestInitHookFailsWhenTheCompiledInHookCannotInitialize(t *testing.T) {
paramtable.Init()
h := &initRecordingHook{initErr: errors.New("the internal port is taken")}
installHook(t, h)
err := initHook()
require.Error(t, err)
assert.ErrorContains(t, err, "the internal port is taken")
_, isDefault := GetHook().(DefaultHook)
assert.True(t, isDefault, "a hook that failed to initialize is not stored")
}
// A compiled-in hook is reconfigured the way a plug-in is too: editing a
// hook.* key re-initializes the hook that is installed, whichever way it got
// there. Without the watcher, a form compiled into the binary would keep the
// configuration it was started with forever while a plug-in picked the change
// up, and the two would answer the same config edit differently.
func TestConfigChangeReinitializesTheCompiledInHook(t *testing.T) {
paramtable.Init()
h := &storeAwareHook{initRecordingHook: initRecordingHook{MockAPIHook: MockAPIHook{User: "root"}}}
installHook(t, h)
require.NoError(t, initHook())
require.Equal(t, 1, h.initsBeforeStore(), "initHook initializes the hook once, before it is stored")
saveHookKey(t, "reloadedkey", "reloadedValue")
assert.Eventually(t, func() bool {
return h.initCount() > 1 && h.initParams()["reloadedkey"] == "reloadedValue"
}, 10*time.Second, 10*time.Millisecond,
"a hook.* config edit must re-initialize the compiled-in hook with the new configuration")
assert.Same(t, h, GetHook(), "the re-initialized hook is the one that stays installed")
}
// A hook.* edit the compiled-in hook refuses must not take the proxy down: it
// is already serving, and the edit can be made at any moment. The refusal is
// reported and the configuration that was working stays in place. (Start-up is
// the other way round, and is covered above: a hook that cannot initialize
// there is a proxy that does not start.)
func TestAConfigChangeRefusedByTheCompiledInHookKeepsTheProxyUp(t *testing.T) {
paramtable.Init()
h := &initRecordingHook{MockAPIHook: MockAPIHook{User: "root"}}
installHook(t, h)
require.NoError(t, initHook())
// Start-up is over before any edit arrives, as it is in a running proxy:
// the refusal below has to be judged on the refresh path alone.
require.Same(t, h, GetHook())
initsBeforeTheEdit := h.initCount()
h.failInitsWith(errors.New("hook.someKey is not a duration"))
saveHookKey(t, "refusedkey", "nonsense")
assert.Eventually(t, func() bool { return h.initCount() > initsBeforeTheEdit },
10*time.Second, 10*time.Millisecond,
"the refused edit must still have been offered to the hook")
assert.Same(t, h, GetHook(), "the hook that refused the new configuration stays installed")
h.failInitsWith(nil)
}
// reportingHook is a compiled-in hook that is also the form's hook.Extension,
// which is how a compiled-in form ships the second half a plug-in exports as a
// separate symbol.
type reportingHook struct {
MockAPIHook
reports int
}
func (h *reportingHook) Report(any) int { h.reports++; return h.reports }
func (h *reportingHook) ReportAction(context.Context, interface{}, interface{}, error, string, string) error {
h.reports++
return nil
}
var _ hook.Extension = (*reportingHook)(nil)
// A compiled-in hook that implements hook.Extension is stored as the extension
// too, so the DML, DQL and authorization paths' Report/ReportAction reach it
// exactly as they reach a plug-in's MilvusExtension symbol.
func TestInitHookStoresTheCompiledInHooksExtensionHalf(t *testing.T) {
paramtable.Init()
h := &reportingHook{MockAPIHook: MockAPIHook{User: "root"}}
installHook(t, h)
require.NoError(t, initHook())
assert.Same(t, h, GetExtension(), "the compiled-in hook is the extension when it implements one")
GetExtension().Report(nil)
assert.Equal(t, 1, h.reports, "a report must reach the form, not a default that drops it")
}
// A compiled-in hook that is only a hook.Hook leaves the default extension in
// place, which drops reports as a stock binary does.
func TestInitHookWithoutAnExtensionKeepsTheDefault(t *testing.T) {
paramtable.Init()
installHook(t, MockAPIHook{User: "root"})
require.NoError(t, initHook())
_, isDefault := GetExtension().(DefaultExtension)
assert.True(t, isDefault)
}
// common.panicWhenPluginFail lets a deployment carry on without a plug-in that
// failed to load. A compiled-in hook is treated by two rules. Setting
// proxy.soPath beside it is a contradiction in the deployment - both answer
// VerifyAPIKey and the request interception - so it stops the proxy whatever
// the setting says, form or not. Any other failure of a compiled-in hook stops
// the proxy only when a form is installed: the distribution switched the
// coordinators' behaviors on too, so serving through the default hook would
// run half of it. A plug-in's failure keeps the setting's meaning
// (TestHookInitLogError).
func TestInitOnceHookIsFatalForACompiledInHookWhateverPanicWhenPluginFailSays(t *testing.T) {
paramtable.Init()
p := paramtable.Get()
require.NoError(t, p.Save(p.CommonCfg.PanicWhenPluginFail.Key, "false"))
t.Cleanup(func() { p.Reset(p.CommonCfg.PanicWhenPluginFail.Key) })
// Each case leaves initOnce done, as the plug-in cases in hook_test.go do:
// a Do whose function panics counts as done. Resetting it here instead
// would make the next GetHook anywhere in the package - a config watcher's
// goroutine included - initialize the hook again behind another test.
t.Run("it cannot initialize", func(t *testing.T) {
installHook(t, &initRecordingHook{initErr: errors.New("the internal port is taken")})
ext.SetForm()
assert.Panics(t, func() {
initOnce = sync.Once{}
InitOnceHook()
})
})
// No form here: the soPath conflict is fatal on its own, not by way of the
// form rule.
t.Run("it is configured beside a plug-in", func(t *testing.T) {
installHook(t, MockAPIHook{User: "root"})
require.NoError(t, p.Save(p.ProxyCfg.SoPath.Key, "/tmp/some-hook.so"))
t.Cleanup(func() { p.Reset(p.ProxyCfg.SoPath.Key) })
assert.Panics(t, func() {
initOnce = sync.Once{}
InitOnceHook()
})
})
}
// A distribution that installs a hook WITHOUT declaring a form has switched
// nothing in the coordinators, so a hook that cannot initialize is a plug-in's
// failure: it follows panicWhenPluginFail instead of always stopping the
// proxy. (A soPath beside it is still fatal - see the test above.)
func TestInitOnceHookFollowsPanicWhenPluginFailWithoutAForm(t *testing.T) {
paramtable.Init()
p := paramtable.Get()
require.NoError(t, p.Save(p.CommonCfg.PanicWhenPluginFail.Key, "false"))
t.Cleanup(func() { p.Reset(p.CommonCfg.PanicWhenPluginFail.Key) })
installHook(t, &initRecordingHook{initErr: errors.New("the internal port is taken")})
assert.NotPanics(t, func() {
initOnce = sync.Once{}
InitOnceHook()
})
}
internal/util/hookutil/hook.go
+126 -17
View File
@@ -16,6 +16,17 @@
* limitations under the License.
*/
// Package hookutil installs and serves the proxy's request hook: the extension
// seam a deployment form uses to answer, inspect or refuse an RPC without
// forking milvus. A form supplies its hook either compiled into the binary
// (pkg/extension.SetHook) or as a plug-in loaded from proxy.soPath; both are
// initialized with the hook.* configuration, reconfigured when it changes, and
// consulted by the same proxy interceptor (internal/proxy/hook_interceptor.go).
// With no hook installed the default one does nothing and every RPC behaves as
// it always did.
//
// The mechanism, and what a form may and may not do with it, is described in
// docs/design-docs/design_docs/20260831-in_tree_extension_mechanism.md.
package hookutil
import (
@@ -23,10 +34,12 @@ import (
"sync"
"sync/atomic"
"github.com/cockroachdb/errors"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v3/hook"
"github.com/milvus-io/milvus/pkg/v3/config"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -36,6 +49,7 @@ var (
hoo atomic.Value // hook.Hook
extension atomic.Value // hook.Extension
initOnce sync.Once
watchOnce sync.Once
)
// hookContainer is Container to wrap hook.Hook interface
@@ -65,12 +79,59 @@ func storeExtension(ext hook.Extension) {
extension.Store(extensionContainer{extension: ext})
}
// errHookConflict is returned when a hook is compiled in and proxy.soPath is
// also set. Both answer VerifyAPIKey and the request interception, only one
// can, and picking silently would make the winner depend on start-up order.
// It is fatal whatever common.panicWhenPluginFail says: it is a contradiction
// in the deployment, not an optional plug-in that failed to load.
var errHookConflict = errors.New("proxy.soPath and a compiled-in hook are both configured")
func initHook() error {
// setup default hook & extension
storeHook(DefaultHook{})
storeExtension(DefaultExtension{})
path := paramtable.Get().ProxyCfg.SoPath.GetValue()
// A form compiled into this binary installs its hook by filling in the
// Hook capability, which is the same interface a plug-in provides and is
// consulted by the same interceptor. It is not merged with a plug-in:
// both answer VerifyAPIKey and the request interception, only one can,
// and picking silently would make which one wins depend on start-up
// order rather than on the deployment.
if compiled := ext.InstalledHook(); compiled != nil {
if path != "" {
return merr.Wrapf(errHookConflict,
"hookutil: proxy.soPath is set to %q and a hook is also compiled in; "+
"both answer VerifyAPIKey and the request interception, and only one can", path)
}
// The compiled-in hook is initialized exactly as a plug-in is, with the
// same configuration and the same consequence: a hook that cannot
// initialize is a proxy that does not start. It is also the one call
// that tells the hook it lives in the proxy process.
if err := compiled.Init(paramtable.GetHookParams().SoConfig.GetValue()); err != nil {
return merr.Wrap(err, "fail to init configs for the compiled-in hook")
}
storeHook(compiled)
// A plug-in ships its hook.Extension as a second symbol; a compiled-in
// hook ships it, if at all, as the same value implementing both
// interfaces. Store it when it does, or every Report/ReportAction on
// the DML, DQL and authorization paths silently no-ops for the form
// while a plug-in's would be heard.
if reporter, ok := compiled.(hook.Extension); ok {
storeExtension(reporter)
} else {
mlog.Warn(context.TODO(), "the compiled-in hook implements no hook.Extension, reports are dropped")
}
// And it is reconfigured exactly as a plug-in is: a hook.* edit
// re-initializes whichever hook is installed. Skipping this for the
// compiled-in hook would leave it on the configuration it started with
// while a plug-in picked the same edit up.
watchHookConfigReload()
mlog.Info(context.TODO(), "using the compiled-in hook")
return nil
}
if path == "" {
mlog.Info(context.TODO(), "empty so path, skip to load plugin")
return nil
@@ -85,22 +146,7 @@ func initHook() error {
return merr.Wrap(err, "fail to init configs for the hook")
}
storeHook((hookVal))
paramtable.GetHookParams().WatchHookWithPrefix("watch_hook", "", func(event *config.Event) {
mlog.Info(context.TODO(), "receive the hook refresh event",
mlog.String("eventSource", event.EventSource),
mlog.String("eventType", event.EventType))
go func() {
hookVal := GetHook()
soConfig := paramtable.GetHookParams().SoConfig.GetValue()
mlog.Info(context.TODO(), "refresh hook configs", mlog.Int("entries", len(soConfig)))
if err := hookVal.Init(soConfig); err != nil {
// Plugin errors can echo the opaque, sensitive initialization map.
mlog.Panic(context.TODO(), "fail to init configs for the hook when refreshing",
mlog.String("error", config.RedactedValue))
}
storeHook(hookVal)
}()
})
watchHookConfigReload()
extVal, err := LoadPlugin[hook.Extension](path, "MilvusExtension")
if err != nil {
@@ -111,6 +157,59 @@ func initHook() error {
return nil
}
// watchHookConfigReload re-initializes the installed hook whenever a hook.*
// configuration key changes, which is how a hook is reconfigured without
// restarting the proxy.
//
// It reads GetHook() rather than closing over one hook value, so it serves
// whichever hook is installed - a plug-in loaded from proxy.soPath or a hook
// compiled into this binary - and both branches register it for that reason.
//
// Registered at most once per process: the dispatcher appends handlers, and a
// second registration would re-initialize the hook twice per config edit.
func watchHookConfigReload() {
watchOnce.Do(func() {
paramtable.GetHookParams().WatchHookWithPrefix("watch_hook", "", func(event *config.Event) {
mlog.Info(context.TODO(), "receive the hook refresh event",
mlog.String("eventSource", event.EventSource),
mlog.String("eventType", event.EventType))
go func() {
// The installed hook is reconfigured in place, and is not
// stored back afterwards: it is already the installed one, so
// storing it could only overwrite a hook installed while this
// goroutine was running.
hookVal := GetHook()
soConfig := paramtable.GetHookParams().SoConfig.GetValue()
mlog.Info(context.TODO(), "refresh hook configs", mlog.Int("entries", len(soConfig)))
if err := hookVal.Init(soConfig); err != nil {
if ext.InstalledHook() != nil {
// A refusal at start-up is a deployment that does not
// start, which is the plug-in's rule too. A refusal
// HERE is different: the proxy is already serving, and
// the only thing that changed is a configuration edit
// an operator can make at any moment. Killing every
// proxy of a deployment over one bad value is a worse
// answer than keeping the configuration that was
// working and saying loudly that the new one was not
// taken.
// The error is redacted for the reason the plug-in
// path redacts it: an initialization error can echo
// the opaque, sensitive configuration map back.
mlog.Error(context.TODO(),
"fail to init configs for the compiled-in hook when refreshing, the new configuration is refused",
mlog.String("error", config.RedactedValue))
return
}
// The plug-in path keeps the behavior it has always had.
// Plugin errors can echo the opaque, sensitive initialization map.
mlog.Panic(context.TODO(), "fail to init configs for the hook when refreshing",
mlog.String("error", config.RedactedValue))
}
}()
})
})
}
func SetHook(connectionManager any) {
hookVal := GetHook()
if setter, ok := hookVal.(hookSetter); ok {
@@ -125,7 +224,17 @@ func InitOnceHook() {
err := initHook()
if err != nil {
soPath := paramtable.Get().ProxyCfg.SoPath.GetValue()
if paramtable.Get().CommonCfg.PanicWhenPluginFail.GetAsBool() {
// A soPath configured beside a compiled-in hook is a contradiction
// in the deployment, not an optional plug-in that failed to load:
// it is fatal whatever the setting says. Any other failure of a
// compiled-in hook follows the form rule - the distribution that
// installed it switched the coordinators' behaviors on too
// (extension.FormInstalled), so a proxy that carried on through
// the default hook would run half of that distribution, with its
// request policy missing - and a plug-in's failure keeps
// common.panicWhenPluginFail's meaning.
if errors.Is(err, errHookConflict) || ext.FormInstalled() ||
paramtable.Get().CommonCfg.PanicWhenPluginFail.GetAsBool() {
mlog.Panic(context.TODO(), "fail to init hook",
mlog.String("so_path", soPath), mlog.String("error", config.RedactedValue))
}
internal/util/searchutil/optimizers/query_hook.go
+5 -8
View File
@@ -10,6 +10,7 @@ import (
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/extension"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
@@ -19,14 +20,10 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// QueryHook is the interface for search/query parameter optimizer.
type QueryHook interface {
Run(map[string]any) error
Init(string) error
InitTuningConfig(map[string]string) error
DeleteTuningConfig(string) error
CalculateEffectiveSegmentNum(rowCounts []int64, topk int64) int
}
// QueryHook is extension.QueryHook: the tuning hook a queryNode.soPath plug-in
// exports, or the one a distribution compiled in. It lives in pkg/extension so
// a distribution can implement it; every consumer in the tree keeps this name.
type QueryHook = extension.QueryHook
// OptimizeSearchParams optimizes search parameters using the query hook and applies Knowhere search defaults.
// numSegments is the effective segment number, pre-computed by the caller via CalculateEffectiveSegmentNum.
internal/util/streamingutil/env.go
+22
View File
@@ -2,6 +2,8 @@ package streamingutil
import (
"os"
"github.com/milvus-io/milvus/pkg/v3/extension"
)
const MilvusStreamingServiceEnabled = "MILVUS_STREAMING_SERVICE_ENABLED"
@@ -12,6 +14,26 @@ func IsStreamingServiceEnabled() bool {
return os.Getenv(MilvusStreamingServiceEnabled) == "1"
}
// UseStreamingQueryNodeAsDelegator reports whether the query coordinator places
// shard delegators on streaming query nodes - the query node embedded in every
// streaming node - rather than on a replica's regular query nodes.
//
// With the streaming service on, that is what milvus does: a delegator lives
// beside the WAL it reads, a replica needs a streaming query node of its own,
// and a collection can hold no more replicas than there are streaming nodes.
//
// An installed form (extension.FormInstalled) keeps its streaming node for DDL
// and the write ahead log only. Its query clusters are resource groups of
// regular query nodes, several of which load the same collection, while the
// whole instance runs one streaming node; bound by the rule above, the second
// cluster's load would be refused. So under a form the delegators go where
// they went before the streaming service existed: onto the replica's regular
// RW query nodes, which watch the channel and read the WAL remotely. Every
// site that asks this already carries both placements; this only picks one.
func UseStreamingQueryNodeAsDelegator() bool {
return IsStreamingServiceEnabled() && !extension.FormInstalled()
}
// SetStreamingServiceEnabled set the env that indicates whether the streaming service is enabled.
func SetStreamingServiceEnabled() {
err := os.Setenv(MilvusStreamingServiceEnabled, "1")
internal/util/streamingutil/env_test.go
+53
View File
@@ -0,0 +1,53 @@
// 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 streamingutil
import (
"testing"
"github.com/stretchr/testify/assert"
ext "github.com/milvus-io/milvus/pkg/v3/extension"
)
// Delegators live on streaming query nodes exactly when the streaming service
// is on and no form is installed. A form keeps its one streaming node for DDL
// and the write ahead log, and serves queries from regular query nodes.
func TestUseStreamingQueryNodeAsDelegator(t *testing.T) {
for name, tc := range map[string]struct {
streaming, form, want bool
}{
"stock binary, streaming on": {streaming: true, form: false, want: true},
"stock binary, streaming off": {streaming: false, form: false, want: false},
"installed form, streaming on": {streaming: true, form: true, want: false},
"installed form, streaming off": {streaming: false, form: true, want: false},
} {
t.Run(name, func(t *testing.T) {
ext.ResetForTest()
t.Cleanup(ext.ResetForTest)
if tc.form {
ext.SetForm()
}
if tc.streaming {
t.Setenv(MilvusStreamingServiceEnabled, "1")
} else {
t.Setenv(MilvusStreamingServiceEnabled, "")
}
assert.Equal(t, tc.want, UseStreamingQueryNodeAsDelegator())
})
}
}
pkg/extension/cipher.go
+38
View File
@@ -0,0 +1,38 @@
// 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 extension
import "github.com/milvus-io/milvus-proto/go-api/v3/hook"
// SetCipher installs a compiled-in cipher: the Go half of the cipher plug-in
// pair, which hookutil prefers over cipherPlugin.soPathGo and refuses beside
// it. Only that half: the C++ half is still loaded by the core from
// cipherPlugin.soPathCpp, and that path stays the deployment's declaration
// that encryption is on - a compiled-in cipher with no soPathCpp is left
// idle, exactly as no plug-in would be loaded. Call it before milvus starts;
// nil installs nothing. It is independent of SetHook and of FormInstalled.
func SetCipher(c hook.Cipher) {
installedCipher.Store(&cipherBox{cipher: c})
}
// InstalledCipher returns the installed cipher, or nil.
func InstalledCipher() hook.Cipher {
if b := installedCipher.Load(); b != nil {
return b.cipher
}
return nil
}
pkg/extension/coordinator_engine.go
+54
View File
@@ -0,0 +1,54 @@
// 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 extension
import (
"context"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/proto/rootcoordpb"
)
// Coordinator is the whole coordinator as its own clients see it: every
// root, query and data coordinator RPC. An engine reaches the coordinator
// through nothing else, so what it can do is exactly what a proxy can do.
type Coordinator interface {
rootcoordpb.RootCoordClient
querypb.QueryCoordClient
datapb.DataCoordClient
}
// CoordinatorEngine is control-plane machinery a deployment hosts inside the
// coordinator process. The coordinator starts it once this replica becomes
// ACTIVE (a standby never starts it) and stops it on shutdown; between the two
// the engine drives the coordinator through coord.
//
// The lifecycle contract:
//
// - Start is called at most once, and only on an ACTIVE replica.
// - Stop is called at most once, and only after Start was called. A standby
// that shuts down never had its engine started and does not have it
// stopped; an engine should nevertheless treat a Stop without a Start as
// a no-op, since nothing but ordering guarantees it.
// - Stop may be called while Start is still running, because a shutdown does
// not wait for a slow activation. Stop must then make Start return, and
// leave the engine stopped; the two run on different goroutines.
type CoordinatorEngine interface {
Start(ctx context.Context, coord Coordinator) error
Stop() error
}
pkg/extension/coordinator_engine_test.go
+72
View File
@@ -0,0 +1,72 @@
// 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 extension
import (
"context"
"testing"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
)
type fakeCoordinatorEngine struct {
startErr error
stopErr error
seenCtx context.Context
seenCoord Coordinator
stopped bool
}
func (f *fakeCoordinatorEngine) Start(ctx context.Context, coord Coordinator) error {
f.seenCtx = ctx
f.seenCoord = coord
return f.startErr
}
func (f *fakeCoordinatorEngine) Stop() error {
f.stopped = true
return f.stopErr
}
func TestInstalledCoordinatorEngineReceivesTheCallersContext(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
engine := &fakeCoordinatorEngine{}
SetCoordinatorEngine(engine)
type ctxKey struct{}
ctx := context.WithValue(context.Background(), ctxKey{}, "coordinator")
assert.NoError(t, InstalledCoordinatorEngine().Start(ctx, nil))
assert.Equal(t, "coordinator", engine.seenCtx.Value(ctxKey{}),
"the context must reach the implementation unchanged")
}
func TestCoordinatorEngineLifecycleErrorsArePropagated(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
wantStartErr := errors.New("engine failed to start")
wantStopErr := errors.New("engine failed to stop")
engine := &fakeCoordinatorEngine{startErr: wantStartErr, stopErr: wantStopErr}
SetCoordinatorEngine(engine)
assert.ErrorIs(t, InstalledCoordinatorEngine().Start(context.Background(), nil), wantStartErr,
"an error from Start must survive install and the call unwrapped and unreplaced")
assert.ErrorIs(t, InstalledCoordinatorEngine().Stop(), wantStopErr,
"an error from Stop must survive install and the call unwrapped and unreplaced")
assert.True(t, engine.stopped)
}
pkg/extension/extension.go
+106
View File
@@ -0,0 +1,106 @@
// 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 extension is what a distribution that compiles its own behavior into
// the milvus binary installs at boot: a request hook, a query hook, a cipher,
// and an engine the coordinator runs while it is active. Everything else such
// a distribution needs is either the hook's own reach (every proxy RPC passes
// Mock, Before and After), a coordinator RPC, or a configuration item.
//
// A stock binary installs nothing: InstalledHook, InstalledQueryHook,
// InstalledCipher and InstalledCoordinatorEngine answer nil, and milvus
// behaves as it always did. The one context mark this package defines,
// WithQueryResourceGroup, pins a query to a resource group for routing;
// nothing in a stock binary sets it.
package extension
import (
"sync/atomic"
"github.com/milvus-io/milvus-proto/go-api/v3/hook"
)
var (
installedHook atomic.Pointer[hookBox]
installedEngine atomic.Pointer[engineBox]
installedQueryHook atomic.Pointer[queryHookBox]
installedCipher atomic.Pointer[cipherBox]
formInstalled atomic.Bool
)
type (
hookBox struct{ hook hook.Hook }
engineBox struct{ engine CoordinatorEngine }
queryHookBox struct{ hook QueryHook }
cipherBox struct{ cipher hook.Cipher }
)
// SetHook installs a compiled-in request hook. hookutil prefers it over
// proxy.soPath and refuses a deployment that configures both. Call it before
// milvus starts; a nil hook leaves the stock behavior in place.
//
// Init is called again whenever the hook's configuration is refreshed. A
// compiled-in hook whose Init returns an error then must leave the
// configuration it was running with in effect: hookutil logs the refusal and
// keeps the process up, where a plug-in's failed refresh stops it, and it has
// no way to inspect or roll back what the hook did with the new values.
func SetHook(h hook.Hook) {
installedHook.Store(&hookBox{hook: h})
}
// InstalledHook returns the installed request hook, or nil.
func InstalledHook() hook.Hook {
if b := installedHook.Load(); b != nil {
return b.hook
}
return nil
}
// SetForm marks this binary as running the deployment form the coordinators'
// form-gated behaviors are written for: one streaming node, several resource
// groups loading the same collection, every role from one image. It switches
// on, in the query coordinator and the data coordinator: shard delegators on
// the regular query nodes of the resource groups, with the streaming node kept
// for DDL and the write ahead log, a load that names the resource groups it
// speaks for, an index engine version answered before any QueryNode registers.
// Call it in EVERY role, before milvus starts - a role whose main forgets the
// call runs as stock with no signal at all. It is independent of SetHook: a
// distribution may install a request hook without switching any coordinator
// behavior on.
func SetForm() {
formInstalled.Store(true)
}
// FormInstalled reports whether SetForm was called. A stock binary answers
// false everywhere and keeps the behavior it always had.
func FormInstalled() bool {
return formInstalled.Load()
}
// SetCoordinatorEngine installs the engine the coordinator starts when it
// becomes active and stops on shutdown. Call it before milvus starts; nil
// installs nothing.
func SetCoordinatorEngine(e CoordinatorEngine) {
installedEngine.Store(&engineBox{engine: e})
}
// InstalledCoordinatorEngine returns the installed engine, or nil.
func InstalledCoordinatorEngine() CoordinatorEngine {
if b := installedEngine.Load(); b != nil {
return b.engine
}
return nil
}
pkg/extension/extension_test.go
+169
View File
@@ -0,0 +1,169 @@
// 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 extension
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus-proto/go-api/v3/hook"
)
// stubHook is a hook.Hook that answers nothing; the tests only ask whether it
// is the one installed.
type stubHook struct{ name string }
func (stubHook) Init(map[string]string) error { return nil }
func (stubHook) Mock(context.Context, interface{}, string) (bool, interface{}, error) {
return false, nil, nil
}
func (stubHook) Before(ctx context.Context, _ interface{}, _ string) (context.Context, error) {
return ctx, nil
}
func (stubHook) After(context.Context, interface{}, error, string) error { return nil }
func (stubHook) VerifyAPIKey(string) (string, error) { return "", nil }
func (stubHook) Release() {}
var _ hook.Hook = stubHook{}
func TestNothingIsInstalledByDefault(t *testing.T) {
ResetForTest()
assert.Nil(t, InstalledHook(), "a stock binary has no compiled-in hook")
assert.Nil(t, InstalledCoordinatorEngine(), "a stock binary has no coordinator engine")
assert.Nil(t, InstalledQueryHook(), "a stock binary has no compiled-in query hook")
assert.Nil(t, InstalledCipher(), "a stock binary has no compiled-in cipher")
}
func TestSetHookInstallsTheHook(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
SetHook(stubHook{name: "form"})
assert.Equal(t, stubHook{name: "form"}, InstalledHook())
assert.Nil(t, InstalledCoordinatorEngine(), "installing a hook must not conjure an engine")
}
// FormInstalled is the one question the coordinators ask, and it must follow
// SetForm alone: nothing else a distribution sets counts, and SetForm must not
// conjure anything either.
func TestFormInstalledFollowsSetFormAlone(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
assert.False(t, FormInstalled(), "a stock binary has no form installed")
SetHook(stubHook{name: "form"})
assert.False(t, FormInstalled(), "a hook alone is not a form: the hook is not the mark")
SetCoordinatorEngine(&fakeCoordinatorEngine{})
assert.False(t, FormInstalled(), "an engine alone is not a form")
SetForm()
assert.True(t, FormInstalled(), "SetForm declares the form")
ResetForTest()
assert.False(t, FormInstalled(), "reset clears the form mark")
}
func TestSetFormDoesNotConjureAnything(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
SetForm()
assert.Nil(t, InstalledHook(), "declaring a form must not conjure a hook")
assert.Nil(t, InstalledCoordinatorEngine(), "declaring a form must not conjure an engine")
assert.Nil(t, InstalledQueryHook(), "declaring a form must not conjure a query hook")
assert.Nil(t, InstalledCipher(), "declaring a form must not conjure a cipher")
}
func TestSetCoordinatorEngineInstallsTheEngine(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
engine := &fakeCoordinatorEngine{}
SetCoordinatorEngine(engine)
assert.Same(t, engine, InstalledCoordinatorEngine())
assert.Nil(t, InstalledHook(), "installing an engine must not conjure a hook")
}
func TestSetNilLeavesNothingInstalled(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
SetHook(stubHook{name: "form"})
SetCoordinatorEngine(&fakeCoordinatorEngine{})
SetQueryHook(stubQueryHook{name: "tuner"})
SetCipher(stubCipher{name: "kms"})
SetHook(nil)
SetCoordinatorEngine(nil)
SetQueryHook(nil)
SetCipher(nil)
assert.Nil(t, InstalledHook())
assert.Nil(t, InstalledCoordinatorEngine())
assert.Nil(t, InstalledQueryHook())
assert.Nil(t, InstalledCipher())
}
// stubQueryHook is a QueryHook that tunes nothing; the tests only ask
// whether it is the one installed.
type stubQueryHook struct{ name string }
func (stubQueryHook) Run(map[string]any) error { return nil }
func (stubQueryHook) Init(string) error { return nil }
func (stubQueryHook) InitTuningConfig(map[string]string) error { return nil }
func (stubQueryHook) DeleteTuningConfig(string) error { return nil }
func (stubQueryHook) CalculateEffectiveSegmentNum([]int64, int64) int { return 0 }
var _ QueryHook = stubQueryHook{}
func TestSetQueryHookInstallsTheQueryHook(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
SetQueryHook(stubQueryHook{name: "tuner"})
assert.Equal(t, stubQueryHook{name: "tuner"}, InstalledQueryHook())
assert.Nil(t, InstalledHook(), "installing a query hook must not conjure a request hook")
assert.Nil(t, InstalledCoordinatorEngine(), "installing a query hook must not conjure an engine")
assert.False(t, FormInstalled(), "a query hook is not the mark of a form")
SetQueryHook(nil)
assert.Nil(t, InstalledQueryHook())
}
// stubCipher is a hook.Cipher that encrypts nothing; the tests only ask
// whether it is the one installed.
type stubCipher struct{ name string }
func (stubCipher) Init(map[string]string) error { return nil }
func (stubCipher) GetEncryptor(int64, int64) (hook.Encryptor, []byte, error) {
return nil, nil, nil
}
func (stubCipher) GetDecryptor(int64, int64, []byte) (hook.Decryptor, error) { return nil, nil }
func (stubCipher) GetUnsafeKey(int64, int64) []byte { return nil }
var _ hook.Cipher = stubCipher{}
func TestSetCipherInstallsTheCipher(t *testing.T) {
ResetForTest()
t.Cleanup(ResetForTest)
SetCipher(stubCipher{name: "kms"})
assert.Equal(t, stubCipher{name: "kms"}, InstalledCipher())
assert.Nil(t, InstalledHook(), "installing a cipher must not conjure a request hook")
assert.Nil(t, InstalledQueryHook(), "installing a cipher must not conjure a query hook")
assert.Nil(t, InstalledCoordinatorEngine(), "installing a cipher must not conjure an engine")
assert.False(t, FormInstalled(), "a cipher is not the mark of a form")
SetCipher(nil)
assert.Nil(t, InstalledCipher())
}
pkg/extension/query_hook.go
+47
View File
@@ -0,0 +1,47 @@
// 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 extension
// QueryHook is the search-parameter tuning hook the QueryNode consults before
// every search: the same method set a queryNode.soPath plug-in exports as its
// QueryNodePlugin symbol, spelled here so that a distribution can compile one
// into the binary instead. Init receives autoIndex.params.search,
// InitTuningConfig and DeleteTuningConfig follow autoIndex.params.tuning, Run
// tunes one search and CalculateEffectiveSegmentNum sizes it.
type QueryHook interface {
Run(map[string]any) error
Init(string) error
InitTuningConfig(map[string]string) error
DeleteTuningConfig(string) error
CalculateEffectiveSegmentNum(rowCounts []int64, topk int64) int
}
// SetQueryHook installs a compiled-in query hook. The QueryNode prefers it
// over queryNode.soPath and refuses a deployment that configures both. Call it
// before milvus starts; a nil hook leaves the stock behavior in place. It is
// independent of SetHook and of FormInstalled: it marks nothing but itself.
func SetQueryHook(h QueryHook) {
installedQueryHook.Store(&queryHookBox{hook: h})
}
// InstalledQueryHook returns the installed query hook, or nil.
func InstalledQueryHook() QueryHook {
if b := installedQueryHook.Load(); b != nil {
return b.hook
}
return nil
}
pkg/extension/query_scope.go
+36
View File
@@ -0,0 +1,36 @@
// 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 extension
import "context"
type queryResourceGroupKey struct{}
// WithQueryResourceGroup pins the query running under ctx to one resource
// group: the proxy routes it to the shard leaders of that group only and
// attributes its latency to that group. A request hook sets it from Before;
// nothing in a stock binary does.
func WithQueryResourceGroup(ctx context.Context, resourceGroup string) context.Context {
return context.WithValue(ctx, queryResourceGroupKey{}, resourceGroup)
}
// QueryResourceGroupFromContext returns the resource group WithQueryResourceGroup
// pinned, or "" when the query is not pinned.
func QueryResourceGroupFromContext(ctx context.Context) string {
resourceGroup, _ := ctx.Value(queryResourceGroupKey{}).(string)
return resourceGroup
}
pkg/extension/query_scope_test.go
+65
View File
@@ -0,0 +1,65 @@
// 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 extension
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
)
func TestQueryResourceGroupSurvivesOnTheContext(t *testing.T) {
ctx := WithQueryResourceGroup(context.Background(), "rg-a")
assert.Equal(t, "rg-a", QueryResourceGroupFromContext(ctx),
"the routing scope must be readable by every stage that runs under the request's context")
assert.Equal(t, "", QueryResourceGroupFromContext(context.Background()),
"a context nothing scoped must report no scope, which is what every request in a stock binary looks like")
}
// TestQueryResourceGroupKeyIsUnreachableFromOutside proves nothing but this
// package can plant a routing scope. An extension carries its own values on the
// same context - a hook's Before returns the context the handler runs under -
// and one of them guessing the string "rg" or "resource_group" must not become
// the scope every shard-leader lookup honors. The scope is milvus's decision,
// made in exactly
// one place, and this is what keeps it that way.
func TestQueryResourceGroupKeyIsUnreachableFromOutside(t *testing.T) {
ctx := context.Background()
for _, guess := range []string{"rg", "resource_group", "queryResourceGroupKey"} {
//nolint:staticcheck // planting a string key is exactly what is under test
ctx = context.WithValue(ctx, guess, "forged-"+guess)
}
assert.Equal(t, "", QueryResourceGroupFromContext(ctx),
"only WithQueryResourceGroup may set the scope routing reads back")
}
// TestQueryResourceGroupRebindShadowsTheOuterValue pins the precedence a nested
// bind has: a re-scoped context answers with the inner value while the context
// it derived from keeps answering with its own. A sub-query run under an outer
// request's context - the search-by-primary-key requery does exactly this -
// must not be able to re-route the request it runs inside.
func TestQueryResourceGroupRebindShadowsTheOuterValue(t *testing.T) {
outer := WithQueryResourceGroup(context.Background(), "rg-outer")
inner := WithQueryResourceGroup(outer, "rg-inner")
assert.Equal(t, "rg-inner", QueryResourceGroupFromContext(inner))
assert.Equal(t, "rg-outer", QueryResourceGroupFromContext(outer),
"scoping a derived context must not rewrite the context it derived from")
}
pkg/extension/reset_testonly.go
+29
View File
@@ -0,0 +1,29 @@
// 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.
//go:build test
package extension
// ResetForTest uninstalls everything a distribution can install, so tests
// start from a stock binary.
func ResetForTest() {
installedHook.Store(nil)
installedEngine.Store(nil)
installedQueryHook.Store(nil)
installedCipher.Store(nil)
formInstalled.Store(false)
}
pkg/metrics/metrics_test.go
+1
View File
@@ -290,6 +290,7 @@ func observeProxyCollection(nodeID, db, collection string) {
ProxyRetrySearchResultInsufficientCount.WithLabelValues(nodeID, queryType, db, collection).Add(1)
ProxyRecallSearchCount.WithLabelValues(nodeID, queryType, db, collection).Add(1)
ProxySearchSparseNumNonZeros.WithLabelValues(nodeID, db, collection, queryType, "1").Observe(1)
ProxyResourceGroupSQLatency.WithLabelValues(nodeID, queryType, db, collection, "rg").Observe(1)
}
for _, msgType := range []string{InsertLabel, DeleteLabel, UpsertLabel, SearchLabel, HybridSearchLabel, QueryLabel} {
ProxyMutationLatency.WithLabelValues(nodeID, msgType, db, collection).Observe(1)
pkg/metrics/proxy_metrics.go
+20
View File
@@ -95,6 +95,24 @@ var (
Buckets: buckets,
}, []string{nodeIDLabelName, queryTypeLabelName, databaseLabelName, collectionName})
// ProxyResourceGroupSQLatency records the latency of a search or query
// attributed to the resource group that served it.
//
// It is a family of its own rather than an rg label on ProxySQLatency
// because only a deployment form that scopes its queries to a resource group
// can attribute anything: adding the label to the existing family would give
// every stock milvus an extra, permanently empty dimension on a metric its
// dashboards and recording rules already read. Nothing scopes a query in a
// binary with no provider installed, so it exposes no series here at all.
ProxyResourceGroupSQLatency = prometheus.NewHistogramVec(
prometheus.HistogramOpts{
Namespace: milvusNamespace,
Subsystem: typeutil.ProxyRole,
Name: "resource_group_sq_latency",
Help: "latency of search or query successfully, per resource group",
Buckets: buckets,
}, []string{nodeIDLabelName, queryTypeLabelName, databaseLabelName, collectionName, ResourceGroupLabelName})
// ProxyCollectionSQLatency record the latency of search successfully, per collection
// Deprecated, ProxySQLatency instead of it
ProxyCollectionSQLatency = prometheus.NewHistogramVec(
@@ -512,6 +530,7 @@ func RegisterProxy(registry *prometheus.Registry) {
registry.MustRegister(ProxyPathReplaceMergeLatency)
registry.MustRegister(ProxySQLatency)
registry.MustRegister(ProxyResourceGroupSQLatency)
registry.MustRegister(ProxyCollectionSQLatency)
registry.MustRegister(ProxyMutationLatency)
registry.MustRegister(ProxyCollectionMutationLatency)
@@ -614,6 +633,7 @@ func proxyCollectionScopedMetrics() []partialMatchDeleter {
ProxySearchSparseNumNonZeros,
ProxyScannedRemoteMB,
ProxyScannedTotalMB,
ProxyResourceGroupSQLatency,
}
}
pkg/proto/query_coord.proto
+9
View File
@@ -208,12 +208,21 @@ message ShowCollectionsRequest {
// Not useful for now
int64 dbID = 2;
repeated int64 collectionIDs = 3;
// When set, inMemory_percentages is the load progress of the replicas that
// live in this resource group only, -1 for a collection that has no replica
// there. A collection that is not loaded at all answers the same -1, with
// query_service_available false, refresh_progress 0 and empty load_fields,
// instead of the CollectionNotLoaded error the unscoped call returns.
// Empty keeps the collection-wide figure.
string resource_group = 4;
}
message ShowCollectionsResponse {
common.Status status = 1;
repeated int64 collectionIDs = 2;
repeated int64 inMemory_percentages = 3;
// Follows the request's scope: with resource_group set it says whether
// that group can serve, otherwise whether any replica of the collection can.
repeated bool query_service_available = 4;
repeated int64 refresh_progress = 5;
repeated schema.LongArray load_fields = 6;
pkg/proto/querypb/query_coord.pb.go
+2220 -2202
View File
File diff suppressed because it is too large Load Diff
pkg/util/paramtable/component_param.go
-1
View File
@@ -6980,7 +6980,6 @@ Startup processes fixed-size batches and retries failed reads per segment. An ex
Export: true,
}
p.SegmentIndexManifestLoadConcurrency.Init(base.mgr)
p.HybridIndexLowCardinalityIndexType = ParamItem{
Key: "dataCoord.index.hybridIndex.lowCardinalityIndexType",
Version: "2.6.10",
scripts/run_go_codecov.sh
+1 -1
View File
@@ -40,7 +40,7 @@ export MILVUS_UT_WITHOUT_KAFKA=1 # kafka is not available in the CI environment,
beginTime=$(date +%s)
echo -e "=== Running go unittest (parallel=$PARALLEL) ===\n\n"
for d in cmd/tools internal pkg client; do
for d in cmd/tools cmd/milvus internal pkg client; do
pushd "$d"
PROFILE="profile_${d//\//_}.out"
echo -e "-----------------------------------\nRunning test cases under $d with -p $PARALLEL ..."
scripts/run_go_unittest.sh
+1
View File
@@ -167,6 +167,7 @@ go test -gcflags="all=-N -l" -race -cover -tags dynamic,test "${MILVUS_DIR}/meta
function test_cmd()
{
go test -gcflags="all=-N -l" -race -cover -tags dynamic,test "${ROOT_DIR}/cmd/tools/..." -failfast -count=1 -ldflags="-r ${RPATH}"
go test -gcflags="all=-N -l" -race -cover -tags dynamic,test "${ROOT_DIR}/cmd/milvus/..." -failfast -count=1 -ldflags="-r ${RPATH}"
}
function test_streaming()