mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 02:05:41 +00:00
enhance: optimize RPC protobuf decoding with fastpb codec (#50743)
issue: #50742 ## What this PR does Optimizes protobuf decoding on Milvus RPC read/write hot paths with a hand-written decoder in `pkg/util/fastpb`. The final change contains no gRPC connection-pool or connection-count changes. ## Fast decoder coverage The gRPC `releaseCodec.Unmarshal` path dispatches these top-level messages to fastpb: - `internalpb.RetrieveResults` - `milvuspb.InsertRequest` - `milvuspb.UpsertRequest` `schemapb.SearchResultData` is not a top-level codec-dispatched message. It is decoded directly from `SlicedBlob` at the search/reduce hot paths in: - `internal/proxy/search_reduce_util.go` - `internal/querynodev2/segments/result.go` - `internal/querynodev2/local_worker.go` - `internal/querynodev2/tasks/search_task_go_reduce.go` Nested hot messages such as `FieldData`, scalar/vector arrays, and IDs are decoded by the same package. Main optimizations: - varchar string arrays use one backing byte arena, reducing per-string allocations - packed float vectors use a single memcpy into self-owned memory - cold, complex, unknown, and mismatched fields delegate or merge through the official protobuf codec - no new external dependency or source-buffer aliasing ## Compatibility boundary The compatibility contract depends on the data path: - **Untrusted client ingress (`InsertRequest` and `UpsertRequest`)** validates proto3 strings as UTF-8 and is tested differentially against `proto.Unmarshal`, including malformed input behavior. - **Trusted internal result paths (`SearchResultData`, `RetrieveResults`, and direct `FieldData` decoding)** assume canonical Milvus/segcore protobuf output and intentionally skip UTF-8 validation for performance. Therefore these paths do not claim equivalence for arbitrary adversarial wire containing invalid UTF-8. Within that boundary, the decoder preserves the protobuf behaviors exercised by the differential tests: - proto2 groups fall back to the official codec - known fields with mismatched wire types fall back to the official codec - unknown and future fields are preserved through an official merge pass - repeated singular messages merge from the raw wire, including explicitly encoded proto3 default values - repeated message-valued oneof members merge correctly, while different variants remain last-wins - mixed packed and unpacked scalar encodings preserve wire order - descriptor field-set tripwires force decoder review when the protobuf schema changes The unsafe packed-float copy assumes the currently supported little-endian Milvus targets such as x86-64 and arm64. ## Benchmarks fastpb vs official `proto.Unmarshal`, using 1000-row payloads: | path | payload | speedup | allocations/op | |---|---|---:|---:| | search / query | varchar | about 2x | about 1000-2035 to about 10-17 | | search / query | float vector, dim 768 | about 6x | unchanged | | insert | varchar with UTF-8 validation | about 1.8x | 1019 to about 10 | ## Verification - unit and differential tests cover field descriptors, UTF-8 contracts, wire-type mismatches, proto2 groups, unknown fields, packed/unpacked ordering, oneof behavior, and repeated-message merge semantics - randomized canonical-message equivalence tests cover search, retrieve, insert, and upsert payloads - `go test ./util/fastpb ./util/resource` passes - `go vet ./util/fastpb ./util/resource` passes - current full CI and end-to-end status is tracked by the PR checks 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: xiaofanluan <xiaofan.luan@zilliz.com> Signed-off-by: xiaofanluan <xf@hjjaq.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: xiaofanluan <xf@hjjaq.com>
This commit is contained in:
James
co-authored by
Claude Fable 5
xiaofanluan
Claude Fable 5
xiaofanluan
parent
6b3eabb1ff
commit
89f107fa00
@@ -4,7 +4,6 @@ import (
|
||||
"context"
|
||||
|
||||
"go.opentelemetry.io/otel"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
@@ -12,6 +11,7 @@ import (
|
||||
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/fastpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/metric"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
||||
@@ -840,7 +840,9 @@ func decodeSearchResults(ctx context.Context, searchResults []*internalpb.Search
|
||||
results = append(results, partialSearchResult.ResultData)
|
||||
} else if partialSearchResult.SlicedBlob != nil {
|
||||
var partialResultData schemapb.SearchResultData
|
||||
err := proto.Unmarshal(partialSearchResult.SlicedBlob, &partialResultData)
|
||||
// fastpb: hand-written decoder for the search reduce hot path
|
||||
// (wire-equivalent to proto.Unmarshal, ~2x varchar / ~6x vector).
|
||||
err := fastpb.UnmarshalSearchResultData(partialSearchResult.SlicedBlob, &partialResultData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -19,14 +19,13 @@ package querynodev2
|
||||
import (
|
||||
"context"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/internal/querynodev2/cluster"
|
||||
"github.com/milvus-io/milvus/internal/util/streamrpc"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/fastpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/resource"
|
||||
)
|
||||
@@ -73,7 +72,9 @@ func (w *LocalWorker) SearchSegments(ctx context.Context, req *querypb.SearchReq
|
||||
// use it directly), then release any pinned C memory.
|
||||
if blob := resp.GetSlicedBlob(); len(blob) > 0 {
|
||||
var resultData schemapb.SearchResultData
|
||||
if unmarshalErr := proto.Unmarshal(blob, &resultData); unmarshalErr != nil {
|
||||
// fastpb: wire-equivalent fast decoder for the local (in-process) search
|
||||
// hot path (~2x varchar / ~6x vector vs proto.Unmarshal).
|
||||
if unmarshalErr := fastpb.UnmarshalSearchResultData(blob, &resultData); unmarshalErr != nil {
|
||||
resource.MsgPins.Release(resp) // still release to avoid leak
|
||||
return nil, merr.WrapErrServiceInternal("unmarshal SearchResultData from SlicedBlob", unmarshalErr.Error())
|
||||
}
|
||||
|
||||
@@ -30,6 +30,7 @@ import (
|
||||
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/segcorepb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/fastpb"
|
||||
"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"
|
||||
@@ -233,7 +234,9 @@ func DecodeSearchResults(ctx context.Context, searchResults []*internalpb.Search
|
||||
results = append(results, partialSearchResult.ResultData)
|
||||
} else if partialSearchResult.SlicedBlob != nil {
|
||||
var partialResultData schemapb.SearchResultData
|
||||
err := proto.Unmarshal(partialSearchResult.SlicedBlob, &partialResultData)
|
||||
// fastpb: hand-written decoder for the search reduce hot path
|
||||
// (wire-equivalent to proto.Unmarshal, ~2x varchar / ~6x vector).
|
||||
err := fastpb.UnmarshalSearchResultData(partialSearchResult.SlicedBlob, &partialResultData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -25,7 +25,6 @@ import (
|
||||
"github.com/apache/arrow/go/v17/arrow"
|
||||
"github.com/apache/arrow/go/v17/arrow/array"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
@@ -34,6 +33,7 @@ import (
|
||||
"github.com/milvus-io/milvus/internal/util/segcore"
|
||||
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/fastpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
|
||||
)
|
||||
@@ -361,7 +361,9 @@ func lateMaterializeOutputFields(
|
||||
}
|
||||
|
||||
var fieldResult schemapb.SearchResultData
|
||||
if err := proto.Unmarshal(protoBytes, &fieldResult); err != nil {
|
||||
// fastpb: wire-equivalent fast decoder for the late-materialize output-fields
|
||||
// hot path (~2x varchar / ~6x vector vs proto.Unmarshal).
|
||||
if err := fastpb.UnmarshalSearchResultData(protoBytes, &fieldResult); err != nil {
|
||||
return err
|
||||
}
|
||||
searchResultData.FieldsData = fieldResult.FieldsData
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
)
|
||||
|
||||
// buildVarcharSRD builds a varchar-PK + varchar-output SearchResultData (rows = nq*topk).
|
||||
func buildVarcharSRD(rows, strLen int) *schemapb.SearchResultData {
|
||||
ids := make([]string, rows)
|
||||
for i := range ids {
|
||||
ids[i] = fmt.Sprintf("pk_%0*d", strLen, i)
|
||||
}
|
||||
scores := make([]float32, rows)
|
||||
for i := range scores {
|
||||
scores[i] = float32(i) * 0.5
|
||||
}
|
||||
return &schemapb.SearchResultData{
|
||||
NumQueries: 10, TopK: int64(rows / 10), PrimaryFieldName: "pk",
|
||||
Scores: scores,
|
||||
Ids: &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: ids}}},
|
||||
FieldsData: varcharFieldData(rows, strLen),
|
||||
}
|
||||
}
|
||||
|
||||
// buildVectorSRD builds an int64-PK + float-vector SearchResultData.
|
||||
func buildVectorSRD(rows, dim int) *schemapb.SearchResultData {
|
||||
ids := make([]int64, rows)
|
||||
for i := range ids {
|
||||
ids[i] = int64(i)
|
||||
}
|
||||
scores := make([]float32, rows)
|
||||
return &schemapb.SearchResultData{
|
||||
NumQueries: 10, TopK: int64(rows / 10), PrimaryFieldName: "id",
|
||||
Scores: scores,
|
||||
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: ids}}},
|
||||
FieldsData: vectorFieldData(rows, dim),
|
||||
}
|
||||
}
|
||||
|
||||
func benchSRD(b *testing.B, src *schemapb.SearchResultData) {
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Run("official", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out schemapb.SearchResultData
|
||||
if err := proto.Unmarshal(wire, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
b.Run("fastpb", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out schemapb.SearchResultData
|
||||
if err := UnmarshalSearchResultData(wire, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkSRD_Varchar(b *testing.B) { benchSRD(b, buildVarcharSRD(1000, 12)) }
|
||||
func BenchmarkSRD_Vector(b *testing.B) { benchSRD(b, buildVectorSRD(1000, 768)) }
|
||||
@@ -0,0 +1,372 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
commonpb "github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
// --- wire helpers: append a single unknown (future) field of each wire type ---
|
||||
|
||||
func appendUnknownFixed64(b []byte, num int32) []byte {
|
||||
b = protowire.AppendTag(b, protowire.Number(num), protowire.Fixed64Type)
|
||||
return protowire.AppendFixed64(b, 0xDEADBEEFCAFEF00D)
|
||||
}
|
||||
|
||||
func appendUnknownFixed32(b []byte, num int32) []byte {
|
||||
b = protowire.AppendTag(b, protowire.Number(num), protowire.Fixed32Type)
|
||||
return protowire.AppendFixed32(b, 0xCAFEF00D)
|
||||
}
|
||||
|
||||
func appendUnknownLenDelim(b []byte, num int32) []byte {
|
||||
b = protowire.AppendTag(b, protowire.Number(num), protowire.BytesType)
|
||||
return protowire.AppendBytes(b, []byte("future-field-bytes"))
|
||||
}
|
||||
|
||||
func appendUnknownVarint(b []byte, num int32) []byte {
|
||||
b = protowire.AppendTag(b, protowire.Number(num), protowire.VarintType)
|
||||
return protowire.AppendVarint(b, 0x7FFFFFFF)
|
||||
}
|
||||
|
||||
// TestDuplicatedMessage_ExplicitDefaultScalar reproduces a wire-equivalence gap.
|
||||
// When a nested message field (Base) appears twice and the SECOND occurrence
|
||||
// explicitly encodes a proto3-default scalar (TargetID=0), official proto.Unmarshal
|
||||
// applies it (wire-level merge -> 0). A decoder that accumulates via message-level
|
||||
// proto.Merge instead DROPS it (proto.Merge skips proto3-default scalars), keeping
|
||||
// the stale 7. External/adversarial Upsert/Insert input can craft exactly this.
|
||||
func TestDuplicatedMessage_ExplicitDefaultScalar(t *testing.T) {
|
||||
base7 := protowire.AppendVarint(protowire.AppendTag(nil, 5, protowire.VarintType), 7) // MsgBase{TargetID:7}
|
||||
base0 := protowire.AppendVarint(protowire.AppendTag(nil, 5, protowire.VarintType), 0) // MsgBase{TargetID:0 explicit}
|
||||
dupBase := func() []byte {
|
||||
w := protowire.AppendBytes(protowire.AppendTag(nil, 1, protowire.BytesType), base7)
|
||||
return protowire.AppendBytes(protowire.AppendTag(w, 1, protowire.BytesType), base0)
|
||||
}
|
||||
t.Run("RetrieveResults", func(t *testing.T) { diffDecode(t, dupBase(), newRetrieve, decRetrieve) })
|
||||
t.Run("InsertRequest", func(t *testing.T) { diffDecode(t, dupBase(), newInsertRequest, decInsertRequest) })
|
||||
t.Run("UpsertRequest", func(t *testing.T) { diffDecode(t, dupBase(), newUpsertRequest, decUpsertRequest) })
|
||||
|
||||
// Same class of bug on a hand-decoded nested field: SearchResultData's
|
||||
// GroupByFieldValue (field 8, a FieldData) appearing twice where the second
|
||||
// occurrence explicitly encodes FieldData.FieldId=0.
|
||||
t.Run("SearchResultData/GroupByFieldValue", func(t *testing.T) {
|
||||
fd9 := protowire.AppendVarint(protowire.AppendTag(nil, 5, protowire.VarintType), 9) // FieldData{FieldId:9}
|
||||
fd0 := protowire.AppendVarint(protowire.AppendTag(nil, 5, protowire.VarintType), 0) // FieldData{FieldId:0 explicit}
|
||||
w := protowire.AppendBytes(protowire.AppendTag(nil, 8, protowire.BytesType), fd9)
|
||||
w = protowire.AppendBytes(protowire.AppendTag(w, 8, protowire.BytesType), fd0)
|
||||
diffDecode(t, w, newSearchResult, decSearchResult)
|
||||
})
|
||||
}
|
||||
|
||||
// diffDecode asserts the fastpb decoder matches the official codec for wire,
|
||||
// both in error behavior and (on success) in the decoded message.
|
||||
func diffDecode(t *testing.T, wire []byte, fresh func() proto.Message, fast func([]byte, proto.Message) error) {
|
||||
t.Helper()
|
||||
want := fresh()
|
||||
wantErr := proto.Unmarshal(wire, want)
|
||||
got := fresh()
|
||||
gotErr := fast(wire, got)
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
require.True(t, proto.Equal(want, got), "mismatch:\n want=%v\n got=%v", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
// decoder adapters so each entry point shares diffDecode.
|
||||
var (
|
||||
decFieldData = func(b []byte, m proto.Message) error { return UnmarshalFieldData(b, m.(*schemapb.FieldData)) }
|
||||
decScalarField = func(b []byte, m proto.Message) error { return dec{}.scalarField(b, m.(*schemapb.ScalarField)) }
|
||||
decVectorField = func(b []byte, m proto.Message) error { return unmarshalVectorField(b, m.(*schemapb.VectorField)) }
|
||||
decIDs = func(b []byte, m proto.Message) error { return dec{}.ids(b, m.(*schemapb.IDs)) }
|
||||
decSearchResult = func(b []byte, m proto.Message) error {
|
||||
return UnmarshalSearchResultData(b, m.(*schemapb.SearchResultData))
|
||||
}
|
||||
decRetrieve = func(b []byte, m proto.Message) error {
|
||||
return UnmarshalRetrieveResults(b, m.(*internalpb.RetrieveResults))
|
||||
}
|
||||
decInsertRequest = func(b []byte, m proto.Message) error { return UnmarshalInsertRequest(b, m.(*milvuspb.InsertRequest)) }
|
||||
decUpsertRequest = func(b []byte, m proto.Message) error { return UnmarshalUpsertRequest(b, m.(*milvuspb.UpsertRequest)) }
|
||||
newFieldData = func() proto.Message { return &schemapb.FieldData{} }
|
||||
newScalarField = func() proto.Message { return &schemapb.ScalarField{} }
|
||||
newVectorField = func() proto.Message { return &schemapb.VectorField{} }
|
||||
newIDs = func() proto.Message { return &schemapb.IDs{} }
|
||||
newSearchResult = func() proto.Message { return &schemapb.SearchResultData{} }
|
||||
newRetrieve = func() proto.Message { return &internalpb.RetrieveResults{} }
|
||||
newInsertRequest = func() proto.Message { return &milvuspb.InsertRequest{} }
|
||||
newUpsertRequest = func() proto.Message { return &milvuspb.UpsertRequest{} }
|
||||
)
|
||||
|
||||
// TestUnknownFieldSkipAndMerge feeds each entry point a canonical message with a
|
||||
// trailing unknown (future) field of every wire type. This exercises skipField
|
||||
// (all wire-type cases) plus the protoMerge "rest" tail that folds the unknown
|
||||
// bytes back via the official codec, so the result must equal proto.Unmarshal.
|
||||
func TestUnknownFieldSkipAndMerge(t *testing.T) {
|
||||
canonFieldData, err := proto.Marshal(&schemapb.FieldData{
|
||||
Type: schemapb.DataType_Int64, FieldName: "f", FieldId: 3,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2, 3}}}}},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
canonScalar, err := proto.Marshal(&schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{4, 5}}}})
|
||||
require.NoError(t, err)
|
||||
canonVector, err := proto.Marshal(&schemapb.VectorField{Dim: 4, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}}}})
|
||||
require.NoError(t, err)
|
||||
canonIDs, err := proto.Marshal(&schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{7, 8}}}})
|
||||
require.NoError(t, err)
|
||||
canonSRD, err := proto.Marshal(&schemapb.SearchResultData{NumQueries: 2, TopK: 5, Scores: []float32{0.1, 0.2}})
|
||||
require.NoError(t, err)
|
||||
canonRR, err := proto.Marshal(&internalpb.RetrieveResults{ReqID: 7, AllRetrieveCount: 9})
|
||||
require.NoError(t, err)
|
||||
canonIR, err := proto.Marshal(&milvuspb.InsertRequest{CollectionName: "c", DbName: "db", NumRows: 3})
|
||||
require.NoError(t, err)
|
||||
canonUR, err := proto.Marshal(&milvuspb.UpsertRequest{CollectionName: "c", DbName: "db", NumRows: 3, PartialUpdate: true})
|
||||
require.NoError(t, err)
|
||||
|
||||
entries := []struct {
|
||||
name string
|
||||
canon []byte
|
||||
unknown int32 // an unused field number for this message
|
||||
fresh func() proto.Message
|
||||
fast func([]byte, proto.Message) error
|
||||
}{
|
||||
{"FieldData", canonFieldData, 20, newFieldData, decFieldData},
|
||||
{"ScalarField", canonScalar, 30, newScalarField, decScalarField},
|
||||
{"VectorField", canonVector, 20, newVectorField, decVectorField},
|
||||
{"IDs", canonIDs, 30, newIDs, decIDs},
|
||||
{"SearchResultData", canonSRD, 30, newSearchResult, decSearchResult},
|
||||
{"RetrieveResults", canonRR, 30, newRetrieve, decRetrieve},
|
||||
{"InsertRequest", canonIR, 30, newInsertRequest, decInsertRequest},
|
||||
{"UpsertRequest", canonUR, 30, newUpsertRequest, decUpsertRequest},
|
||||
}
|
||||
for _, e := range entries {
|
||||
for _, w := range []struct {
|
||||
name string
|
||||
append func([]byte, int32) []byte
|
||||
}{
|
||||
{"fixed64", appendUnknownFixed64},
|
||||
{"fixed32", appendUnknownFixed32},
|
||||
{"lendelim", appendUnknownLenDelim},
|
||||
{"varint", appendUnknownVarint},
|
||||
} {
|
||||
t.Run(e.name+"/"+w.name, func(t *testing.T) {
|
||||
wire := w.append(append([]byte{}, e.canon...), e.unknown)
|
||||
diffDecode(t, wire, e.fresh, e.fast)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestDecodePackedNonPackedFixed exercises the single-element (non-packed)
|
||||
// fixed32/fixed64 paths (decodePackedF32 case 5 → le32, decodePackedF64 case 1
|
||||
// → le64), which the packed-only fuzz tests never reach.
|
||||
func TestDecodePackedNonPackedFixed(t *testing.T) {
|
||||
t.Run("float32 non-packed fixed32", func(t *testing.T) {
|
||||
var f []byte
|
||||
f = protowire.AppendTag(f, 1, protowire.Fixed32Type)
|
||||
f = protowire.AppendFixed32(f, math.Float32bits(1.5))
|
||||
f = protowire.AppendTag(f, 1, protowire.Fixed32Type)
|
||||
f = protowire.AppendFixed32(f, math.Float32bits(-2.25))
|
||||
|
||||
want := &schemapb.FloatArray{}
|
||||
require.NoError(t, proto.Unmarshal(f, want))
|
||||
got := &schemapb.FloatArray{}
|
||||
require.NoError(t, decodePackedF32(f, &got.Data, got))
|
||||
require.True(t, proto.Equal(want, got))
|
||||
})
|
||||
t.Run("float64 non-packed fixed64", func(t *testing.T) {
|
||||
var d []byte
|
||||
d = protowire.AppendTag(d, 1, protowire.Fixed64Type)
|
||||
d = protowire.AppendFixed64(d, math.Float64bits(1.5))
|
||||
d = protowire.AppendTag(d, 1, protowire.Fixed64Type)
|
||||
d = protowire.AppendFixed64(d, math.Float64bits(-2.25))
|
||||
|
||||
want := &schemapb.DoubleArray{}
|
||||
require.NoError(t, proto.Unmarshal(d, want))
|
||||
got := &schemapb.DoubleArray{}
|
||||
require.NoError(t, decodePackedF64(d, &got.Data, got))
|
||||
require.True(t, proto.Equal(want, got))
|
||||
})
|
||||
}
|
||||
|
||||
// TestColdScalarVariants pins the rare ScalarField oneof variants that delegate
|
||||
// to the official codec (decodeScalarFallback cases 8/10/11/12/13/14).
|
||||
func TestColdScalarVariants(t *testing.T) {
|
||||
cases := map[string]*schemapb.ScalarField{
|
||||
"array": {Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{ElementType: schemapb.DataType_Int64}}},
|
||||
"geometry": {Data: &schemapb.ScalarField_GeometryData{GeometryData: &schemapb.GeometryArray{}}},
|
||||
"timestamptz": {Data: &schemapb.ScalarField_TimestamptzData{TimestamptzData: &schemapb.TimestamptzArray{}}},
|
||||
"geometrywkt": {Data: &schemapb.ScalarField_GeometryWktData{GeometryWktData: &schemapb.GeometryWktArray{}}},
|
||||
"mol": {Data: &schemapb.ScalarField_MolData{MolData: &schemapb.MolArray{}}},
|
||||
"molsmiles": {Data: &schemapb.ScalarField_MolSmilesData{MolSmilesData: &schemapb.MolSmilesArray{}}},
|
||||
}
|
||||
for name, sf := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
roundTripFieldData(t, &schemapb.FieldData{
|
||||
FieldName: name, FieldId: 5,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: sf},
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestColdSearchResultFields pins the delegate / rarer SearchResultData fields
|
||||
// that the common-case tests skip: iterator results (11), recalls (12),
|
||||
// highlights (14), element_indices (15), group_by_field_values (17),
|
||||
// agg_buckets (18), agg_topks (19).
|
||||
func TestColdSearchResultFields(t *testing.T) {
|
||||
roundTripSRD(t, &schemapb.SearchResultData{
|
||||
NumQueries: 1,
|
||||
TopK: 2,
|
||||
Recalls: []float32{0.95, 0.9},
|
||||
ElementIndices: &schemapb.LongArray{Data: []int64{0, 1, 2}},
|
||||
AggTopks: []int64{3, 4},
|
||||
SearchIteratorV2Results: &schemapb.SearchIteratorV2Results{Token: "tok", LastBound: 1.5},
|
||||
GroupByFieldValues: []*schemapb.FieldData{
|
||||
{FieldName: "g", FieldId: 9, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1}}}}}},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// TestRepeatedSingularMessageMerge: a singular message field encoded twice on the
|
||||
// wire is merged by proto3 (last-wins per scalar, concatenated repeated). This hits
|
||||
// the "already set → raw-wire merge" branches for Base/Status/Ids/CostAggregation.
|
||||
func TestRepeatedSingularMessageMerge(t *testing.T) {
|
||||
t.Run("RetrieveResults Base/Ids/Cost twice", func(t *testing.T) {
|
||||
// Build manually: field 1 (Base) twice, field 4 (Ids) twice, field 13 (Cost) twice.
|
||||
var wire []byte
|
||||
appendMsgField := func(num int32, m proto.Message) {
|
||||
bb, err := proto.Marshal(m)
|
||||
require.NoError(t, err)
|
||||
wire = protowire.AppendTag(wire, protowire.Number(num), protowire.BytesType)
|
||||
wire = protowire.AppendBytes(wire, bb)
|
||||
}
|
||||
appendMsgField(1, &commonpb.MsgBase{MsgID: 1, SourceID: 10})
|
||||
appendMsgField(1, &commonpb.MsgBase{TargetID: 2})
|
||||
appendMsgField(4, &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}})
|
||||
appendMsgField(4, &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2}}}})
|
||||
appendMsgField(13, &internalpb.CostAggregation{ResponseTime: 5})
|
||||
appendMsgField(13, &internalpb.CostAggregation{TotalNQ: 7})
|
||||
diffDecode(t, wire, newRetrieve, decRetrieve)
|
||||
})
|
||||
t.Run("SearchResultData Ids/iterator twice", func(t *testing.T) {
|
||||
var wire []byte
|
||||
appendMsgField := func(num int32, m proto.Message) {
|
||||
bb, err := proto.Marshal(m)
|
||||
require.NoError(t, err)
|
||||
wire = protowire.AppendTag(wire, protowire.Number(num), protowire.BytesType)
|
||||
wire = protowire.AppendBytes(wire, bb)
|
||||
}
|
||||
appendMsgField(5, &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}})
|
||||
appendMsgField(5, &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2}}}})
|
||||
appendMsgField(11, &schemapb.SearchIteratorV2Results{Token: "t1"})
|
||||
appendMsgField(11, &schemapb.SearchIteratorV2Results{LastBound: 2.0})
|
||||
diffDecode(t, wire, newSearchResult, decSearchResult)
|
||||
})
|
||||
}
|
||||
|
||||
// TestPackedFieldsAsSingleVarint: proto allows a packed repeated field to also
|
||||
// appear as a sequence of single varints. This hits the wtype==0 packed-field
|
||||
// branches in searchResultData (6/19), retrieveResults (6/8), insertRequest (6/7/8).
|
||||
func TestPackedFieldsAsSingleVarint(t *testing.T) {
|
||||
t.Run("SearchResultData topks/agg_topks", func(t *testing.T) {
|
||||
var wire []byte
|
||||
wire = protowire.AppendTag(wire, 6, protowire.VarintType) // topks
|
||||
wire = protowire.AppendVarint(wire, 3)
|
||||
wire = protowire.AppendTag(wire, 19, protowire.VarintType) // agg_topks
|
||||
wire = protowire.AppendVarint(wire, 4)
|
||||
diffDecode(t, wire, newSearchResult, decSearchResult)
|
||||
})
|
||||
t.Run("RetrieveResults sealed/global segIDs", func(t *testing.T) {
|
||||
var wire []byte
|
||||
wire = protowire.AppendTag(wire, 6, protowire.VarintType) // sealed_segmentIDs_retrieved
|
||||
wire = protowire.AppendVarint(wire, 11)
|
||||
wire = protowire.AppendTag(wire, 8, protowire.VarintType) // global_sealed_segmentIDs
|
||||
wire = protowire.AppendVarint(wire, 22)
|
||||
diffDecode(t, wire, newRetrieve, decRetrieve)
|
||||
})
|
||||
t.Run("InsertRequest hashkeys single", func(t *testing.T) {
|
||||
var wire []byte
|
||||
wire = protowire.AppendTag(wire, 6, protowire.VarintType) // hash_keys
|
||||
wire = protowire.AppendVarint(wire, 99)
|
||||
diffDecode(t, wire, newInsertRequest, decInsertRequest)
|
||||
})
|
||||
}
|
||||
|
||||
// TestMalformedInputs feeds truncated / overflowing wire data to the public entry
|
||||
// points; fastpb must report an error (matching the official codec), never panic.
|
||||
func TestMalformedInputs(t *testing.T) {
|
||||
overflowVarint := []byte{0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x02} // 10th byte > 1
|
||||
truncatedVarint := []byte{0x80, 0x80} // never terminates
|
||||
// field 2 (length-delimited) declaring length 10 but only 2 bytes follow
|
||||
truncatedBytes := append(protowire.AppendTag(nil, 2, protowire.BytesType), 0x0A, 0x01, 0x02)
|
||||
|
||||
cases := map[string][]byte{
|
||||
"overflow-varint": overflowVarint,
|
||||
"truncated-varint": truncatedVarint,
|
||||
"truncated-bytes": truncatedBytes,
|
||||
}
|
||||
for name, wire := range cases {
|
||||
t.Run("FieldData/"+name, func(t *testing.T) { diffDecode(t, wire, newFieldData, decFieldData) })
|
||||
t.Run("SearchResultData/"+name, func(t *testing.T) { diffDecode(t, wire, newSearchResult, decSearchResult) })
|
||||
t.Run("RetrieveResults/"+name, func(t *testing.T) { diffDecode(t, wire, newRetrieve, decRetrieve) })
|
||||
t.Run("InsertRequest/"+name, func(t *testing.T) { diffDecode(t, wire, newInsertRequest, decInsertRequest) })
|
||||
}
|
||||
}
|
||||
|
||||
// TestInvalidUTF8Ingress: InsertRequest is untrusted ingress, so strings are
|
||||
// UTF-8 validated. An invalid byte sequence must be rejected, matching proto3's
|
||||
// official decoder, which also rejects invalid UTF-8 in string fields.
|
||||
func TestInvalidUTF8Ingress(t *testing.T) {
|
||||
t.Run("scalar string field (DbName)", func(t *testing.T) {
|
||||
var wire []byte
|
||||
wire = protowire.AppendTag(wire, 2, protowire.BytesType) // DbName
|
||||
wire = protowire.AppendBytes(wire, []byte{0xff, 0xfe}) // invalid UTF-8
|
||||
diffDecode(t, wire, newInsertRequest, decInsertRequest)
|
||||
})
|
||||
t.Run("string inside fields_data StringArray", func(t *testing.T) {
|
||||
// fields_data(5) → ScalarField(3) → StringData(6) → StringArray data(1) = invalid utf8
|
||||
var sa []byte
|
||||
sa = protowire.AppendTag(sa, 1, protowire.BytesType)
|
||||
sa = protowire.AppendBytes(sa, []byte{0xff})
|
||||
var sf []byte
|
||||
sf = protowire.AppendTag(sf, 6, protowire.BytesType)
|
||||
sf = protowire.AppendBytes(sf, sa)
|
||||
var fd []byte
|
||||
fd = protowire.AppendTag(fd, 3, protowire.BytesType) // Scalars
|
||||
fd = protowire.AppendBytes(fd, sf)
|
||||
var wire []byte
|
||||
wire = protowire.AppendTag(wire, 5, protowire.BytesType) // fields_data
|
||||
wire = protowire.AppendBytes(wire, fd)
|
||||
diffDecode(t, wire, newInsertRequest, decInsertRequest)
|
||||
})
|
||||
}
|
||||
|
||||
// TestTryUnmarshalDispatch covers representative TryUnmarshal fast paths;
|
||||
// unsupported types report (false, nil) so the caller uses the official codec.
|
||||
func TestTryUnmarshalDispatch(t *testing.T) {
|
||||
rrWire, _ := proto.Marshal(&internalpb.RetrieveResults{ReqID: 1})
|
||||
irWire, _ := proto.Marshal(&milvuspb.InsertRequest{CollectionName: "c"})
|
||||
|
||||
handled, err := TryUnmarshal(&internalpb.RetrieveResults{}, rrWire)
|
||||
require.True(t, handled)
|
||||
require.NoError(t, err)
|
||||
|
||||
handled, err = TryUnmarshal(&milvuspb.InsertRequest{}, irWire)
|
||||
require.True(t, handled)
|
||||
require.NoError(t, err)
|
||||
|
||||
// SearchResultData and any other proto are NOT top-level fast-pathed.
|
||||
handled, err = TryUnmarshal(&schemapb.SearchResultData{}, nil)
|
||||
require.False(t, handled)
|
||||
require.NoError(t, err)
|
||||
|
||||
handled, err = TryUnmarshal(&milvuspb.SearchResults{}, nil)
|
||||
require.False(t, handled)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/reflect/protoreflect"
|
||||
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
// repValue returns a representative non-default value for a scalar field kind.
|
||||
func repValue(t *testing.T, fd protoreflect.FieldDescriptor, m protoreflect.Message) protoreflect.Value {
|
||||
switch fd.Kind() {
|
||||
case protoreflect.BoolKind:
|
||||
return protoreflect.ValueOfBool(true)
|
||||
case protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Sfixed32Kind:
|
||||
return protoreflect.ValueOfInt32(7)
|
||||
case protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Sfixed64Kind:
|
||||
return protoreflect.ValueOfInt64(7)
|
||||
case protoreflect.Uint32Kind, protoreflect.Fixed32Kind:
|
||||
return protoreflect.ValueOfUint32(7)
|
||||
case protoreflect.Uint64Kind, protoreflect.Fixed64Kind:
|
||||
return protoreflect.ValueOfUint64(7)
|
||||
case protoreflect.FloatKind:
|
||||
return protoreflect.ValueOfFloat32(1.5)
|
||||
case protoreflect.DoubleKind:
|
||||
return protoreflect.ValueOfFloat64(1.5)
|
||||
case protoreflect.StringKind:
|
||||
return protoreflect.ValueOfString("x") // valid UTF-8
|
||||
case protoreflect.BytesKind:
|
||||
return protoreflect.ValueOfBytes([]byte{1, 2, 3})
|
||||
case protoreflect.EnumKind:
|
||||
vals := fd.Enum().Values()
|
||||
if vals.Len() > 1 {
|
||||
return protoreflect.ValueOfEnum(vals.Get(1).Number())
|
||||
}
|
||||
return protoreflect.ValueOfEnum(0)
|
||||
case protoreflect.MessageKind, protoreflect.GroupKind:
|
||||
return protoreflect.ValueOfMessage(m.NewField(fd).Message())
|
||||
default:
|
||||
t.Fatalf("unhandled kind %v for %s", fd.Kind(), fd.FullName())
|
||||
return protoreflect.Value{}
|
||||
}
|
||||
}
|
||||
|
||||
func setOneField(t *testing.T, m protoreflect.Message, fd protoreflect.FieldDescriptor) {
|
||||
switch {
|
||||
case fd.IsList():
|
||||
list := m.Mutable(fd).List()
|
||||
switch fd.Kind() {
|
||||
case protoreflect.MessageKind, protoreflect.GroupKind:
|
||||
list.Append(protoreflect.ValueOfMessage(list.NewElement().Message()))
|
||||
default:
|
||||
list.Append(repValue(t, fd, m))
|
||||
}
|
||||
case fd.IsMap():
|
||||
t.Logf("skipping map field %s (none expected in decoded types)", fd.FullName())
|
||||
default:
|
||||
m.Set(fd, repValue(t, fd, m))
|
||||
}
|
||||
}
|
||||
|
||||
// assertFieldCoverage sets each field of fresh's descriptor individually,
|
||||
// marshals with the official codec, decodes with decode, and asserts proto.Equal.
|
||||
// A dropped or mis-decoded field — including a newly-added proto field — fails here.
|
||||
func assertFieldCoverage(t *testing.T, fresh proto.Message, decode func([]byte) (proto.Message, error)) {
|
||||
t.Helper()
|
||||
fields := fresh.ProtoReflect().Descriptor().Fields()
|
||||
for i := 0; i < fields.Len(); i++ {
|
||||
fd := fields.Get(i)
|
||||
msg := fresh.ProtoReflect().New()
|
||||
setOneField(t, msg, fd)
|
||||
src := msg.Interface()
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal %s: %v", fd.FullName(), err)
|
||||
}
|
||||
got, err := decode(wire)
|
||||
if err != nil {
|
||||
t.Fatalf("decode %s: %v", fd.FullName(), err)
|
||||
}
|
||||
if !proto.Equal(src, got) {
|
||||
t.Fatalf("field %s NOT round-tripped by fastpb (dropped or mis-decoded):\n src=%v\n got=%v",
|
||||
fd.FullName(), src, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCoverage_FieldData(t *testing.T) {
|
||||
assertFieldCoverage(t, &schemapb.FieldData{}, func(b []byte) (proto.Message, error) {
|
||||
m := &schemapb.FieldData{}
|
||||
return m, UnmarshalFieldData(b, m)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCoverage_ScalarField(t *testing.T) {
|
||||
assertFieldCoverage(t, &schemapb.ScalarField{}, func(b []byte) (proto.Message, error) {
|
||||
m := &schemapb.ScalarField{}
|
||||
return m, dec{}.scalarField(b, m)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCoverage_VectorField(t *testing.T) {
|
||||
assertFieldCoverage(t, &schemapb.VectorField{}, func(b []byte) (proto.Message, error) {
|
||||
m := &schemapb.VectorField{}
|
||||
return m, unmarshalVectorField(b, m)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCoverage_IDs(t *testing.T) {
|
||||
assertFieldCoverage(t, &schemapb.IDs{}, func(b []byte) (proto.Message, error) {
|
||||
m := &schemapb.IDs{}
|
||||
return m, dec{}.ids(b, m)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCoverage_SearchResultData(t *testing.T) {
|
||||
assertFieldCoverage(t, &schemapb.SearchResultData{}, func(b []byte) (proto.Message, error) {
|
||||
m := &schemapb.SearchResultData{}
|
||||
return m, UnmarshalSearchResultData(b, m)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCoverage_RetrieveResults(t *testing.T) {
|
||||
assertFieldCoverage(t, &internalpb.RetrieveResults{}, func(b []byte) (proto.Message, error) {
|
||||
m := &internalpb.RetrieveResults{}
|
||||
return m, UnmarshalRetrieveResults(b, m)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCoverage_InsertRequest(t *testing.T) {
|
||||
assertFieldCoverage(t, &milvuspb.InsertRequest{}, func(b []byte) (proto.Message, error) {
|
||||
m := &milvuspb.InsertRequest{}
|
||||
return m, UnmarshalInsertRequest(b, m)
|
||||
})
|
||||
}
|
||||
|
||||
// TestStructArrays_RoundTrips documents the previously-dropped field explicitly.
|
||||
func TestStructArrays_RoundTrips(t *testing.T) {
|
||||
roundTripFieldData(t, &schemapb.FieldData{
|
||||
Type: schemapb.DataType_ArrayOfStruct,
|
||||
FieldName: "structs",
|
||||
FieldId: 200,
|
||||
Field: &schemapb.FieldData_StructArrays{StructArrays: &schemapb.StructArrayField{
|
||||
Fields: []*schemapb.FieldData{
|
||||
{Type: schemapb.DataType_Int64, FieldName: "sub", FieldId: 201, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2}}}}}},
|
||||
},
|
||||
}},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"math"
|
||||
"math/rand"
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
)
|
||||
|
||||
// roundTripSRD pins SearchResultData decode equivalence against official proto.
|
||||
func roundTripSRD(t *testing.T, src *schemapb.SearchResultData) {
|
||||
t.Helper()
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatalf("official marshal: %v", err)
|
||||
}
|
||||
var got schemapb.SearchResultData
|
||||
if err := UnmarshalSearchResultData(wire, &got); err != nil {
|
||||
t.Fatalf("UnmarshalSearchResultData: %v", err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch:\n src = %v\n got = %v", src, &got)
|
||||
}
|
||||
}
|
||||
|
||||
// --- one explicit case per scalar data type ---
|
||||
|
||||
func TestEquiv_ScalarTypes(t *testing.T) {
|
||||
cases := map[string]*schemapb.ScalarField{
|
||||
"bool": {Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true, false, true, true}}}},
|
||||
"int": {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{0, -1, 2147483647, -2147483648, 42}}}},
|
||||
"long": {Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{0, -1, 9223372036854775807, -9223372036854775808}}}},
|
||||
"float": {Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1.5, -2.25, 0, 3e9}}}},
|
||||
"double": {Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{1.5, -2.25, 0, 3e300}}}},
|
||||
"string": {Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"a", "", "世界", "long varchar value here"}}}},
|
||||
"bytes": {Data: &schemapb.ScalarField_BytesData{BytesData: &schemapb.BytesArray{Data: [][]byte{{1, 2, 3}, {}, {255, 0, 128}}}}},
|
||||
"json": {Data: &schemapb.ScalarField_JsonData{JsonData: &schemapb.JSONArray{Data: [][]byte{[]byte(`{"a":1}`), []byte(`[]`)}}}},
|
||||
}
|
||||
for name, sf := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
roundTripFieldData(t, &schemapb.FieldData{
|
||||
FieldName: name, FieldId: 7,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: sf},
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEquiv_VectorTypes(t *testing.T) {
|
||||
cases := map[string]*schemapb.VectorField{
|
||||
"float": {Dim: 4, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6, 7, 8}}}},
|
||||
"binary": {Dim: 16, Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0xAB, 0xCD, 0x00, 0xFF}}},
|
||||
"fp16": {Dim: 2, Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4}}},
|
||||
"bf16": {Dim: 2, Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{5, 6, 7, 8}}},
|
||||
"int8": {Dim: 4, Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{250, 1, 0, 200}}},
|
||||
"sparse": {Dim: 100, Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{Dim: 100, Contents: [][]byte{{1, 2}, {3, 4, 5}}}}},
|
||||
}
|
||||
for name, vf := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
roundTripFieldData(t, &schemapb.FieldData{
|
||||
FieldName: name, FieldId: 9,
|
||||
Field: &schemapb.FieldData_Vectors{Vectors: vf},
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEquiv_FieldData_ValidData(t *testing.T) {
|
||||
roundTripFieldData(t, &schemapb.FieldData{
|
||||
Type: schemapb.DataType_Int64, FieldName: "x", FieldId: 3, IsDynamic: true,
|
||||
ValidData: []bool{true, false, false, true},
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2, 3, 4}}}}},
|
||||
})
|
||||
}
|
||||
|
||||
func TestEquiv_IDs(t *testing.T) {
|
||||
roundTripSRD(t, &schemapb.SearchResultData{Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{10, 20, 30}}}}})
|
||||
roundTripSRD(t, &schemapb.SearchResultData{Ids: &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{"k1", "k2"}}}}})
|
||||
}
|
||||
|
||||
func TestEquiv_SearchResultData_Full(t *testing.T) {
|
||||
roundTripSRD(t, &schemapb.SearchResultData{
|
||||
NumQueries: 2,
|
||||
TopK: 3,
|
||||
Scores: []float32{0.9, 0.8, 0.7, 0.6, 0.5, 0.4},
|
||||
Topks: []int64{3, 3},
|
||||
OutputFields: []string{"pk", "embedding", "title"},
|
||||
Distances: []float32{1.1, 2.2, 3.3},
|
||||
PrimaryFieldName: "pk",
|
||||
AllSearchCount: 12345,
|
||||
Ids: &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{"a", "b", "c", "d", "e", "f"}}}},
|
||||
FieldsData: []*schemapb.FieldData{
|
||||
{Type: schemapb.DataType_VarChar, FieldName: "title", FieldId: 101, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"t1", "t2", "t3", "t4", "t5", "t6"}}}}}},
|
||||
{Type: schemapb.DataType_FloatVector, FieldName: "embedding", FieldId: 102, Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: 2, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}}}}}},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// --- randomized differential fuzz: the equivalence guarantee at volume ---
|
||||
|
||||
func randString(r *rand.Rand) string {
|
||||
n := r.Intn(12)
|
||||
b := make([]byte, n)
|
||||
for i := range b {
|
||||
b[i] = byte('a' + r.Intn(26))
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func randScalarField(r *rand.Rand, rows int) *schemapb.ScalarField {
|
||||
switch r.Intn(7) {
|
||||
case 0:
|
||||
d := make([]bool, rows)
|
||||
for i := range d {
|
||||
d[i] = r.Intn(2) == 0
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: d}}}
|
||||
case 1:
|
||||
d := make([]int32, rows)
|
||||
for i := range d {
|
||||
d[i] = int32(r.Uint32())
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: d}}}
|
||||
case 2:
|
||||
d := make([]int64, rows)
|
||||
for i := range d {
|
||||
d[i] = int64(r.Uint64())
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: d}}}
|
||||
case 3:
|
||||
d := make([]float32, rows)
|
||||
for i := range d {
|
||||
d[i] = math.Float32frombits(r.Uint32())
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: d}}}
|
||||
case 4:
|
||||
d := make([]float64, rows)
|
||||
for i := range d {
|
||||
d[i] = math.Float64frombits(r.Uint64())
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: d}}}
|
||||
case 5:
|
||||
d := make([]string, rows)
|
||||
for i := range d {
|
||||
d[i] = randString(r)
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: d}}}
|
||||
default:
|
||||
d := make([][]byte, rows)
|
||||
for i := range d {
|
||||
d[i] = []byte(randString(r))
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_BytesData{BytesData: &schemapb.BytesArray{Data: d}}}
|
||||
}
|
||||
}
|
||||
|
||||
func randFieldData(r *rand.Rand, rows int) *schemapb.FieldData {
|
||||
fd := &schemapb.FieldData{FieldName: randString(r), FieldId: int64(r.Intn(1000))}
|
||||
if r.Intn(2) == 0 {
|
||||
fd.Field = &schemapb.FieldData_Scalars{Scalars: randScalarField(r, rows)}
|
||||
} else {
|
||||
dim := 1 + r.Intn(8)
|
||||
d := make([]float32, rows*dim)
|
||||
for i := range d {
|
||||
d[i] = math.Float32frombits(r.Uint32())
|
||||
}
|
||||
fd.Field = &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: int64(dim), Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: d}}}}
|
||||
}
|
||||
if r.Intn(2) == 0 {
|
||||
vd := make([]bool, rows)
|
||||
for i := range vd {
|
||||
vd[i] = r.Intn(2) == 0
|
||||
}
|
||||
fd.ValidData = vd
|
||||
}
|
||||
return fd
|
||||
}
|
||||
|
||||
func randSRD(r *rand.Rand) *schemapb.SearchResultData {
|
||||
rows := r.Intn(20)
|
||||
srd := &schemapb.SearchResultData{
|
||||
NumQueries: int64(r.Intn(5)),
|
||||
TopK: int64(r.Intn(10)),
|
||||
PrimaryFieldName: randString(r),
|
||||
AllSearchCount: int64(r.Uint32()),
|
||||
}
|
||||
srd.Scores = make([]float32, rows)
|
||||
for i := range srd.Scores {
|
||||
srd.Scores[i] = math.Float32frombits(r.Uint32())
|
||||
}
|
||||
if r.Intn(2) == 0 {
|
||||
ids := make([]int64, rows)
|
||||
for i := range ids {
|
||||
ids[i] = int64(r.Uint64())
|
||||
}
|
||||
srd.Ids = &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: ids}}}
|
||||
} else {
|
||||
ids := make([]string, rows)
|
||||
for i := range ids {
|
||||
ids[i] = randString(r)
|
||||
}
|
||||
srd.Ids = &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: ids}}}
|
||||
}
|
||||
for i := 0; i < r.Intn(4); i++ {
|
||||
srd.FieldsData = append(srd.FieldsData, randFieldData(r, rows))
|
||||
}
|
||||
for i := 0; i < r.Intn(3); i++ {
|
||||
srd.OutputFields = append(srd.OutputFields, randString(r))
|
||||
}
|
||||
return srd
|
||||
}
|
||||
|
||||
func TestEquiv_Fuzz(t *testing.T) {
|
||||
r := rand.New(rand.NewSource(0xC0FFEE))
|
||||
for i := 0; i < 5000; i++ {
|
||||
src := randSRD(r)
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal iter %d: %v", i, err)
|
||||
}
|
||||
var got schemapb.SearchResultData
|
||||
if err := UnmarshalSearchResultData(wire, &got); err != nil {
|
||||
t.Fatalf("decode iter %d: %v", i, err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch iter %d:\n src = %v\n got = %v", i, src, &got)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,303 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
// --- tiny wire builders (each comment says which wire rule the bytes violate) ---
|
||||
|
||||
// wtag emits just a field tag (no value) — a value that should follow is missing.
|
||||
func wtag(num int32, wt protowire.Type) []byte {
|
||||
return protowire.AppendTag(nil, protowire.Number(num), wt)
|
||||
}
|
||||
|
||||
// wfield emits a well-formed length-delimited field wrapping payload.
|
||||
func wfield(num int32, payload []byte) []byte {
|
||||
return protowire.AppendBytes(protowire.AppendTag(nil, protowire.Number(num), protowire.BytesType), payload)
|
||||
}
|
||||
|
||||
func cat(parts ...[]byte) []byte {
|
||||
var out []byte
|
||||
for _, p := range parts {
|
||||
out = append(out, p...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func f32le(vals ...float32) []byte {
|
||||
var b []byte
|
||||
for _, v := range vals {
|
||||
b = protowire.AppendFixed32(b, math.Float32bits(v))
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func f64le(vals ...float64) []byte {
|
||||
var b []byte
|
||||
for _, v := range vals {
|
||||
b = protowire.AppendFixed64(b, math.Float64bits(v))
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// truncTag is a varint (tag or value) whose continuation bit never terminates.
|
||||
var truncTag = []byte{0x80}
|
||||
|
||||
// TestMalformedWirePerField drives every decoder over hand-crafted wire bytes
|
||||
// that violate one specific wire rule each (truncated varints, truncated
|
||||
// length prefixes, length prefixes overrunning the buffer, bad packed payload
|
||||
// sizes, invalid wire types, malformed delegated submessages). Every case is a
|
||||
// differential assertion: fastpb must match the official codec both in error
|
||||
// behavior and, when the official codec accepts (wire-type-mismatch fallbacks,
|
||||
// unpacked encodings of packed fields), in the decoded message.
|
||||
func TestMalformedWirePerField(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
wire []byte
|
||||
fresh func() proto.Message
|
||||
fast func([]byte, proto.Message) error
|
||||
}{
|
||||
// --- FieldData: one truncation per field ---
|
||||
{"FieldData/type-truncated-varint", cat(wtag(1, protowire.VarintType), truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/fieldname-truncated-len-prefix", cat(wtag(2, protowire.BytesType), truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/fieldname-as-varint-fallback", protowire.AppendVarint(wtag(2, protowire.VarintType), 5), newFieldData, decFieldData},
|
||||
{"FieldData/scalars-len-overruns-buffer", cat(wtag(3, protowire.BytesType), []byte{0x05, 0x01}), newFieldData, decFieldData},
|
||||
{"FieldData/scalars-as-varint-fallback", protowire.AppendVarint(wtag(3, protowire.VarintType), 5), newFieldData, decFieldData},
|
||||
{"FieldData/vectors-len-overruns-buffer", cat(wtag(4, protowire.BytesType), []byte{0x05}), newFieldData, decFieldData},
|
||||
{"FieldData/vectors-malformed-submsg", wfield(4, truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/fieldid-truncated-varint", cat(wtag(5, protowire.VarintType), truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/isdynamic-truncated-varint", cat(wtag(6, protowire.VarintType), truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/validdata-truncated-len-prefix", cat(wtag(7, protowire.BytesType), truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/validdata-truncated-packed-varint", wfield(7, truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/validdata-single-varint-ok", protowire.AppendVarint(wtag(7, protowire.VarintType), 1), newFieldData, decFieldData},
|
||||
{"FieldData/validdata-single-varint-truncated", cat(wtag(7, protowire.VarintType), truncTag), newFieldData, decFieldData},
|
||||
// ValidData accepts varint and bytes; fixed32/fixed64 must fall back to the
|
||||
// official codec (which keeps them as unknown fields) instead of misreading
|
||||
// the payload as varints.
|
||||
{"FieldData/validdata-fixed32-fallback", protowire.AppendFixed32(wtag(7, protowire.Fixed32Type), 1), newFieldData, decFieldData},
|
||||
{"FieldData/validdata-fixed64-fallback", protowire.AppendFixed64(wtag(7, protowire.Fixed64Type), 1), newFieldData, decFieldData},
|
||||
{"FieldData/structarrays-len-overruns-buffer", cat(wtag(8, protowire.BytesType), []byte{0x05}), newFieldData, decFieldData},
|
||||
{"FieldData/structarrays-malformed-submsg", wfield(8, truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/unknown-truncated-varint", cat(wtag(99, protowire.VarintType), truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/unknown-fixed64-short", cat(wtag(99, protowire.Fixed64Type), []byte{1, 2, 3, 4}), newFieldData, decFieldData},
|
||||
{"FieldData/unknown-fixed32-short", cat(wtag(99, protowire.Fixed32Type), []byte{1, 2}), newFieldData, decFieldData},
|
||||
{"FieldData/unknown-lendelim-truncated-len", cat(wtag(99, protowire.BytesType), truncTag), newFieldData, decFieldData},
|
||||
{"FieldData/unknown-lendelim-overruns-buffer", cat(wtag(99, protowire.BytesType), []byte{0x05, 0x01}), newFieldData, decFieldData},
|
||||
{"FieldData/invalid-wire-type-6", wtag(99, protowire.Type(6)), newFieldData, decFieldData},
|
||||
|
||||
// --- ScalarField: per-variant malformed payloads (leaf array decoders) ---
|
||||
{"Scalar/truncated-tag", truncTag, newScalarField, decScalarField},
|
||||
{"Scalar/unknown-varint-truncated", cat(wtag(20, protowire.VarintType), truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/member-len-overruns-buffer", cat(wtag(1, protowire.BytesType), []byte{0x05}), newScalarField, decScalarField},
|
||||
// BoolArray (field 1): every decodePackedBool branch
|
||||
{"Scalar/bool-truncated-inner-tag", wfield(1, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/bool-truncated-packed-len", wfield(1, cat(wtag(1, protowire.BytesType), truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/bool-truncated-packed-varint", wfield(1, wfield(1, truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/bool-unpacked-varint-ok", wfield(1, protowire.AppendVarint(wtag(1, protowire.VarintType), 1)), newScalarField, decScalarField},
|
||||
{"Scalar/bool-unpacked-varint-truncated", wfield(1, wtag(1, protowire.VarintType)), newScalarField, decScalarField},
|
||||
// IntArray (field 2): decodePackedI32 branches
|
||||
{"Scalar/int-truncated-inner-tag", wfield(2, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/int-truncated-packed-len", wfield(2, cat(wtag(1, protowire.BytesType), truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/int-truncated-packed-varint", wfield(2, wfield(1, truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/int-unpacked-varint-ok", wfield(2, protowire.AppendVarint(wtag(1, protowire.VarintType), 7)), newScalarField, decScalarField},
|
||||
{"Scalar/int-unpacked-varint-truncated", wfield(2, wtag(1, protowire.VarintType)), newScalarField, decScalarField},
|
||||
// LongArray (field 3): decodePackedI64 branches
|
||||
{"Scalar/long-truncated-inner-tag", wfield(3, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/long-truncated-packed-len", wfield(3, cat(wtag(1, protowire.BytesType), truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/long-truncated-packed-varint", wfield(3, wfield(1, truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/long-unpacked-varint-ok", wfield(3, protowire.AppendVarint(wtag(1, protowire.VarintType), 9)), newScalarField, decScalarField},
|
||||
{"Scalar/long-unpacked-varint-truncated", wfield(3, wtag(1, protowire.VarintType)), newScalarField, decScalarField},
|
||||
// FloatArray (field 4): decodePackedF32 branches
|
||||
{"Scalar/float-truncated-inner-tag", wfield(4, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/float-truncated-packed-len", wfield(4, cat(wtag(1, protowire.BytesType), truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/float-packed-len-not-multiple-of-4", wfield(4, wfield(1, []byte{1, 2, 3})), newScalarField, decScalarField},
|
||||
{"Scalar/float-two-packed-chunks-ok", wfield(4, cat(wfield(1, f32le(1.5, 2.5)), wfield(1, f32le(3.5)))), newScalarField, decScalarField},
|
||||
{"Scalar/float-single-fixed32-short", wfield(4, cat(wtag(1, protowire.Fixed32Type), []byte{1, 2})), newScalarField, decScalarField},
|
||||
{"Scalar/float-data-as-varint-fallback", wfield(4, protowire.AppendVarint(wtag(1, protowire.VarintType), 7)), newScalarField, decScalarField},
|
||||
// DoubleArray (field 5): decodePackedF64 branches
|
||||
{"Scalar/double-truncated-inner-tag", wfield(5, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/double-truncated-packed-len", wfield(5, cat(wtag(1, protowire.BytesType), truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/double-packed-len-not-multiple-of-8", wfield(5, wfield(1, []byte{1, 2, 3})), newScalarField, decScalarField},
|
||||
{"Scalar/double-two-packed-chunks-ok", wfield(5, cat(wfield(1, f64le(1.5)), wfield(1, f64le(2.5)))), newScalarField, decScalarField},
|
||||
{"Scalar/double-single-fixed64-short", wfield(5, cat(wtag(1, protowire.Fixed64Type), []byte{1, 2})), newScalarField, decScalarField},
|
||||
{"Scalar/double-data-as-varint-fallback", wfield(5, protowire.AppendVarint(wtag(1, protowire.VarintType), 7)), newScalarField, decScalarField},
|
||||
// StringArray (field 6): strings arena pass-1 error branches
|
||||
{"Scalar/string-truncated-inner-tag", wfield(6, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/string-truncated-len-prefix", wfield(6, cat(wtag(1, protowire.BytesType), truncTag)), newScalarField, decScalarField},
|
||||
// BytesArray (field 7) / JSONArray (field 9): decodeRepeatedBytes branches
|
||||
{"Scalar/bytes-truncated-inner-tag", wfield(7, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/bytes-truncated-len-prefix", wfield(7, cat(wtag(1, protowire.BytesType), truncTag)), newScalarField, decScalarField},
|
||||
{"Scalar/json-truncated-inner-tag", wfield(9, truncTag), newScalarField, decScalarField},
|
||||
// cold oneof variants (decodeScalarFallback 8/10/11/12/13/14): malformed
|
||||
// submessage bytes must surface the official codec's decode error
|
||||
{"Scalar/array-data-malformed", wfield(8, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/geometry-data-malformed", wfield(10, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/timestamptz-data-malformed", wfield(11, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/geometry-wkt-data-malformed", wfield(12, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/mol-data-malformed", wfield(13, truncTag), newScalarField, decScalarField},
|
||||
{"Scalar/mol-smiles-data-malformed", wfield(14, truncTag), newScalarField, decScalarField},
|
||||
|
||||
// --- VectorField: one malformed case per field ---
|
||||
{"Vector/truncated-tag", truncTag, newVectorField, decVectorField},
|
||||
{"Vector/dim-truncated-varint", cat(wtag(1, protowire.VarintType), truncTag), newVectorField, decVectorField},
|
||||
{"Vector/floatvector-len-overruns-buffer", cat(wtag(2, protowire.BytesType), []byte{0x05}), newVectorField, decVectorField},
|
||||
{"Vector/floatvector-malformed", wfield(2, truncTag), newVectorField, decVectorField},
|
||||
{"Vector/binary-truncated-len-prefix", cat(wtag(3, protowire.BytesType), truncTag), newVectorField, decVectorField},
|
||||
{"Vector/float16-len-overruns-buffer", cat(wtag(4, protowire.BytesType), []byte{0x05}), newVectorField, decVectorField},
|
||||
{"Vector/bfloat16-len-overruns-buffer", cat(wtag(5, protowire.BytesType), []byte{0x05}), newVectorField, decVectorField},
|
||||
{"Vector/sparse-len-overruns-buffer", cat(wtag(6, protowire.BytesType), []byte{0x05}), newVectorField, decVectorField},
|
||||
{"Vector/sparse-truncated-inner-tag", wfield(6, truncTag), newVectorField, decVectorField},
|
||||
{"Vector/sparse-contents-truncated-len", wfield(6, cat(wtag(1, protowire.BytesType), truncTag)), newVectorField, decVectorField},
|
||||
{"Vector/sparse-dim-truncated-varint", wfield(6, wtag(2, protowire.VarintType)), newVectorField, decVectorField},
|
||||
{"Vector/int8-truncated-len-prefix", cat(wtag(7, protowire.BytesType), truncTag), newVectorField, decVectorField},
|
||||
{"Vector/vectorarray-len-overruns-buffer", cat(wtag(8, protowire.BytesType), []byte{0x05}), newVectorField, decVectorField},
|
||||
{"Vector/vectorarray-malformed-submsg", wfield(8, truncTag), newVectorField, decVectorField},
|
||||
{"Vector/unknown-truncated-varint", cat(wtag(20, protowire.VarintType), truncTag), newVectorField, decVectorField},
|
||||
|
||||
// --- IDs ---
|
||||
{"IDs/truncated-tag", truncTag, newIDs, decIDs},
|
||||
{"IDs/unknown-varint-truncated", cat(wtag(9, protowire.VarintType), truncTag), newIDs, decIDs},
|
||||
{"IDs/member-len-overruns-buffer", cat(wtag(1, protowire.BytesType), []byte{0x05}), newIDs, decIDs},
|
||||
{"IDs/intid-malformed", wfield(1, truncTag), newIDs, decIDs},
|
||||
{"IDs/strid-malformed", wfield(2, truncTag), newIDs, decIDs},
|
||||
|
||||
// --- SearchResultData: error propagation per field ---
|
||||
{"SRD/numqueries-truncated-varint", cat(wtag(1, protowire.VarintType), truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/unknown-fixed64-short", cat(wtag(99, protowire.Fixed64Type), []byte{1, 2}), newSearchResult, decSearchResult},
|
||||
{"SRD/fieldsdata-malformed", wfield(3, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/scores-len-not-multiple-of-4", wfield(4, []byte{1, 2, 3}), newSearchResult, decSearchResult},
|
||||
{"SRD/scores-two-packed-chunks-ok", cat(wfield(4, f32le(0.5)), wfield(4, f32le(1.5, 2.5))), newSearchResult, decSearchResult},
|
||||
{"SRD/ids-malformed", wfield(5, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/topks-truncated-packed-varint", wfield(6, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/groupby-malformed", wfield(8, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/distances-len-not-multiple-of-4", wfield(10, []byte{1, 2, 3}), newSearchResult, decSearchResult},
|
||||
{"SRD/iterator-malformed", wfield(11, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/recalls-len-not-multiple-of-4", wfield(12, []byte{1, 2, 3}), newSearchResult, decSearchResult},
|
||||
{"SRD/highlight-malformed", wfield(14, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/elementindices-malformed", wfield(15, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/groupbyvalues-malformed", wfield(17, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/aggbuckets-malformed", wfield(18, truncTag), newSearchResult, decSearchResult},
|
||||
{"SRD/aggtopks-truncated-packed-varint", wfield(19, truncTag), newSearchResult, decSearchResult},
|
||||
|
||||
// --- RetrieveResults: error propagation per field ---
|
||||
{"RR/reqid-truncated-varint", cat(wtag(3, protowire.VarintType), truncTag), newRetrieve, decRetrieve},
|
||||
{"RR/unknown-fixed32-short", cat(wtag(99, protowire.Fixed32Type), []byte{1}), newRetrieve, decRetrieve},
|
||||
{"RR/base-malformed", wfield(1, truncTag), newRetrieve, decRetrieve},
|
||||
{"RR/status-malformed", wfield(2, truncTag), newRetrieve, decRetrieve},
|
||||
{"RR/ids-malformed", wfield(4, truncTag), newRetrieve, decRetrieve},
|
||||
{"RR/fieldsdata-malformed", wfield(5, truncTag), newRetrieve, decRetrieve},
|
||||
{"RR/sealedsegids-truncated-packed-varint", wfield(6, truncTag), newRetrieve, decRetrieve},
|
||||
{"RR/globalsegids-truncated-packed-varint", wfield(8, truncTag), newRetrieve, decRetrieve},
|
||||
{"RR/cost-malformed", wfield(13, truncTag), newRetrieve, decRetrieve},
|
||||
{"RR/elementindices-malformed", wfield(19, truncTag), newRetrieve, decRetrieve},
|
||||
|
||||
// --- InsertRequest: error propagation per field ---
|
||||
{"IR/numrows-truncated-varint", cat(wtag(7, protowire.VarintType), truncTag), newInsertRequest, decInsertRequest},
|
||||
{"IR/unknown-fixed64-short", cat(wtag(99, protowire.Fixed64Type), []byte{1, 2, 3}), newInsertRequest, decInsertRequest},
|
||||
{"IR/base-malformed", wfield(1, truncTag), newInsertRequest, decInsertRequest},
|
||||
{"IR/fieldsdata-malformed", wfield(5, truncTag), newInsertRequest, decInsertRequest},
|
||||
{"IR/hashkeys-truncated-packed-varint", wfield(6, truncTag), newInsertRequest, decInsertRequest},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
diffDecode(t, c.wire, c.fresh, c.fast)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestNestedGroupFallbacks places a well-formed proto2 group (wire types 3/4,
|
||||
// which proto3 never emits) inside every nested decoder that has its own
|
||||
// errProto2 check. The error must propagate to the public entry point, which
|
||||
// redoes the decode with the official codec — so the result must equal
|
||||
// proto.Unmarshal exactly (the group is preserved as an unknown field).
|
||||
func TestNestedGroupFallbacks(t *testing.T) {
|
||||
canonFD, err := proto.Marshal(&schemapb.FieldData{FieldName: "f", FieldId: 3})
|
||||
require.NoError(t, err)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
wire []byte
|
||||
fresh func() proto.Message
|
||||
fast func([]byte, proto.Message) error
|
||||
}{
|
||||
{"FieldData/top-level", appendGroup(append([]byte{}, canonFD...), 999), newFieldData, decFieldData},
|
||||
{"FieldData/in-scalarfield", wfield(3, appendGroup(nil, 20)), newFieldData, decFieldData},
|
||||
{"FieldData/in-boolarray", wfield(3, wfield(1, appendGroup(nil, 15))), newFieldData, decFieldData},
|
||||
{"FieldData/in-intarray", wfield(3, wfield(2, appendGroup(nil, 15))), newFieldData, decFieldData},
|
||||
{"FieldData/in-longarray", wfield(3, wfield(3, appendGroup(nil, 15))), newFieldData, decFieldData},
|
||||
{"FieldData/in-floatarray", wfield(3, wfield(4, appendGroup(nil, 15))), newFieldData, decFieldData},
|
||||
{"FieldData/in-doublearray", wfield(3, wfield(5, appendGroup(nil, 15))), newFieldData, decFieldData},
|
||||
{"FieldData/in-vectorfield", wfield(4, appendGroup(nil, 20)), newFieldData, decFieldData},
|
||||
{"FieldData/in-sparsefloatarray", wfield(4, wfield(6, appendGroup(nil, 9))), newFieldData, decFieldData},
|
||||
{"SearchResultData/in-ids", wfield(5, appendGroup(nil, 9)), newSearchResult, decSearchResult},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
diffDecode(t, c.wire, c.fresh, c.fast)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestUTF8ErrorBranches covers the remaining invalid-UTF-8 rejection branches.
|
||||
// InsertRequest is untrusted ingress, so official proto3 behavior (reject) must
|
||||
// be matched field by field.
|
||||
func TestUTF8ErrorBranches(t *testing.T) {
|
||||
bad := []byte{0xff, 0xfe} // not a valid UTF-8 sequence
|
||||
cases := map[string][]byte{
|
||||
"collection-name": wfield(3, bad),
|
||||
"partition-name": wfield(4, bad),
|
||||
"namespace": wfield(9, bad),
|
||||
"nested-fielddata-name": wfield(5, wfield(2, bad)), // fields_data → FieldData.field_name
|
||||
}
|
||||
for name, wire := range cases {
|
||||
t.Run("InsertRequest/"+name, func(t *testing.T) {
|
||||
diffDecode(t, wire, newInsertRequest, decInsertRequest)
|
||||
})
|
||||
}
|
||||
|
||||
// The internal result decoders run with utf8=false via the public entry
|
||||
// points; exercise their validation branches directly with a utf8 dec so
|
||||
// the code stays correct if a validated entry point is ever added.
|
||||
t.Run("searchResultData/output-fields", func(t *testing.T) {
|
||||
srd := &schemapb.SearchResultData{}
|
||||
err := dec{utf8: true}.searchResultData(wfield(7, bad), srd)
|
||||
require.ErrorIs(t, err, errInvalidUTF8)
|
||||
})
|
||||
t.Run("searchResultData/primary-field-name", func(t *testing.T) {
|
||||
srd := &schemapb.SearchResultData{}
|
||||
err := dec{utf8: true}.searchResultData(wfield(13, bad), srd)
|
||||
require.ErrorIs(t, err, errInvalidUTF8)
|
||||
})
|
||||
t.Run("retrieveResults/channel-ids", func(t *testing.T) {
|
||||
rr := &internalpb.RetrieveResults{}
|
||||
err := dec{utf8: true}.retrieveResults(wfield(7, bad), rr)
|
||||
require.ErrorIs(t, err, errInvalidUTF8)
|
||||
})
|
||||
}
|
||||
|
||||
// TestTrustedDecodersSkipUTF8Validation pins the intended divergence from the
|
||||
// official codec: internal/trusted result decoders do NOT validate UTF-8 (the
|
||||
// bytes were produced by Milvus itself), while official proto3 rejects them.
|
||||
func TestTrustedDecodersSkipUTF8Validation(t *testing.T) {
|
||||
bad := []byte{0xff}
|
||||
|
||||
srdWire := wfield(7, bad) // output_fields
|
||||
require.Error(t, proto.Unmarshal(srdWire, &schemapb.SearchResultData{}), "official proto3 rejects invalid UTF-8")
|
||||
srd := &schemapb.SearchResultData{}
|
||||
require.NoError(t, UnmarshalSearchResultData(srdWire, srd), "trusted fastpb path skips validation by design")
|
||||
require.Equal(t, []string{"\xff"}, srd.OutputFields)
|
||||
|
||||
rrWire := wfield(7, bad) // channelIDs_retrieved
|
||||
require.Error(t, proto.Unmarshal(rrWire, &internalpb.RetrieveResults{}))
|
||||
rr := &internalpb.RetrieveResults{}
|
||||
require.NoError(t, UnmarshalRetrieveResults(rrWire, rr))
|
||||
require.Equal(t, []string{"\xff"}, rr.ChannelIDsRetrieved)
|
||||
}
|
||||
@@ -0,0 +1,539 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
)
|
||||
|
||||
var (
|
||||
errMalformed = errors.New("fastpb: malformed wire data")
|
||||
errInvalidUTF8 = errors.New("fastpb: invalid UTF-8 in string field")
|
||||
// errProto2 signals a proto2-only wire feature — a group, encoded with the
|
||||
// start-group (3) / end-group (4) wire types, which proto3 does not have.
|
||||
// fastpb only decodes canonical proto3; the public entry points catch this
|
||||
// and fall back to the official codec, which handles groups. Milvus itself
|
||||
// only ever produces proto3, so this never fires on trusted traffic — it is a
|
||||
// robustness path for non-Milvus / adversarial encoders.
|
||||
errProto2 = errors.New("fastpb: proto2 group wire type, fallback to official codec")
|
||||
)
|
||||
|
||||
// isProto2Group reports whether wtype is a proto2 group delimiter (start=3,
|
||||
// end=4). Encountering one means the whole message must be decoded by the
|
||||
// official codec rather than fast-pathed.
|
||||
func isProto2Group(wtype int) bool { return wtype == 3 || wtype == 4 }
|
||||
|
||||
// fallbackOnProto2 reports whether err means "fastpb hit a proto2 group". The
|
||||
// public entry points use it to redo the decode with the official codec.
|
||||
func fallbackOnProto2(err error) bool { return errors.Is(err, errProto2) }
|
||||
|
||||
// dec carries decode options down the recursive string-bearing decoders.
|
||||
// utf8=true validates every decoded string (used on untrusted ingress, e.g.
|
||||
// client InsertRequest) to match official proto3 behavior; internal/trusted
|
||||
// result messages decode with utf8=false for maximum speed.
|
||||
type dec struct {
|
||||
utf8 bool
|
||||
}
|
||||
|
||||
// str converts a wire byte span to a string, validating UTF-8 if requested.
|
||||
func (d dec) str(v []byte) (string, error) {
|
||||
if d.utf8 && !validUTF8(v) {
|
||||
return "", errInvalidUTF8
|
||||
}
|
||||
return string(v), nil
|
||||
}
|
||||
|
||||
// UnmarshalFieldData decodes a wire-format schemapb.FieldData into fd
|
||||
// (internal/trusted; no UTF-8 validation).
|
||||
func UnmarshalFieldData(b []byte, fd *schemapb.FieldData) error {
|
||||
proto.Reset(fd) // match official proto.Unmarshal: clear target before decode
|
||||
if err := (dec{}).fieldData(b, fd); err != nil {
|
||||
if fallbackOnProto2(err) {
|
||||
proto.Reset(fd) // discard partial decode before the authoritative pass
|
||||
return proto.Unmarshal(b, fd)
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d dec) fieldData(b []byte, fd *schemapb.FieldData) error {
|
||||
full := b
|
||||
var rest []byte
|
||||
oneofNum := 0
|
||||
for len(b) > 0 {
|
||||
start := b
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
// Known field with an unexpected wire type → official codec, to match
|
||||
// proto.Unmarshal instead of misdecoding the wrong shape. Field 7 (ValidData)
|
||||
// is a packed repeated bool and legitimately accepts varint (unpacked) and
|
||||
// bytes (packed); any other wire type falls back like the rest.
|
||||
if (num == 1 || num == 5 || num == 6) && wtype != 0 {
|
||||
return fallbackUnmarshal(full, fd)
|
||||
}
|
||||
if (num == 2 || num == 3 || num == 4 || num == 8) && wtype != 2 {
|
||||
return fallbackUnmarshal(full, fd)
|
||||
}
|
||||
if num == 7 && wtype != 0 && wtype != 2 {
|
||||
return fallbackUnmarshal(full, fd)
|
||||
}
|
||||
if num == 3 || num == 4 || num == 8 {
|
||||
if oneofNum == num {
|
||||
return fallbackUnmarshal(full, fd)
|
||||
}
|
||||
oneofNum = num
|
||||
}
|
||||
switch num {
|
||||
case 1: // Type (enum, varint)
|
||||
v, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
fd.Type = schemapb.DataType(v)
|
||||
b = b[n:]
|
||||
case 2: // FieldName (string)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fd.FieldName = s
|
||||
b = b[n:]
|
||||
case 3: // Scalars (oneof, nested ScalarField)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
sf := &schemapb.ScalarField{}
|
||||
if err := d.scalarField(v, sf); err != nil {
|
||||
return err
|
||||
}
|
||||
fd.Field = &schemapb.FieldData_Scalars{Scalars: sf}
|
||||
b = b[n:]
|
||||
case 4: // Vectors (oneof, nested VectorField)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
vf := &schemapb.VectorField{}
|
||||
if err := unmarshalVectorField(v, vf); err != nil {
|
||||
return err
|
||||
}
|
||||
fd.Field = &schemapb.FieldData_Vectors{Vectors: vf}
|
||||
b = b[n:]
|
||||
case 5: // FieldId (varint)
|
||||
v, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
fd.FieldId = int64(v)
|
||||
b = b[n:]
|
||||
case 6: // IsDynamic (bool, varint)
|
||||
v, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
fd.IsDynamic = v != 0
|
||||
b = b[n:]
|
||||
case 7: // ValidData (repeated bool, packed)
|
||||
if wtype == 2 {
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
if err := appendPackedBool(v, &fd.ValidData); err != nil {
|
||||
return err
|
||||
}
|
||||
b = b[n:]
|
||||
} else {
|
||||
v, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
fd.ValidData = append(fd.ValidData, v != 0)
|
||||
b = b[n:]
|
||||
}
|
||||
case 8: // StructArrays (oneof) — delegate to official; must stay in this
|
||||
// pass so oneof last-wins ordering matches official.
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
m := &schemapb.StructArrayField{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
fd.Field = &schemapb.FieldData_StructArrays{StructArrays: m}
|
||||
b = b[n:]
|
||||
default: // unhandled (future) non-oneof field → fold into official merge
|
||||
n2 := skipField(b, wtype)
|
||||
if n2 <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
rest = append(rest, start[:n+n2]...)
|
||||
b = b[n2:]
|
||||
}
|
||||
}
|
||||
if len(rest) > 0 {
|
||||
return protoMerge(rest, fd)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// scalarField decodes a wire-format schemapb.ScalarField into sf.
|
||||
// Hot array types (bool/int/long/float/double/string/bytes/json) are hand-decoded;
|
||||
// rarer variants (array/geometry/timestamptz/mol/...) fall back to the official codec.
|
||||
func (d dec) scalarField(b []byte, sf *schemapb.ScalarField) error {
|
||||
full := b
|
||||
var rest []byte
|
||||
oneofNum := 0
|
||||
for len(b) > 0 {
|
||||
start := b
|
||||
num, wtype, tn := consumeTag(b)
|
||||
if tn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[tn:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if wtype != 2 { // members are all length-delimited; anything else is unknown
|
||||
sn := skipField(b, wtype)
|
||||
if sn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
rest = append(rest, start[:tn+sn]...)
|
||||
b = b[sn:]
|
||||
continue
|
||||
}
|
||||
v, vn := consumeBytes(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
if num >= 1 && num <= 14 {
|
||||
if oneofNum == num {
|
||||
return fallbackUnmarshal(full, sf)
|
||||
}
|
||||
oneofNum = num
|
||||
}
|
||||
switch num {
|
||||
case 1: // BoolData
|
||||
a := &schemapb.BoolArray{}
|
||||
if err := decodePackedBool(v, &a.Data, a); err != nil {
|
||||
return err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_BoolData{BoolData: a}
|
||||
case 2: // IntData
|
||||
a := &schemapb.IntArray{}
|
||||
if err := decodePackedI32(v, &a.Data, a); err != nil {
|
||||
return err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_IntData{IntData: a}
|
||||
case 3: // LongData
|
||||
a := &schemapb.LongArray{}
|
||||
if err := decodePackedI64(v, &a.Data, a); err != nil {
|
||||
return err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_LongData{LongData: a}
|
||||
case 4: // FloatData
|
||||
a := &schemapb.FloatArray{}
|
||||
if err := decodePackedF32(v, &a.Data, a); err != nil {
|
||||
return err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_FloatData{FloatData: a}
|
||||
case 5: // DoubleData
|
||||
a := &schemapb.DoubleArray{}
|
||||
if err := decodePackedF64(v, &a.Data, a); err != nil {
|
||||
return err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_DoubleData{DoubleData: a}
|
||||
case 6: // StringData
|
||||
sa := &schemapb.StringArray{}
|
||||
if err := d.stringArray(v, sa); err != nil {
|
||||
return err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_StringData{StringData: sa}
|
||||
case 7: // BytesData
|
||||
a := &schemapb.BytesArray{}
|
||||
if err := decodeRepeatedBytes(v, &a.Data, a); err != nil {
|
||||
return err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_BytesData{BytesData: a}
|
||||
case 9: // JsonData
|
||||
a := &schemapb.JSONArray{}
|
||||
if err := decodeRepeatedBytes(v, &a.Data, a); err != nil {
|
||||
return err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_JsonData{JsonData: a}
|
||||
default: // known cold oneof variants (array/geometry/mol/...) stay in-pass;
|
||||
// truly-unknown fields fold into official merge.
|
||||
handled, err := decodeScalarFallback(num, v, sf)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !handled {
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(rest) > 0 {
|
||||
return protoMerge(rest, sf)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// stringArray decodes a wire-format schemapb.StringArray into sa
|
||||
// (repeated string data = 1) via a single byte arena (N allocations → 2).
|
||||
func (d dec) stringArray(b []byte, sa *schemapb.StringArray) error {
|
||||
strs, ok, err := d.strings(b)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return fallbackUnmarshal(b, sa) // unknown field → official codec, preserves it
|
||||
}
|
||||
sa.Data = strs
|
||||
return nil
|
||||
}
|
||||
|
||||
// unmarshalVectorField decodes a wire-format schemapb.VectorField into vf
|
||||
// (no string fields, so no UTF-8 validation needed).
|
||||
func unmarshalVectorField(b []byte, vf *schemapb.VectorField) error {
|
||||
full := b
|
||||
var rest []byte
|
||||
oneofNum := 0
|
||||
for len(b) > 0 {
|
||||
start := b
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
// Known field with an unexpected wire type (Dim=varint; 2..8=length-delimited)
|
||||
// → official codec, to match proto.Unmarshal instead of misdecoding it.
|
||||
if (num == 1 && wtype != 0) || (num >= 2 && num <= 8 && wtype != 2) {
|
||||
return fallbackUnmarshal(full, vf)
|
||||
}
|
||||
if num >= 2 && num <= 8 {
|
||||
if oneofNum == num && (num == 2 || num == 6 || num == 8) {
|
||||
return fallbackUnmarshal(full, vf)
|
||||
}
|
||||
oneofNum = num
|
||||
}
|
||||
switch num {
|
||||
case 1: // Dim (varint)
|
||||
v, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
vf.Dim = int64(v)
|
||||
b = b[n:]
|
||||
case 2: // FloatVector (nested FloatArray)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
fa := &schemapb.FloatArray{}
|
||||
if err := unmarshalFloatArray(v, fa); err != nil {
|
||||
return err
|
||||
}
|
||||
vf.Data = &schemapb.VectorField_FloatVector{FloatVector: fa}
|
||||
b = b[n:]
|
||||
case 3: // BinaryVector (bytes)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
vf.Data = &schemapb.VectorField_BinaryVector{BinaryVector: cloneBytes(v)}
|
||||
b = b[n:]
|
||||
case 4: // Float16Vector (bytes)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
vf.Data = &schemapb.VectorField_Float16Vector{Float16Vector: cloneBytes(v)}
|
||||
b = b[n:]
|
||||
case 5: // Bfloat16Vector (bytes)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
vf.Data = &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: cloneBytes(v)}
|
||||
b = b[n:]
|
||||
case 6: // SparseFloatVector (nested SparseFloatArray)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
sa := &schemapb.SparseFloatArray{}
|
||||
if err := unmarshalSparseFloatArray(v, sa); err != nil {
|
||||
return err
|
||||
}
|
||||
vf.Data = &schemapb.VectorField_SparseFloatVector{SparseFloatVector: sa}
|
||||
b = b[n:]
|
||||
case 7: // Int8Vector (bytes)
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
vf.Data = &schemapb.VectorField_Int8Vector{Int8Vector: cloneBytes(v)}
|
||||
b = b[n:]
|
||||
case 8: // VectorArray — delegate to official
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
va := &schemapb.VectorArray{}
|
||||
if err := protoUnmarshal(v, va); err != nil {
|
||||
return err
|
||||
}
|
||||
vf.Data = &schemapb.VectorField_VectorArray{VectorArray: va}
|
||||
b = b[n:]
|
||||
default: // unhandled (future) field → fold into official merge
|
||||
n2 := skipField(b, wtype)
|
||||
if n2 <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
rest = append(rest, start[:n+n2]...)
|
||||
b = b[n2:]
|
||||
}
|
||||
}
|
||||
if len(rest) > 0 {
|
||||
return protoMerge(rest, vf)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// unmarshalSparseFloatArray: repeated bytes contents = 1; int64 dim = 2;
|
||||
func unmarshalSparseFloatArray(b []byte, sa *schemapb.SparseFloatArray) error {
|
||||
full := b
|
||||
for len(b) > 0 {
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
// A known field with an unexpected wire type is not what proto3 produces;
|
||||
// hand the whole message to the official codec so the outcome matches
|
||||
// proto.Unmarshal exactly instead of blindly decoding the wrong shape.
|
||||
if (num == 1 && wtype != 2) || (num == 2 && wtype != 0) {
|
||||
return fallbackUnmarshal(full, sa)
|
||||
}
|
||||
switch num {
|
||||
case 1:
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
sa.Contents = append(sa.Contents, cloneBytes(v))
|
||||
b = b[n:]
|
||||
case 2:
|
||||
v, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
sa.Dim = int64(v)
|
||||
b = b[n:]
|
||||
default:
|
||||
return fallbackUnmarshal(full, sa) // unknown field → official, preserves it
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// unmarshalFloatArray decodes a wire-format schemapb.FloatArray into fa
|
||||
// (repeated float data = 1, packed fixed32).
|
||||
func unmarshalFloatArray(b []byte, fa *schemapb.FloatArray) error {
|
||||
return decodePackedF32(b, &fa.Data, fa)
|
||||
}
|
||||
|
||||
// --- wire primitives ---
|
||||
|
||||
func consumeVarint(b []byte) (uint64, int) {
|
||||
var x uint64
|
||||
var s uint
|
||||
for i := 0; i < len(b); i++ {
|
||||
c := b[i]
|
||||
if c < 0x80 {
|
||||
if i > 9 || (i == 9 && c > 1) {
|
||||
return 0, -1 // overflow
|
||||
}
|
||||
return x | uint64(c)<<s, i + 1
|
||||
}
|
||||
x |= uint64(c&0x7f) << s
|
||||
s += 7
|
||||
}
|
||||
return 0, 0 // truncated
|
||||
}
|
||||
|
||||
func consumeTag(b []byte) (num int, wtype int, n int) {
|
||||
v, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return 0, 0, n
|
||||
}
|
||||
return int(v >> 3), int(v & 7), n
|
||||
}
|
||||
|
||||
func consumeBytes(b []byte) ([]byte, int) {
|
||||
l, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return nil, n
|
||||
}
|
||||
if uint64(len(b)-n) < l {
|
||||
return nil, 0 // truncated
|
||||
}
|
||||
return b[n : n+int(l)], n + int(l)
|
||||
}
|
||||
|
||||
func skipField(b []byte, wtype int) int {
|
||||
switch wtype {
|
||||
case 0: // varint
|
||||
_, n := consumeVarint(b)
|
||||
return n
|
||||
case 1: // fixed64
|
||||
if len(b) < 8 {
|
||||
return 0
|
||||
}
|
||||
return 8
|
||||
case 2: // length-delimited
|
||||
l, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return n
|
||||
}
|
||||
if uint64(len(b)-n) < l {
|
||||
return 0
|
||||
}
|
||||
return n + int(l)
|
||||
case 5: // fixed32
|
||||
if len(b) < 4 {
|
||||
return 0
|
||||
}
|
||||
return 4
|
||||
default:
|
||||
return -1
|
||||
}
|
||||
}
|
||||
|
||||
// little-endian fixed readers (used by packed vector/scalar decode single paths)
|
||||
func le32(b []byte) uint32 { return binary.LittleEndian.Uint32(b) }
|
||||
func le64(b []byte) uint64 { return binary.LittleEndian.Uint64(b) }
|
||||
@@ -0,0 +1,67 @@
|
||||
// Package fastpb hand-writes unmarshal for the hot schemapb types
|
||||
// (FieldData / SearchResultData). Correctness is pinned by differential
|
||||
// testing against the official proto.Unmarshal: for any input, our decode
|
||||
// must produce a result that proto.Equal-s the official one.
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
)
|
||||
|
||||
// roundTrip marshals src with the official codec, decodes with our fast path,
|
||||
// and fails unless the result is proto.Equal to src.
|
||||
func roundTripFieldData(t *testing.T, src *schemapb.FieldData) {
|
||||
t.Helper()
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatalf("official marshal: %v", err)
|
||||
}
|
||||
var got schemapb.FieldData
|
||||
if err := UnmarshalFieldData(wire, &got); err != nil {
|
||||
t.Fatalf("UnmarshalFieldData: %v", err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch:\n src = %v\n got = %v", src, &got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnmarshalFieldData_ScalarHeaderFields(t *testing.T) {
|
||||
roundTripFieldData(t, &schemapb.FieldData{
|
||||
Type: schemapb.DataType_VarChar,
|
||||
FieldName: "pk",
|
||||
FieldId: 100,
|
||||
IsDynamic: true,
|
||||
})
|
||||
}
|
||||
|
||||
func TestUnmarshalFieldData_FloatVector(t *testing.T) {
|
||||
roundTripFieldData(t, &schemapb.FieldData{
|
||||
Type: schemapb.DataType_FloatVector,
|
||||
FieldName: "embedding",
|
||||
FieldId: 102,
|
||||
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
|
||||
Dim: 4,
|
||||
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{
|
||||
Data: []float32{1.5, -2.25, 0, 3.0e9, 0.001, float32(math.Copysign(0, -1)), 42, 7},
|
||||
}},
|
||||
}},
|
||||
})
|
||||
}
|
||||
|
||||
func TestUnmarshalFieldData_StringData(t *testing.T) {
|
||||
roundTripFieldData(t, &schemapb.FieldData{
|
||||
Type: schemapb.DataType_VarChar,
|
||||
FieldName: "title",
|
||||
FieldId: 101,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{
|
||||
Data: []string{"hello", "", "世界", "a longer varchar value to copy"},
|
||||
}},
|
||||
}},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
)
|
||||
|
||||
// withUnknownField appends an unknown field (number 15, varint 42) to a wire blob.
|
||||
func withUnknownField(b []byte) []byte {
|
||||
return append(append([]byte{}, b...), 0x78, 0x2A) // tag (15<<3|0)=0x78, value 42
|
||||
}
|
||||
|
||||
// TestLeafUnknownFieldPreserved guards the fix: leaf *Array messages carrying an
|
||||
// unknown/future field must decode identically to the official codec (the field
|
||||
// is preserved, not dropped). Previously these helpers silently skipped it.
|
||||
func TestLeafUnknownFieldPreserved(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
base proto.Message
|
||||
fresh func() proto.Message
|
||||
decode func(b []byte, m proto.Message) error
|
||||
}{
|
||||
{
|
||||
"StringArray", &schemapb.StringArray{Data: []string{"a", "b"}},
|
||||
func() proto.Message { return &schemapb.StringArray{} },
|
||||
func(b []byte, m proto.Message) error { return dec{}.stringArray(b, m.(*schemapb.StringArray)) },
|
||||
},
|
||||
{
|
||||
"LongArray", &schemapb.LongArray{Data: []int64{1, 2, 3}},
|
||||
func() proto.Message { return &schemapb.LongArray{} },
|
||||
func(b []byte, m proto.Message) error {
|
||||
a := m.(*schemapb.LongArray)
|
||||
return decodePackedI64(b, &a.Data, a)
|
||||
},
|
||||
},
|
||||
{
|
||||
"IntArray", &schemapb.IntArray{Data: []int32{1, 2, 3}},
|
||||
func() proto.Message { return &schemapb.IntArray{} },
|
||||
func(b []byte, m proto.Message) error {
|
||||
a := m.(*schemapb.IntArray)
|
||||
return decodePackedI32(b, &a.Data, a)
|
||||
},
|
||||
},
|
||||
{
|
||||
"BoolArray", &schemapb.BoolArray{Data: []bool{true, false}},
|
||||
func() proto.Message { return &schemapb.BoolArray{} },
|
||||
func(b []byte, m proto.Message) error {
|
||||
a := m.(*schemapb.BoolArray)
|
||||
return decodePackedBool(b, &a.Data, a)
|
||||
},
|
||||
},
|
||||
{
|
||||
"FloatArray", &schemapb.FloatArray{Data: []float32{1, 2, 3}},
|
||||
func() proto.Message { return &schemapb.FloatArray{} },
|
||||
func(b []byte, m proto.Message) error {
|
||||
a := m.(*schemapb.FloatArray)
|
||||
return decodePackedF32(b, &a.Data, a)
|
||||
},
|
||||
},
|
||||
{
|
||||
"DoubleArray", &schemapb.DoubleArray{Data: []float64{1, 2, 3}},
|
||||
func() proto.Message { return &schemapb.DoubleArray{} },
|
||||
func(b []byte, m proto.Message) error {
|
||||
a := m.(*schemapb.DoubleArray)
|
||||
return decodePackedF64(b, &a.Data, a)
|
||||
},
|
||||
},
|
||||
{
|
||||
"BytesArray", &schemapb.BytesArray{Data: [][]byte{{1, 2}, {3}}},
|
||||
func() proto.Message { return &schemapb.BytesArray{} },
|
||||
func(b []byte, m proto.Message) error {
|
||||
a := m.(*schemapb.BytesArray)
|
||||
return decodeRepeatedBytes(b, &a.Data, a)
|
||||
},
|
||||
},
|
||||
{
|
||||
"SparseFloatArray", &schemapb.SparseFloatArray{Dim: 9, Contents: [][]byte{{1}, {2}}},
|
||||
func() proto.Message { return &schemapb.SparseFloatArray{} },
|
||||
func(b []byte, m proto.Message) error {
|
||||
return unmarshalSparseFloatArray(b, m.(*schemapb.SparseFloatArray))
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
base, err := proto.Marshal(c.base)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wire := withUnknownField(base)
|
||||
|
||||
official := c.fresh()
|
||||
if err := proto.Unmarshal(wire, official); err != nil {
|
||||
t.Fatalf("official: %v", err)
|
||||
}
|
||||
got := c.fresh()
|
||||
if err := c.decode(wire, got); err != nil {
|
||||
t.Fatalf("fastpb: %v", err)
|
||||
}
|
||||
if !proto.Equal(official, got) {
|
||||
t.Fatalf("unknown field NOT preserved:\n official=%v\n got=%v", official, got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestUnmarshalResetsTarget guards the fix: the public entry points must clear a
|
||||
// pre-populated target (matching official proto.Unmarshal semantics) so stale
|
||||
// fields cannot leak when the codec reuses a message.
|
||||
func TestUnmarshalResetsTarget(t *testing.T) {
|
||||
src := &schemapb.SearchResultData{TopK: 5, NumQueries: 2}
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// target pre-populated with stale data that is NOT present in src
|
||||
got := &schemapb.SearchResultData{
|
||||
PrimaryFieldName: "stale",
|
||||
Scores: []float32{9, 9, 9},
|
||||
AllSearchCount: 123,
|
||||
}
|
||||
if err := UnmarshalSearchResultData(wire, got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !proto.Equal(src, got) {
|
||||
t.Fatalf("stale fields not cleared before decode:\n want=%v\n got=%v", src, got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
// appendGroup appends a proto2 group-encoded field (start-group ... end-group,
|
||||
// wire types 3/4) carrying one nested varint. proto3 never produces this; the
|
||||
// official codec preserves it as an unknown field.
|
||||
func appendGroup(b []byte, fieldNum int32) []byte {
|
||||
b = protowire.AppendTag(b, protowire.Number(fieldNum), protowire.StartGroupType)
|
||||
b = protowire.AppendTag(b, 1, protowire.VarintType)
|
||||
b = protowire.AppendVarint(b, 42)
|
||||
b = protowire.AppendTag(b, protowire.Number(fieldNum), protowire.EndGroupType)
|
||||
return b
|
||||
}
|
||||
|
||||
// TestProto2GroupFallback: a wire buffer containing a proto2 group (which fastpb
|
||||
// does not decode) must fall back to the official codec and yield a byte-for-byte
|
||||
// identical message, never an error (issue: fastpb proto2 group handling).
|
||||
func TestProto2GroupFallback(t *testing.T) {
|
||||
t.Run("InsertRequest/top-level group", func(t *testing.T) {
|
||||
canonical, err := proto.Marshal(&milvuspb.InsertRequest{
|
||||
CollectionName: "c", DbName: "db", NumRows: 3,
|
||||
FieldsData: []*schemapb.FieldData{{FieldId: 1, Type: schemapb.DataType_Int64, FieldName: "id"}},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
wire := appendGroup(append([]byte{}, canonical...), 1000)
|
||||
|
||||
want := &milvuspb.InsertRequest{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want), "official codec must accept the group")
|
||||
got := &milvuspb.InsertRequest{}
|
||||
require.NoError(t, UnmarshalInsertRequest(wire, got), "fastpb must fall back, not error")
|
||||
assert.True(t, proto.Equal(got, want), "fallback result must equal official decode")
|
||||
})
|
||||
|
||||
t.Run("InsertRequest/group nested in fields_data", func(t *testing.T) {
|
||||
// Build a FieldData submessage whose bytes carry a nested group, exercising
|
||||
// errProto2 propagation up from the nested fieldData decoder to the entry.
|
||||
fdBytes, err := proto.Marshal(&schemapb.FieldData{FieldId: 1, Type: schemapb.DataType_Int64, FieldName: "id"})
|
||||
require.NoError(t, err)
|
||||
fdBytes = appendGroup(fdBytes, 999)
|
||||
|
||||
var wire []byte
|
||||
wire = protowire.AppendTag(wire, 3, protowire.BytesType) // collection_name
|
||||
wire = protowire.AppendString(wire, "c")
|
||||
wire = protowire.AppendTag(wire, 5, protowire.BytesType) // fields_data (5)
|
||||
wire = protowire.AppendBytes(wire, fdBytes)
|
||||
|
||||
want := &milvuspb.InsertRequest{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &milvuspb.InsertRequest{}
|
||||
require.NoError(t, UnmarshalInsertRequest(wire, got))
|
||||
assert.True(t, proto.Equal(got, want))
|
||||
})
|
||||
|
||||
t.Run("SearchResultData/top-level group", func(t *testing.T) {
|
||||
canonical, err := proto.Marshal(&schemapb.SearchResultData{NumQueries: 2, TopK: 5})
|
||||
require.NoError(t, err)
|
||||
wire := appendGroup(append([]byte{}, canonical...), 1000)
|
||||
|
||||
want := &schemapb.SearchResultData{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.SearchResultData{}
|
||||
require.NoError(t, UnmarshalSearchResultData(wire, got))
|
||||
assert.True(t, proto.Equal(got, want))
|
||||
})
|
||||
|
||||
t.Run("RetrieveResults/top-level group", func(t *testing.T) {
|
||||
canonical, err := proto.Marshal(&internalpb.RetrieveResults{ReqID: 7, AllRetrieveCount: 9})
|
||||
require.NoError(t, err)
|
||||
wire := appendGroup(append([]byte{}, canonical...), 1000)
|
||||
|
||||
want := &internalpb.RetrieveResults{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &internalpb.RetrieveResults{}
|
||||
require.NoError(t, UnmarshalRetrieveResults(wire, got))
|
||||
assert.True(t, proto.Equal(got, want))
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,379 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"math"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
)
|
||||
|
||||
// cloneBytes returns a heap copy of v (v aliases the source buffer).
|
||||
func cloneBytes(v []byte) []byte {
|
||||
out := make([]byte, len(v))
|
||||
copy(out, v)
|
||||
return out
|
||||
}
|
||||
|
||||
// protoUnmarshal delegates a cold/complex sub-message to the official codec
|
||||
// (into a fresh message).
|
||||
func protoUnmarshal(b []byte, m proto.Message) error { return proto.Unmarshal(b, m) }
|
||||
|
||||
// protoMerge decodes b into m WITHOUT resetting m first — used to fold the raw
|
||||
// bytes of unhandled (cold/future) fields back into a message whose hot fields
|
||||
// were already populated by the fast path. proto wire merge semantics make this
|
||||
// byte-for-byte equivalent to a single official proto.Unmarshal of the whole
|
||||
// message, as long as no oneof member was deferred here (see callers).
|
||||
func protoMerge(b []byte, m proto.Message) error {
|
||||
return proto.UnmarshalOptions{Merge: true}.Unmarshal(b, m)
|
||||
}
|
||||
|
||||
// fallbackUnmarshal decodes the whole submessage with the official codec. Leaf
|
||||
// *Array decoders call this the moment they see any field other than the single
|
||||
// `data` field, so an unknown/future field on a leaf message is preserved exactly
|
||||
// (rather than dropped). The fast path is unaffected for the common all-data case.
|
||||
func fallbackUnmarshal(b []byte, m proto.Message) error {
|
||||
proto.Reset(m)
|
||||
return proto.Unmarshal(b, m)
|
||||
}
|
||||
|
||||
// decodeScalarFallback decodes the rare ScalarField oneof variants via the
|
||||
// official codec. Returns false if num is not a known cold variant, so the
|
||||
// caller can route the raw field bytes to protoMerge instead of dropping it.
|
||||
func decodeScalarFallback(num int, v []byte, sf *schemapb.ScalarField) (bool, error) {
|
||||
switch num {
|
||||
case 8:
|
||||
m := &schemapb.ArrayArray{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return true, err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_ArrayData{ArrayData: m}
|
||||
case 10:
|
||||
m := &schemapb.GeometryArray{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return true, err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_GeometryData{GeometryData: m}
|
||||
case 11:
|
||||
m := &schemapb.TimestamptzArray{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return true, err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_TimestamptzData{TimestamptzData: m}
|
||||
case 12:
|
||||
m := &schemapb.GeometryWktArray{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return true, err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_GeometryWktData{GeometryWktData: m}
|
||||
case 13:
|
||||
m := &schemapb.MolArray{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return true, err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_MolData{MolData: m}
|
||||
case 14:
|
||||
m := &schemapb.MolSmilesArray{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return true, err
|
||||
}
|
||||
sf.Data = &schemapb.ScalarField_MolSmilesData{MolSmilesData: m}
|
||||
default:
|
||||
return false, nil
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// --- array-message decoders: b is the *Array submessage; field 1 holds the data ---
|
||||
|
||||
func decodePackedBool(b []byte, dst *[]bool, m proto.Message) error {
|
||||
full := b
|
||||
for len(b) > 0 {
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if num != 1 {
|
||||
return fallbackUnmarshal(full, m) // unknown field → official codec, preserves it
|
||||
}
|
||||
switch wtype {
|
||||
case 2: // packed
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
for len(v) > 0 {
|
||||
x, m := consumeVarint(v)
|
||||
if m <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, x != 0)
|
||||
v = v[m:]
|
||||
}
|
||||
b = b[n:]
|
||||
case 0: // single varint
|
||||
x, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, x != 0)
|
||||
b = b[n:]
|
||||
default:
|
||||
return fallbackUnmarshal(full, m) // field 1 with unexpected wire type → official
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodePackedI32(b []byte, dst *[]int32, m proto.Message) error {
|
||||
full := b
|
||||
for len(b) > 0 {
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if num != 1 {
|
||||
return fallbackUnmarshal(full, m) // unknown field → official codec, preserves it
|
||||
}
|
||||
switch wtype {
|
||||
case 2: // packed
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
for len(v) > 0 {
|
||||
x, m := consumeVarint(v)
|
||||
if m <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, int32(x))
|
||||
v = v[m:]
|
||||
}
|
||||
b = b[n:]
|
||||
case 0: // single varint
|
||||
x, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, int32(x))
|
||||
b = b[n:]
|
||||
default:
|
||||
return fallbackUnmarshal(full, m) // field 1 with unexpected wire type → official
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodePackedI64(b []byte, dst *[]int64, m proto.Message) error {
|
||||
full := b
|
||||
for len(b) > 0 {
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if num != 1 {
|
||||
return fallbackUnmarshal(full, m) // unknown field → official codec, preserves it
|
||||
}
|
||||
switch wtype {
|
||||
case 2: // packed
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
for len(v) > 0 {
|
||||
x, m := consumeVarint(v)
|
||||
if m <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, int64(x))
|
||||
v = v[m:]
|
||||
}
|
||||
b = b[n:]
|
||||
case 0: // single varint
|
||||
x, n := consumeVarint(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, int64(x))
|
||||
b = b[n:]
|
||||
default:
|
||||
return fallbackUnmarshal(full, m) // field 1 with unexpected wire type → official
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodePackedF32(b []byte, dst *[]float32, m proto.Message) error {
|
||||
full := b
|
||||
for len(b) > 0 {
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if num != 1 {
|
||||
return fallbackUnmarshal(full, m) // unknown field → official codec, preserves it
|
||||
}
|
||||
switch wtype {
|
||||
case 2:
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
if len(v)%4 != 0 {
|
||||
return errMalformed
|
||||
}
|
||||
chunk := bytesToF32(v) // single memcpy
|
||||
if len(*dst) == 0 {
|
||||
*dst = chunk
|
||||
} else {
|
||||
*dst = append(*dst, chunk...)
|
||||
}
|
||||
b = b[n:]
|
||||
case 5:
|
||||
if len(b) < 4 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, math.Float32frombits(le32(b)))
|
||||
b = b[4:]
|
||||
default:
|
||||
return fallbackUnmarshal(full, m) // field 1 with unexpected wire type → official
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodePackedF64(b []byte, dst *[]float64, m proto.Message) error {
|
||||
full := b
|
||||
for len(b) > 0 {
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if num != 1 {
|
||||
return fallbackUnmarshal(full, m) // unknown field → official codec, preserves it
|
||||
}
|
||||
switch wtype {
|
||||
case 2:
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
if len(v)%8 != 0 {
|
||||
return errMalformed
|
||||
}
|
||||
chunk := bytesToF64(v) // single memcpy
|
||||
if len(*dst) == 0 {
|
||||
*dst = chunk
|
||||
} else {
|
||||
*dst = append(*dst, chunk...)
|
||||
}
|
||||
b = b[n:]
|
||||
case 1:
|
||||
if len(b) < 8 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, math.Float64frombits(le64(b)))
|
||||
b = b[8:]
|
||||
default:
|
||||
return fallbackUnmarshal(full, m) // field 1 with unexpected wire type → official
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// appendPackedF32 appends a raw packed-fixed32 payload (the field value itself,
|
||||
// not wrapped in an *Array submessage) to dst.
|
||||
func appendPackedF32(v []byte, dst *[]float32) error {
|
||||
if len(v)%4 != 0 {
|
||||
return errMalformed
|
||||
}
|
||||
chunk := bytesToF32(v) // single memcpy
|
||||
if len(*dst) == 0 {
|
||||
*dst = chunk
|
||||
} else {
|
||||
*dst = append(*dst, chunk...)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// appendPackedBool appends a raw packed-varint bool payload to dst.
|
||||
func appendPackedBool(v []byte, dst *[]bool) error {
|
||||
for len(v) > 0 {
|
||||
x, m := consumeVarint(v)
|
||||
if m <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, x != 0)
|
||||
v = v[m:]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// appendPackedU32 appends a raw packed-varint payload to dst ([]uint32).
|
||||
func appendPackedU32(v []byte, dst *[]uint32) error {
|
||||
for len(v) > 0 {
|
||||
x, m := consumeVarint(v)
|
||||
if m <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, uint32(x))
|
||||
v = v[m:]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// appendPackedI64 appends a raw packed-varint payload to dst.
|
||||
func appendPackedI64(v []byte, dst *[]int64) error {
|
||||
for len(v) > 0 {
|
||||
x, m := consumeVarint(v)
|
||||
if m <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, int64(x))
|
||||
v = v[m:]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// decodeRepeatedBytes: field 1 is repeated bytes (BytesArray / JSONArray).
|
||||
func decodeRepeatedBytes(b []byte, dst *[][]byte, m proto.Message) error {
|
||||
full := b
|
||||
for len(b) > 0 {
|
||||
num, wtype, n := consumeTag(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[n:]
|
||||
if num == 1 && wtype == 2 {
|
||||
v, n := consumeBytes(b)
|
||||
if n <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
*dst = append(*dst, cloneBytes(v))
|
||||
b = b[n:]
|
||||
} else {
|
||||
return fallbackUnmarshal(full, m) // field 1 with unexpected wire type → official
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
commonpb "github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
// concat marshals each message and concatenates the bytes. proto3 defines the
|
||||
// concatenation of two serialized messages to be equivalent to merging them, so
|
||||
// this produces a wire buffer where a singular embedded message field appears
|
||||
// twice — the case fastpb must merge (not replace) to match proto.Unmarshal.
|
||||
func concat(t *testing.T, msgs ...proto.Message) []byte {
|
||||
var out []byte
|
||||
for _, m := range msgs {
|
||||
b, err := proto.Marshal(m)
|
||||
require.NoError(t, err)
|
||||
out = append(out, b...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestSingularMessageMerge verifies that a singular non-oneof embedded message
|
||||
// appearing more than once on the wire is merged field-by-field (proto3 semantics),
|
||||
// matching proto.Unmarshal exactly, rather than last-wins replaced.
|
||||
func TestSingularMessageMerge(t *testing.T) {
|
||||
t.Run("InsertRequest.Base", func(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&milvuspb.InsertRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Insert}, CollectionName: "c"},
|
||||
&milvuspb.InsertRequest{Base: &commonpb.MsgBase{Timestamp: 123}, NumRows: 5},
|
||||
)
|
||||
want := &milvuspb.InsertRequest{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &milvuspb.InsertRequest{}
|
||||
require.NoError(t, UnmarshalInsertRequest(wire, got))
|
||||
// sanity: the two Base occurrences must have merged, not last-wins replaced.
|
||||
require.Equal(t, commonpb.MsgType_Insert, want.Base.MsgType)
|
||||
require.Equal(t, uint64(123), want.Base.Timestamp)
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
|
||||
t.Run("RetrieveResults.Base/Status/Ids/CostAggregation", func(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&internalpb.RetrieveResults{
|
||||
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Retrieve},
|
||||
Status: &commonpb.Status{Code: 1},
|
||||
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
|
||||
CostAggregation: &internalpb.CostAggregation{ResponseTime: 10},
|
||||
},
|
||||
&internalpb.RetrieveResults{
|
||||
Base: &commonpb.MsgBase{Timestamp: 7},
|
||||
Status: &commonpb.Status{Reason: "x"},
|
||||
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{3}}}},
|
||||
CostAggregation: &internalpb.CostAggregation{TotalNQ: 2},
|
||||
},
|
||||
)
|
||||
want := &internalpb.RetrieveResults{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &internalpb.RetrieveResults{}
|
||||
require.NoError(t, UnmarshalRetrieveResults(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
|
||||
t.Run("SearchResultData.Ids/GroupByFieldValue", func(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&schemapb.SearchResultData{
|
||||
NumQueries: 1,
|
||||
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
|
||||
GroupByFieldValue: &schemapb.FieldData{FieldId: 10, Type: schemapb.DataType_Int64},
|
||||
},
|
||||
&schemapb.SearchResultData{
|
||||
TopK: 5,
|
||||
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2}}}},
|
||||
GroupByFieldValue: &schemapb.FieldData{FieldName: "g"},
|
||||
},
|
||||
)
|
||||
want := &schemapb.SearchResultData{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.SearchResultData{}
|
||||
require.NoError(t, UnmarshalSearchResultData(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
}
|
||||
|
||||
func TestRepeatedMessageOneofMerge(t *testing.T) {
|
||||
t.Run("FieldData.Scalars", func(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1}}},
|
||||
}}},
|
||||
&schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{2}}},
|
||||
}}},
|
||||
)
|
||||
want := &schemapb.FieldData{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.FieldData{}
|
||||
require.NoError(t, UnmarshalFieldData(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
|
||||
t.Run("ScalarField.LongData", func(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1}}}},
|
||||
&schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{2}}}},
|
||||
)
|
||||
want := &schemapb.ScalarField{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.ScalarField{}
|
||||
require.NoError(t, dec{}.scalarField(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
|
||||
t.Run("VectorField.FloatVector", func(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&schemapb.VectorField{Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1}}}},
|
||||
&schemapb.VectorField{Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{2}}}},
|
||||
)
|
||||
want := &schemapb.VectorField{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.VectorField{}
|
||||
require.NoError(t, unmarshalVectorField(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
|
||||
t.Run("IDs.IntId", func(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
|
||||
&schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2}}}},
|
||||
)
|
||||
want := &schemapb.IDs{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.IDs{}
|
||||
require.NoError(t, dec{}.ids(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
|
||||
t.Run("different variants remain last-wins", func(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{}}},
|
||||
&schemapb.FieldData{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: 8}}},
|
||||
)
|
||||
want := &schemapb.FieldData{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.FieldData{}
|
||||
require.NoError(t, UnmarshalFieldData(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
}
|
||||
|
||||
func TestElementIndicesMerge(t *testing.T) {
|
||||
wire := concat(t,
|
||||
&schemapb.SearchResultData{ElementIndices: &schemapb.LongArray{Data: []int64{1}}},
|
||||
&schemapb.SearchResultData{ElementIndices: &schemapb.LongArray{Data: []int64{2}}},
|
||||
)
|
||||
want := &schemapb.SearchResultData{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.SearchResultData{}
|
||||
require.NoError(t, UnmarshalSearchResultData(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
}
|
||||
|
||||
func TestMixedPackedUnpackedFixed32Order(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
field protowire.Number
|
||||
}{
|
||||
{name: "scores", field: 4},
|
||||
{name: "distances", field: 10},
|
||||
{name: "recalls", field: 12},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
packedFirst := protowire.AppendFixed32(nil, math.Float32bits(1))
|
||||
packedLast := protowire.AppendFixed32(nil, math.Float32bits(3))
|
||||
var wire []byte
|
||||
wire = protowire.AppendTag(wire, tc.field, protowire.BytesType)
|
||||
wire = protowire.AppendBytes(wire, packedFirst)
|
||||
wire = protowire.AppendTag(wire, tc.field, protowire.Fixed32Type)
|
||||
wire = protowire.AppendFixed32(wire, math.Float32bits(2))
|
||||
wire = protowire.AppendTag(wire, tc.field, protowire.BytesType)
|
||||
wire = protowire.AppendBytes(wire, packedLast)
|
||||
|
||||
want := &schemapb.SearchResultData{}
|
||||
require.NoError(t, proto.Unmarshal(wire, want))
|
||||
got := &schemapb.SearchResultData{}
|
||||
require.NoError(t, UnmarshalSearchResultData(wire, got))
|
||||
assert.True(t, proto.Equal(got, want), "got=%v want=%v", got, want)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
// fieldNumbers returns the sorted set of field numbers declared on m's
|
||||
// descriptor, INCLUDING oneof members and proto3-optional fields (the
|
||||
// descriptor enumerates them, unlike struct-tag scraping).
|
||||
func fieldNumbers(m proto.Message) []int {
|
||||
fds := m.ProtoReflect().Descriptor().Fields()
|
||||
out := make([]int, 0, fds.Len())
|
||||
for i := 0; i < fds.Len(); i++ {
|
||||
out = append(out, int(fds.Get(i).Number()))
|
||||
}
|
||||
sort.Ints(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// TestProtoContract_FieldSetsPinned is a tripwire. fastpb hand-writes wire
|
||||
// decoders for the message types below; each decoder switches on hard-coded
|
||||
// field numbers. The golden field set for every such type is pinned here, so
|
||||
// ANY edit to these protos (a new field, a removed field, a renumber) flips
|
||||
// this test red and blocks CI.
|
||||
//
|
||||
// When it fails, do NOT just bump the golden set. First go read the matching
|
||||
// hand-decoder in pkg/util/fastpb and decide what the proto change requires:
|
||||
//
|
||||
// - A new plain (non-oneof) scalar/message field is correct-by-construction:
|
||||
// unknown field numbers fold into the official codec via the decoder's
|
||||
// `default -> rest -> protoMerge` fallback. You may still want to hand-decode
|
||||
// it if it lands on a hot path, but correctness is already preserved.
|
||||
//
|
||||
// - A NEW oneof VARIANT is the dangerous case. It also falls through to the
|
||||
// deferred protoMerge, but that breaks oneof last-wins ordering relative to
|
||||
// the variants decoded in-pass (see the StructArrays/case-8 comment in
|
||||
// fielddata.go). New oneof variants on FieldData / ScalarField / VectorField
|
||||
// / IDs MUST be handled in-pass, not left to the fallback.
|
||||
//
|
||||
// - Changing the wire type or meaning of an EXISTING hand-decoded number
|
||||
// silently corrupts the fast path — the hard-coded case still fires.
|
||||
//
|
||||
// Only after the decoder is correct should you update the golden set below.
|
||||
func TestProtoContract_FieldSetsPinned(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
msg proto.Message
|
||||
want []int
|
||||
}{
|
||||
// top-level fast-pathed types (TryUnmarshal dispatch)
|
||||
{"internalpb.RetrieveResults", &internalpb.RetrieveResults{}, []int{1, 2, 3, 4, 5, 6, 7, 8, 13, 14, 15, 16, 17, 18, 19}},
|
||||
{"milvuspb.InsertRequest", &milvuspb.InsertRequest{}, []int{1, 2, 3, 4, 5, 6, 7, 8, 9}},
|
||||
{"milvuspb.UpsertRequest", &milvuspb.UpsertRequest{}, []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}},
|
||||
{"schemapb.SearchResultData", &schemapb.SearchResultData{}, []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 17, 18, 19}},
|
||||
// nested hand-decoded types (oneof-bearing — highest divergence risk)
|
||||
{"schemapb.FieldData", &schemapb.FieldData{}, []int{1, 2, 3, 4, 5, 6, 7, 8}},
|
||||
{"schemapb.ScalarField", &schemapb.ScalarField{}, []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14}},
|
||||
{"schemapb.VectorField", &schemapb.VectorField{}, []int{1, 2, 3, 4, 5, 6, 7, 8}},
|
||||
{"schemapb.IDs", &schemapb.IDs{}, []int{1, 2}},
|
||||
// leaf arrays with their own hand-written field switches (no unknown-field
|
||||
// fallback in some — a new field here would mis-parse)
|
||||
{"schemapb.SparseFloatArray", &schemapb.SparseFloatArray{}, []int{1, 2}},
|
||||
{"schemapb.FloatArray", &schemapb.FloatArray{}, []int{1}},
|
||||
{"schemapb.LongArray", &schemapb.LongArray{}, []int{1}},
|
||||
{"schemapb.IntArray", &schemapb.IntArray{}, []int{1}},
|
||||
{"schemapb.BoolArray", &schemapb.BoolArray{}, []int{1}},
|
||||
{"schemapb.DoubleArray", &schemapb.DoubleArray{}, []int{1}},
|
||||
{"schemapb.BytesArray", &schemapb.BytesArray{}, []int{1}},
|
||||
{"schemapb.JSONArray", &schemapb.JSONArray{}, []int{1}},
|
||||
{"schemapb.StringArray", &schemapb.StringArray{}, []int{1}},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := fieldNumbers(c.msg)
|
||||
assert.Equalf(t, c.want, got,
|
||||
"%s proto field set changed (want %v, got %v).\n"+
|
||||
"This type is hand-decoded in pkg/util/fastpb. Read the decoder and the\n"+
|
||||
"guidance on TestProtoContract_FieldSetsPinned BEFORE updating this golden set —\n"+
|
||||
"a new oneof variant in particular must be handled in-pass, not by the protoMerge fallback.",
|
||||
c.name, c.want, got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,416 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
commonpb "github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
// TryUnmarshal fast-paths the hot read/write message types through the
|
||||
// hand-written decoders and reports whether it handled v. Any other type
|
||||
// returns (false, nil) so the caller falls back to the official proto codec.
|
||||
// Only types pinned by the differential + fuzz tests are fast-pathed.
|
||||
func TryUnmarshal(v any, b []byte) (bool, error) {
|
||||
switch m := v.(type) {
|
||||
case *internalpb.RetrieveResults:
|
||||
return true, UnmarshalRetrieveResults(b, m)
|
||||
case *milvuspb.InsertRequest:
|
||||
return true, UnmarshalInsertRequest(b, m)
|
||||
case *milvuspb.UpsertRequest:
|
||||
return true, UnmarshalUpsertRequest(b, m)
|
||||
default:
|
||||
// Everything else (incl. SearchResultData, which is never a top-level gRPC
|
||||
// message — it travels as the SlicedBlob bytes of internalpb.SearchResults and
|
||||
// is fast-decoded directly at the reduce sites via UnmarshalSearchResultData)
|
||||
// falls back to the official codec.
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
|
||||
// UnmarshalRetrieveResults decodes internalpb.RetrieveResults (query read path,
|
||||
// querynode→proxy; internal/trusted). Hot field: fields_data (5).
|
||||
func UnmarshalRetrieveResults(b []byte, rr *internalpb.RetrieveResults) error {
|
||||
proto.Reset(rr) // match official proto.Unmarshal: clear target before decode
|
||||
if err := (dec{}).retrieveResults(b, rr); err != nil {
|
||||
if fallbackOnProto2(err) {
|
||||
proto.Reset(rr) // discard partial decode before the authoritative pass
|
||||
return proto.Unmarshal(b, rr)
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d dec) retrieveResults(b []byte, rr *internalpb.RetrieveResults) error {
|
||||
var rest []byte
|
||||
for len(b) > 0 {
|
||||
start := b
|
||||
num, wtype, tn := consumeTag(b)
|
||||
if tn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[tn:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
// varint scalar fields
|
||||
if wtype == 0 {
|
||||
v, vn := consumeVarint(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
switch num {
|
||||
case 3:
|
||||
rr.ReqID = int64(v)
|
||||
case 14:
|
||||
rr.AllRetrieveCount = int64(v)
|
||||
case 15:
|
||||
rr.HasMoreResult = v != 0
|
||||
case 16:
|
||||
rr.ScannedRemoteBytes = int64(v)
|
||||
case 17:
|
||||
rr.ScannedTotalBytes = int64(v)
|
||||
case 18:
|
||||
rr.ElementLevel = v != 0
|
||||
// fields 6/8 (packed int64) may also appear as a single varint:
|
||||
case 6:
|
||||
rr.SealedSegmentIDsRetrieved = append(rr.SealedSegmentIDsRetrieved, int64(v))
|
||||
case 8:
|
||||
rr.GlobalSealedSegmentIDs = append(rr.GlobalSealedSegmentIDs, int64(v))
|
||||
default:
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if wtype != 2 {
|
||||
sn := skipField(b, wtype)
|
||||
if sn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
rest = append(rest, start[:tn+sn]...)
|
||||
b = b[sn:]
|
||||
continue
|
||||
}
|
||||
v, vn := consumeBytes(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
switch num {
|
||||
case 1: // Base (delegate)
|
||||
m := &commonpb.MsgBase{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
if rr.Base == nil {
|
||||
rr.Base = m
|
||||
} else {
|
||||
// wire-level merge: a later occurrence's explicit proto3-default
|
||||
// scalar (e.g. TargetID=0) must overwrite, matching proto.Unmarshal.
|
||||
if err := protoMerge(v, rr.Base); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 2: // Status (delegate)
|
||||
m := &commonpb.Status{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
if rr.Status == nil {
|
||||
rr.Status = m
|
||||
} else {
|
||||
if err := protoMerge(v, rr.Status); err != nil { // wire-level merge (see Base)
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 4: // Ids
|
||||
ids := &schemapb.IDs{}
|
||||
if err := d.ids(v, ids); err != nil {
|
||||
return err
|
||||
}
|
||||
if rr.Ids == nil {
|
||||
rr.Ids = ids
|
||||
} else {
|
||||
if err := protoMerge(v, rr.Ids); err != nil { // wire-level merge (see Base)
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 5: // fields_data (HOT)
|
||||
fd := &schemapb.FieldData{}
|
||||
if err := d.fieldData(v, fd); err != nil {
|
||||
return err
|
||||
}
|
||||
rr.FieldsData = append(rr.FieldsData, fd)
|
||||
case 6: // sealed_segmentIDs_retrieved (packed int64)
|
||||
if err := appendPackedI64(v, &rr.SealedSegmentIDsRetrieved); err != nil {
|
||||
return err
|
||||
}
|
||||
case 7: // channelIDs_retrieved (repeated string)
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rr.ChannelIDsRetrieved = append(rr.ChannelIDsRetrieved, s)
|
||||
case 8: // global_sealed_segmentIDs (packed int64)
|
||||
if err := appendPackedI64(v, &rr.GlobalSealedSegmentIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
case 13: // CostAggregation (delegate)
|
||||
m := &internalpb.CostAggregation{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
if rr.CostAggregation == nil {
|
||||
rr.CostAggregation = m
|
||||
} else {
|
||||
if err := protoMerge(v, rr.CostAggregation); err != nil { // wire-level merge (see Base)
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 19: // element_indices (repeated message, delegate)
|
||||
m := &internalpb.ElementIndices{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
rr.ElementIndices = append(rr.ElementIndices, m)
|
||||
default: // unhandled (future) field → fold into official merge
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
}
|
||||
if len(rest) > 0 {
|
||||
return protoMerge(rest, rr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UnmarshalInsertRequest decodes milvuspb.InsertRequest (write path, client→proxy;
|
||||
// UNTRUSTED ingress → strings are UTF-8 validated). Hot field: fields_data (5).
|
||||
func UnmarshalInsertRequest(b []byte, ir *milvuspb.InsertRequest) error {
|
||||
proto.Reset(ir) // match official proto.Unmarshal: clear target before decode
|
||||
if err := (dec{utf8: true}).insertRequest(b, ir); err != nil {
|
||||
if fallbackOnProto2(err) {
|
||||
proto.Reset(ir) // discard partial decode before the authoritative pass
|
||||
return proto.Unmarshal(b, ir)
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d dec) insertRequest(b []byte, ir *milvuspb.InsertRequest) error {
|
||||
var rest []byte
|
||||
for len(b) > 0 {
|
||||
start := b
|
||||
num, wtype, tn := consumeTag(b)
|
||||
if tn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[tn:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if wtype == 0 {
|
||||
v, vn := consumeVarint(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
switch num {
|
||||
case 6:
|
||||
ir.HashKeys = append(ir.HashKeys, uint32(v))
|
||||
case 7:
|
||||
ir.NumRows = uint32(v)
|
||||
case 8:
|
||||
ir.SchemaTimestamp = v
|
||||
default:
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if wtype != 2 {
|
||||
sn := skipField(b, wtype)
|
||||
if sn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
rest = append(rest, start[:tn+sn]...)
|
||||
b = b[sn:]
|
||||
continue
|
||||
}
|
||||
v, vn := consumeBytes(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
switch num {
|
||||
case 1: // Base (delegate)
|
||||
m := &commonpb.MsgBase{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
if ir.Base == nil {
|
||||
ir.Base = m
|
||||
} else {
|
||||
if err := protoMerge(v, ir.Base); err != nil { // wire-level merge (see rr.Base)
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 2: // DbName
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ir.DbName = s
|
||||
case 3: // CollectionName
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ir.CollectionName = s
|
||||
case 4: // PartitionName
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ir.PartitionName = s
|
||||
case 5: // fields_data (HOT)
|
||||
fd := &schemapb.FieldData{}
|
||||
if err := d.fieldData(v, fd); err != nil {
|
||||
return err
|
||||
}
|
||||
ir.FieldsData = append(ir.FieldsData, fd)
|
||||
case 6: // hash_keys (packed uint32)
|
||||
if err := appendPackedU32(v, &ir.HashKeys); err != nil {
|
||||
return err
|
||||
}
|
||||
case 9: // namespace (proto3 optional string)
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ns := s
|
||||
ir.Namespace = &ns
|
||||
default: // unhandled (future) field → fold into official merge
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
}
|
||||
if len(rest) > 0 {
|
||||
return protoMerge(rest, ir)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UnmarshalUpsertRequest decodes milvuspb.UpsertRequest (write path, client→proxy;
|
||||
// UNTRUSTED ingress → strings are UTF-8 validated). Fields 1-8 share InsertRequest's
|
||||
// exact wire layout (hot field: fields_data (5)); the upsert-only fields
|
||||
// partial_update (9), namespace (10) and field_ops (11) are folded into the official
|
||||
// merge, so they decode via the standard codec and stay wire-equivalent.
|
||||
func UnmarshalUpsertRequest(b []byte, ur *milvuspb.UpsertRequest) error {
|
||||
proto.Reset(ur) // match official proto.Unmarshal: clear target before decode
|
||||
if err := (dec{utf8: true}).upsertRequest(b, ur); err != nil {
|
||||
if fallbackOnProto2(err) {
|
||||
proto.Reset(ur) // discard partial decode before the authoritative pass
|
||||
return proto.Unmarshal(b, ur)
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d dec) upsertRequest(b []byte, ur *milvuspb.UpsertRequest) error {
|
||||
var rest []byte
|
||||
for len(b) > 0 {
|
||||
start := b
|
||||
num, wtype, tn := consumeTag(b)
|
||||
if tn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[tn:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if wtype == 0 {
|
||||
v, vn := consumeVarint(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
switch num {
|
||||
case 6:
|
||||
ur.HashKeys = append(ur.HashKeys, uint32(v))
|
||||
case 7:
|
||||
ur.NumRows = uint32(v)
|
||||
case 8:
|
||||
ur.SchemaTimestamp = v
|
||||
default: // partial_update (9) and any future varint field → official merge
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if wtype != 2 {
|
||||
sn := skipField(b, wtype)
|
||||
if sn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
rest = append(rest, start[:tn+sn]...)
|
||||
b = b[sn:]
|
||||
continue
|
||||
}
|
||||
v, vn := consumeBytes(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
switch num {
|
||||
case 1: // Base (delegate)
|
||||
m := &commonpb.MsgBase{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
if ur.Base == nil {
|
||||
ur.Base = m
|
||||
} else {
|
||||
if err := protoMerge(v, ur.Base); err != nil { // wire-level merge (see rr.Base)
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 2: // DbName
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ur.DbName = s
|
||||
case 3: // CollectionName
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ur.CollectionName = s
|
||||
case 4: // PartitionName
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ur.PartitionName = s
|
||||
case 5: // fields_data (HOT)
|
||||
fd := &schemapb.FieldData{}
|
||||
if err := d.fieldData(v, fd); err != nil {
|
||||
return err
|
||||
}
|
||||
ur.FieldsData = append(ur.FieldsData, fd)
|
||||
case 6: // hash_keys (packed uint32)
|
||||
if err := appendPackedU32(v, &ur.HashKeys); err != nil {
|
||||
return err
|
||||
}
|
||||
default: // namespace (10), field_ops (11), any future field → official merge
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
}
|
||||
if len(rest) > 0 {
|
||||
return protoMerge(rest, ur)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
func varcharFieldData(rows, strLen int) []*schemapb.FieldData {
|
||||
titles := make([]string, rows)
|
||||
for i := range titles {
|
||||
titles[i] = fmt.Sprintf("title_%0*d", strLen, i)
|
||||
}
|
||||
return []*schemapb.FieldData{{
|
||||
Type: schemapb.DataType_VarChar, FieldName: "title", FieldId: 101,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: titles}},
|
||||
}},
|
||||
}}
|
||||
}
|
||||
|
||||
func vectorFieldData(rows, dim int) []*schemapb.FieldData {
|
||||
vec := make([]float32, rows*dim)
|
||||
for i := range vec {
|
||||
vec[i] = float32(i) * 0.001
|
||||
}
|
||||
return []*schemapb.FieldData{{
|
||||
Type: schemapb.DataType_FloatVector, FieldName: "embedding", FieldId: 102,
|
||||
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
|
||||
Dim: int64(dim), Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: vec}},
|
||||
}},
|
||||
}}
|
||||
}
|
||||
|
||||
// --- Query path: RetrieveResults ---
|
||||
|
||||
func benchRetrieve(b *testing.B, rr *internalpb.RetrieveResults) {
|
||||
wire, err := proto.Marshal(rr)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Run("official", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out internalpb.RetrieveResults
|
||||
if err := proto.Unmarshal(wire, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
b.Run("fastpb", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out internalpb.RetrieveResults
|
||||
if err := UnmarshalRetrieveResults(wire, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkRetrieve_Varchar(b *testing.B) {
|
||||
benchRetrieve(b, &internalpb.RetrieveResults{FieldsData: varcharFieldData(1000, 12)})
|
||||
}
|
||||
|
||||
func BenchmarkRetrieve_Vector(b *testing.B) {
|
||||
benchRetrieve(b, &internalpb.RetrieveResults{FieldsData: vectorFieldData(1000, 768)})
|
||||
}
|
||||
|
||||
// --- Insert path: InsertRequest (UTF-8 validated) ---
|
||||
|
||||
func benchInsert(b *testing.B, ir *milvuspb.InsertRequest) {
|
||||
wire, err := proto.Marshal(ir)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Run("official", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out milvuspb.InsertRequest
|
||||
if err := proto.Unmarshal(wire, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
b.Run("fastpb", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out milvuspb.InsertRequest
|
||||
if err := UnmarshalInsertRequest(wire, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkInsert_Varchar(b *testing.B) {
|
||||
ir := &milvuspb.InsertRequest{CollectionName: "c", FieldsData: varcharFieldData(1000, 12)}
|
||||
benchInsert(b, ir)
|
||||
}
|
||||
|
||||
func BenchmarkInsert_Vector(b *testing.B) {
|
||||
ir := &milvuspb.InsertRequest{CollectionName: "c", FieldsData: vectorFieldData(1000, 768)}
|
||||
benchInsert(b, ir)
|
||||
}
|
||||
|
||||
// --- Upsert path: UpsertRequest (UTF-8 validated; fields 1-8 fast, 9/10/11 fold) ---
|
||||
|
||||
func benchUpsert(b *testing.B, ur *milvuspb.UpsertRequest) {
|
||||
wire, err := proto.Marshal(ur)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Run("official", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out milvuspb.UpsertRequest
|
||||
if err := proto.Unmarshal(wire, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
b.Run("fastpb", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out milvuspb.UpsertRequest
|
||||
if err := UnmarshalUpsertRequest(wire, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkUpsert_Varchar(b *testing.B) {
|
||||
benchUpsert(b, &milvuspb.UpsertRequest{CollectionName: "c", FieldsData: varcharFieldData(1000, 12)})
|
||||
}
|
||||
|
||||
func BenchmarkUpsert_Vector(b *testing.B) {
|
||||
benchUpsert(b, &milvuspb.UpsertRequest{CollectionName: "c", FieldsData: vectorFieldData(1000, 768)})
|
||||
}
|
||||
|
||||
// BenchmarkUpsert_ColdFields measures the fold-to-official path: the upsert-only
|
||||
// fields partial_update (9) / namespace (10) / field_ops (11) are not
|
||||
// hand-decoded — they accumulate into `rest` and merge via the official codec on
|
||||
// top of the fast-decoded fields 1-8. This is also the shape unknown/future
|
||||
// fields take.
|
||||
func BenchmarkUpsert_ColdFields(b *testing.B) {
|
||||
ns := "tenant-x"
|
||||
ops := make([]*schemapb.FieldPartialUpdateOp, 16)
|
||||
for i := range ops {
|
||||
ops[i] = &schemapb.FieldPartialUpdateOp{FieldName: fmt.Sprintf("f%d", i)}
|
||||
}
|
||||
benchUpsert(b, &milvuspb.UpsertRequest{
|
||||
CollectionName: "c", NumRows: 1000, PartialUpdate: true, Namespace: &ns,
|
||||
FieldOps: ops, FieldsData: varcharFieldData(1000, 12),
|
||||
})
|
||||
}
|
||||
|
||||
// BenchmarkUpsert_Malformed measures the reject path: a fields_data (5)
|
||||
// sub-message whose declared length overruns the buffer, which both the official
|
||||
// codec and fastpb must reject.
|
||||
func BenchmarkUpsert_Malformed(b *testing.B) {
|
||||
wire := []byte{0x2A, 0x04, 0x08} // field 5, wtype 2, len=4, but only 1 byte follows
|
||||
b.Run("official", func(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out milvuspb.UpsertRequest
|
||||
_ = proto.Unmarshal(wire, &out)
|
||||
}
|
||||
})
|
||||
b.Run("fastpb", func(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out milvuspb.UpsertRequest
|
||||
_ = UnmarshalUpsertRequest(wire, &out)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
commonpb "github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
func TestEquiv_RetrieveResults(t *testing.T) {
|
||||
src := &internalpb.RetrieveResults{
|
||||
Base: &commonpb.MsgBase{MsgID: 7},
|
||||
Status: &commonpb.Status{Code: 0, Reason: "ok"},
|
||||
ReqID: 99,
|
||||
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
|
||||
SealedSegmentIDsRetrieved: []int64{10, 20, 30},
|
||||
ChannelIDsRetrieved: []string{"ch1", "ch2"},
|
||||
GlobalSealedSegmentIDs: []int64{100, 200},
|
||||
AllRetrieveCount: 3,
|
||||
HasMoreResult: true,
|
||||
ScannedTotalBytes: 4096,
|
||||
FieldsData: []*schemapb.FieldData{
|
||||
{Type: schemapb.DataType_VarChar, FieldName: "title", FieldId: 101, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"a", "b", "c"}}}}}},
|
||||
{Type: schemapb.DataType_FloatVector, FieldName: "emb", FieldId: 102, Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: 2, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}}}}}},
|
||||
},
|
||||
}
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got internalpb.RetrieveResults
|
||||
if err := UnmarshalRetrieveResults(wire, &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch:\n src=%v\n got=%v", src, &got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEquiv_InsertRequest(t *testing.T) {
|
||||
ns := "tenant-x"
|
||||
src := &milvuspb.InsertRequest{
|
||||
Base: &commonpb.MsgBase{MsgID: 1},
|
||||
DbName: "db",
|
||||
CollectionName: "coll",
|
||||
PartitionName: "part",
|
||||
HashKeys: []uint32{1, 2, 3, 4},
|
||||
NumRows: 3,
|
||||
SchemaTimestamp: 123456,
|
||||
Namespace: &ns,
|
||||
FieldsData: []*schemapb.FieldData{
|
||||
{Type: schemapb.DataType_VarChar, FieldName: "title", FieldId: 101, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"x", "y", "z"}}}}}},
|
||||
{Type: schemapb.DataType_FloatVector, FieldName: "emb", FieldId: 102, Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: 2, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}}}}}},
|
||||
},
|
||||
}
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got milvuspb.InsertRequest
|
||||
if err := UnmarshalInsertRequest(wire, &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch:\n src=%v\n got=%v", src, &got)
|
||||
}
|
||||
}
|
||||
|
||||
// InsertRequest is untrusted ingress: fastpb must agree with official proto on
|
||||
// whether invalid UTF-8 in a string field is rejected.
|
||||
func TestInsertRequest_UTF8MatchesOfficial(t *testing.T) {
|
||||
// field 3 (collection_name), wire type 2, len 2, bytes 0xFF 0xFE (invalid UTF-8)
|
||||
wire := []byte{0x1A, 0x02, 0xFF, 0xFE}
|
||||
var official milvuspb.InsertRequest
|
||||
officialErr := proto.Unmarshal(wire, &official)
|
||||
var got milvuspb.InsertRequest
|
||||
gotErr := UnmarshalInsertRequest(wire, &got)
|
||||
if (officialErr == nil) != (gotErr == nil) {
|
||||
t.Fatalf("UTF-8 accept/reject diverges: official err=%v, fastpb err=%v", officialErr, gotErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTryUnmarshal_Dispatch(t *testing.T) {
|
||||
// known type → handled
|
||||
rr := &internalpb.RetrieveResults{ReqID: 5}
|
||||
wire, _ := proto.Marshal(rr)
|
||||
var got internalpb.RetrieveResults
|
||||
handled, err := TryUnmarshal(&got, wire)
|
||||
if !handled || err != nil || got.ReqID != 5 {
|
||||
t.Fatalf("expected handled RetrieveResults: handled=%v err=%v reqID=%d", handled, err, got.ReqID)
|
||||
}
|
||||
// unknown type → not handled (caller falls back to official)
|
||||
var status commonpb.Status
|
||||
handled, err = TryUnmarshal(&status, nil)
|
||||
if handled || err != nil {
|
||||
t.Fatalf("expected unhandled for *commonpb.Status: handled=%v err=%v", handled, err)
|
||||
}
|
||||
}
|
||||
|
||||
func randRetrieveResults(r *rand.Rand) *internalpb.RetrieveResults {
|
||||
rows := r.Intn(20)
|
||||
rr := &internalpb.RetrieveResults{
|
||||
ReqID: int64(r.Uint32()),
|
||||
AllRetrieveCount: int64(r.Uint32()),
|
||||
HasMoreResult: r.Intn(2) == 0,
|
||||
}
|
||||
if r.Intn(2) == 0 {
|
||||
ids := make([]int64, rows)
|
||||
for i := range ids {
|
||||
ids[i] = int64(r.Uint64())
|
||||
}
|
||||
rr.Ids = &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: ids}}}
|
||||
}
|
||||
for i := 0; i < r.Intn(3); i++ {
|
||||
rr.FieldsData = append(rr.FieldsData, randFieldData(r, rows))
|
||||
}
|
||||
for i := 0; i < r.Intn(3); i++ {
|
||||
rr.ChannelIDsRetrieved = append(rr.ChannelIDsRetrieved, randString(r))
|
||||
}
|
||||
return rr
|
||||
}
|
||||
|
||||
func randInsertRequest(r *rand.Rand) *milvuspb.InsertRequest {
|
||||
rows := r.Intn(20)
|
||||
ir := &milvuspb.InsertRequest{
|
||||
DbName: randString(r),
|
||||
CollectionName: randString(r),
|
||||
PartitionName: randString(r),
|
||||
NumRows: uint32(rows),
|
||||
}
|
||||
for i := 0; i < r.Intn(3); i++ {
|
||||
ir.FieldsData = append(ir.FieldsData, randFieldData(r, rows))
|
||||
}
|
||||
hk := make([]uint32, rows)
|
||||
for i := range hk {
|
||||
hk[i] = r.Uint32()
|
||||
}
|
||||
ir.HashKeys = hk
|
||||
return ir
|
||||
}
|
||||
|
||||
func TestEquiv_Fuzz_RetrieveResults(t *testing.T) {
|
||||
r := rand.New(rand.NewSource(0xBEEF))
|
||||
for i := 0; i < 3000; i++ {
|
||||
src := randRetrieveResults(r)
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal %d: %v", i, err)
|
||||
}
|
||||
var got internalpb.RetrieveResults
|
||||
if err := UnmarshalRetrieveResults(wire, &got); err != nil {
|
||||
t.Fatalf("decode %d: %v", i, err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch %d:\n src=%v\n got=%v", i, src, &got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEquiv_Fuzz_InsertRequest(t *testing.T) {
|
||||
r := rand.New(rand.NewSource(0xFEED))
|
||||
for i := 0; i < 3000; i++ {
|
||||
src := randInsertRequest(r)
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal %d: %v", i, err)
|
||||
}
|
||||
var got milvuspb.InsertRequest
|
||||
if err := UnmarshalInsertRequest(wire, &got); err != nil {
|
||||
t.Fatalf("decode %d: %v", i, err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch %d:\n src=%v\n got=%v", i, src, &got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEquiv_UpsertRequest(t *testing.T) {
|
||||
ns := "tenant-x"
|
||||
src := &milvuspb.UpsertRequest{
|
||||
Base: &commonpb.MsgBase{MsgID: 1},
|
||||
DbName: "db",
|
||||
CollectionName: "coll",
|
||||
PartitionName: "part",
|
||||
HashKeys: []uint32{1, 2, 3, 4},
|
||||
NumRows: 3,
|
||||
SchemaTimestamp: 123456,
|
||||
PartialUpdate: true, // field 9 (varint) → folded to official merge
|
||||
Namespace: &ns, // field 10 (bytes) → folded to official merge
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{ // field 11 (message) → folded
|
||||
{FieldName: "title"},
|
||||
},
|
||||
FieldsData: []*schemapb.FieldData{
|
||||
{Type: schemapb.DataType_VarChar, FieldName: "title", FieldId: 101, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"x", "y", "z"}}}}}},
|
||||
{Type: schemapb.DataType_FloatVector, FieldName: "emb", FieldId: 102, Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: 2, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}}}}}},
|
||||
},
|
||||
}
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got milvuspb.UpsertRequest
|
||||
if err := UnmarshalUpsertRequest(wire, &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch:\n src=%v\n got=%v", src, &got)
|
||||
}
|
||||
// TryUnmarshal must dispatch UpsertRequest to fastpb (RPC-boundary coverage).
|
||||
var viaDispatch milvuspb.UpsertRequest
|
||||
handled, err := TryUnmarshal(&viaDispatch, wire)
|
||||
if !handled || err != nil || !proto.Equal(src, &viaDispatch) {
|
||||
t.Fatalf("TryUnmarshal(UpsertRequest) dispatch failed: handled=%v err=%v", handled, err)
|
||||
}
|
||||
}
|
||||
|
||||
// UpsertRequest is untrusted ingress too: fastpb must agree with official proto on
|
||||
// whether invalid UTF-8 in a string field is rejected.
|
||||
func TestUpsertRequest_UTF8MatchesOfficial(t *testing.T) {
|
||||
// field 3 (collection_name), wire type 2, len 2, bytes 0xFF 0xFE (invalid UTF-8)
|
||||
wire := []byte{0x1A, 0x02, 0xFF, 0xFE}
|
||||
var official milvuspb.UpsertRequest
|
||||
officialErr := proto.Unmarshal(wire, &official)
|
||||
var got milvuspb.UpsertRequest
|
||||
gotErr := UnmarshalUpsertRequest(wire, &got)
|
||||
if (officialErr == nil) != (gotErr == nil) {
|
||||
t.Fatalf("UTF-8 accept/reject diverges: official err=%v, fastpb err=%v", officialErr, gotErr)
|
||||
}
|
||||
}
|
||||
|
||||
func randUpsertRequest(r *rand.Rand) *milvuspb.UpsertRequest {
|
||||
rows := r.Intn(20)
|
||||
ur := &milvuspb.UpsertRequest{
|
||||
DbName: randString(r),
|
||||
CollectionName: randString(r),
|
||||
PartitionName: randString(r),
|
||||
NumRows: uint32(rows),
|
||||
PartialUpdate: r.Intn(2) == 0, // exercise the field-9 fold path
|
||||
}
|
||||
for i := 0; i < r.Intn(3); i++ {
|
||||
ur.FieldsData = append(ur.FieldsData, randFieldData(r, rows))
|
||||
}
|
||||
hk := make([]uint32, rows)
|
||||
for i := range hk {
|
||||
hk[i] = r.Uint32()
|
||||
}
|
||||
ur.HashKeys = hk
|
||||
if r.Intn(2) == 0 {
|
||||
ns := randString(r) // exercise the field-10 fold path
|
||||
ur.Namespace = &ns
|
||||
}
|
||||
for i := 0; i < r.Intn(3); i++ { // exercise the field-11 (field_ops) fold path
|
||||
ur.FieldOps = append(ur.FieldOps, &schemapb.FieldPartialUpdateOp{
|
||||
FieldName: randString(r),
|
||||
Op: schemapb.FieldPartialUpdateOp_OpType(r.Intn(2)),
|
||||
})
|
||||
}
|
||||
return ur
|
||||
}
|
||||
|
||||
func TestEquiv_Fuzz_UpsertRequest(t *testing.T) {
|
||||
r := rand.New(rand.NewSource(0xCAFE))
|
||||
for i := 0; i < 3000; i++ {
|
||||
src := randUpsertRequest(r)
|
||||
wire, err := proto.Marshal(src)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal %d: %v", i, err)
|
||||
}
|
||||
var got milvuspb.UpsertRequest
|
||||
if err := UnmarshalUpsertRequest(wire, &got); err != nil {
|
||||
t.Fatalf("decode %d: %v", i, err)
|
||||
}
|
||||
if !proto.Equal(src, &got) {
|
||||
t.Fatalf("mismatch %d:\n src=%v\n got=%v", i, src, &got)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"math"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
commonpb "github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
)
|
||||
|
||||
// UnmarshalSearchResultData decodes a wire-format schemapb.SearchResultData into srd.
|
||||
// Hot fields (fields_data, scores, ids, topks, distances, output_fields, ...) are
|
||||
// hand-decoded; rare fields (iterator results, highlights, agg buckets) delegate
|
||||
// to the official codec. Equivalence is pinned by the differential + fuzz tests.
|
||||
func UnmarshalSearchResultData(b []byte, srd *schemapb.SearchResultData) error {
|
||||
proto.Reset(srd) // match official proto.Unmarshal: clear target before decode
|
||||
if err := (dec{}).searchResultData(b, srd); err != nil {
|
||||
if fallbackOnProto2(err) {
|
||||
proto.Reset(srd) // discard partial decode before the authoritative pass
|
||||
return proto.Unmarshal(b, srd)
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d dec) searchResultData(b []byte, srd *schemapb.SearchResultData) error {
|
||||
var rest []byte
|
||||
for len(b) > 0 {
|
||||
start := b
|
||||
num, wtype, tn := consumeTag(b)
|
||||
if tn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[tn:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if wtype == 5 && (num == 4 || num == 10 || num == 12) {
|
||||
if len(b) < 4 {
|
||||
return errMalformed
|
||||
}
|
||||
value := math.Float32frombits(le32(b))
|
||||
switch num {
|
||||
case 4:
|
||||
srd.Scores = append(srd.Scores, value)
|
||||
case 10:
|
||||
srd.Distances = append(srd.Distances, value)
|
||||
case 12:
|
||||
srd.Recalls = append(srd.Recalls, value)
|
||||
}
|
||||
b = b[4:]
|
||||
continue
|
||||
}
|
||||
|
||||
// varint-typed fields
|
||||
if wtype == 0 {
|
||||
v, vn := consumeVarint(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
switch num {
|
||||
case 1:
|
||||
srd.NumQueries = int64(v)
|
||||
case 2:
|
||||
srd.TopK = int64(v)
|
||||
case 9:
|
||||
srd.AllSearchCount = int64(v)
|
||||
// packed repeated fields may also arrive as single varints:
|
||||
case 6:
|
||||
srd.Topks = append(srd.Topks, int64(v))
|
||||
case 19:
|
||||
srd.AggTopks = append(srd.AggTopks, int64(v))
|
||||
default:
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
b = b[vn:]
|
||||
continue
|
||||
}
|
||||
|
||||
// everything else here is length-delimited
|
||||
if wtype != 2 {
|
||||
sn := skipField(b, wtype)
|
||||
if sn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
rest = append(rest, start[:tn+sn]...)
|
||||
b = b[sn:]
|
||||
continue
|
||||
}
|
||||
v, vn := consumeBytes(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
|
||||
switch num {
|
||||
case 3: // fields_data (repeated FieldData)
|
||||
fd := &schemapb.FieldData{}
|
||||
if err := d.fieldData(v, fd); err != nil {
|
||||
return err
|
||||
}
|
||||
srd.FieldsData = append(srd.FieldsData, fd)
|
||||
case 4: // scores (packed fixed32)
|
||||
if err := appendPackedF32(v, &srd.Scores); err != nil {
|
||||
return err
|
||||
}
|
||||
case 5: // ids (IDs)
|
||||
ids := &schemapb.IDs{}
|
||||
if err := d.ids(v, ids); err != nil {
|
||||
return err
|
||||
}
|
||||
if srd.Ids == nil {
|
||||
srd.Ids = ids
|
||||
} else {
|
||||
// wire-level merge: a later occurrence's explicit proto3-default
|
||||
// scalar must overwrite, matching proto.Unmarshal (not proto.Merge).
|
||||
if err := protoMerge(v, srd.Ids); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 6: // topks (packed varint)
|
||||
if err := appendPackedI64(v, &srd.Topks); err != nil {
|
||||
return err
|
||||
}
|
||||
case 7: // output_fields (repeated string)
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
srd.OutputFields = append(srd.OutputFields, s)
|
||||
case 8: // group_by_field_value (FieldData)
|
||||
fd := &schemapb.FieldData{}
|
||||
if err := d.fieldData(v, fd); err != nil {
|
||||
return err
|
||||
}
|
||||
if srd.GroupByFieldValue == nil {
|
||||
srd.GroupByFieldValue = fd
|
||||
} else {
|
||||
if err := protoMerge(v, srd.GroupByFieldValue); err != nil { // wire-level merge (see Ids)
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 10: // distances (packed fixed32)
|
||||
if err := appendPackedF32(v, &srd.Distances); err != nil {
|
||||
return err
|
||||
}
|
||||
case 11: // search_iterator_v2_results (delegate)
|
||||
m := &schemapb.SearchIteratorV2Results{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
if srd.SearchIteratorV2Results == nil {
|
||||
srd.SearchIteratorV2Results = m
|
||||
} else {
|
||||
if err := protoMerge(v, srd.SearchIteratorV2Results); err != nil { // wire-level merge (see Ids)
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 12: // recalls (packed fixed32)
|
||||
if err := appendPackedF32(v, &srd.Recalls); err != nil {
|
||||
return err
|
||||
}
|
||||
case 13: // primary_field_name (string)
|
||||
s, err := d.str(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
srd.PrimaryFieldName = s
|
||||
case 14: // highlight_results (repeated commonpb.HighlightResult, delegate)
|
||||
m := &commonpb.HighlightResult{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
srd.HighlightResults = append(srd.HighlightResults, m)
|
||||
case 15: // element_indices (LongArray)
|
||||
la := &schemapb.LongArray{}
|
||||
if err := decodePackedI64(v, &la.Data, la); err != nil {
|
||||
return err
|
||||
}
|
||||
if srd.ElementIndices == nil {
|
||||
srd.ElementIndices = la
|
||||
} else {
|
||||
if err := protoMerge(v, srd.ElementIndices); err != nil { // wire-level merge (see Ids)
|
||||
return err
|
||||
}
|
||||
}
|
||||
case 17: // group_by_field_values (repeated FieldData)
|
||||
fd := &schemapb.FieldData{}
|
||||
if err := d.fieldData(v, fd); err != nil {
|
||||
return err
|
||||
}
|
||||
srd.GroupByFieldValues = append(srd.GroupByFieldValues, fd)
|
||||
case 18: // agg_buckets (repeated AggBucket, delegate)
|
||||
m := &schemapb.AggBucket{}
|
||||
if err := protoUnmarshal(v, m); err != nil {
|
||||
return err
|
||||
}
|
||||
srd.AggBuckets = append(srd.AggBuckets, m)
|
||||
case 19: // agg_topks (packed varint)
|
||||
if err := appendPackedI64(v, &srd.AggTopks); err != nil {
|
||||
return err
|
||||
}
|
||||
default: // unhandled (future) field → fold into official merge
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
}
|
||||
if len(rest) > 0 {
|
||||
return protoMerge(rest, srd)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// unmarshalIDs decodes schemapb.IDs: oneof int_id (LongArray, 1) / str_id (StringArray, 2).
|
||||
func (d dec) ids(b []byte, ids *schemapb.IDs) error {
|
||||
full := b
|
||||
var rest []byte
|
||||
oneofNum := 0
|
||||
for len(b) > 0 {
|
||||
start := b
|
||||
num, wtype, tn := consumeTag(b)
|
||||
if tn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[tn:]
|
||||
if isProto2Group(wtype) {
|
||||
return errProto2
|
||||
}
|
||||
if wtype != 2 {
|
||||
sn := skipField(b, wtype)
|
||||
if sn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
rest = append(rest, start[:tn+sn]...)
|
||||
b = b[sn:]
|
||||
continue
|
||||
}
|
||||
v, vn := consumeBytes(b)
|
||||
if vn <= 0 {
|
||||
return errMalformed
|
||||
}
|
||||
b = b[vn:]
|
||||
if num == 1 || num == 2 {
|
||||
if oneofNum == num {
|
||||
return fallbackUnmarshal(full, ids)
|
||||
}
|
||||
oneofNum = num
|
||||
}
|
||||
switch num {
|
||||
case 1:
|
||||
la := &schemapb.LongArray{}
|
||||
if err := decodePackedI64(v, &la.Data, la); err != nil {
|
||||
return err
|
||||
}
|
||||
ids.IdField = &schemapb.IDs_IntId{IntId: la}
|
||||
case 2:
|
||||
sa := &schemapb.StringArray{}
|
||||
if err := d.stringArray(v, sa); err != nil {
|
||||
return err
|
||||
}
|
||||
ids.IdField = &schemapb.IDs_StrId{StrId: sa}
|
||||
default:
|
||||
rest = append(rest, start[:tn+vn]...)
|
||||
}
|
||||
}
|
||||
if len(rest) > 0 {
|
||||
return protoMerge(rest, ids)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"unicode/utf8"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// validUTF8 reports whether v is valid UTF-8 (matches proto3 string validation).
|
||||
func validUTF8(v []byte) bool { return utf8.Valid(v) }
|
||||
|
||||
// This file isolates the only unsafe in the package. Both helpers are
|
||||
// little-endian (x86/arm64) and produce freshly-allocated, self-owned memory —
|
||||
// they do NOT alias the source buffer, so they are safe regardless of the
|
||||
// source buffer's lifetime (gRPC-pooled or owned). A separate zero-copy
|
||||
// aliasing variant (for the owned sliced_blob path) can be added later.
|
||||
|
||||
// bytesToF32 returns a new []float32 holding the little-endian float32s in v
|
||||
// via a single memcpy (len(v) must be a multiple of 4).
|
||||
func bytesToF32(v []byte) []float32 {
|
||||
n := len(v) / 4
|
||||
out := make([]float32, n)
|
||||
if n > 0 {
|
||||
copy(unsafe.Slice((*byte)(unsafe.Pointer(&out[0])), n*4), v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// bytesToF64 returns a new []float64 holding the little-endian float64s in v
|
||||
// via a single memcpy (len(v) must be a multiple of 8).
|
||||
func bytesToF64(v []byte) []float64 {
|
||||
n := len(v) / 8
|
||||
out := make([]float64, n)
|
||||
if n > 0 {
|
||||
copy(unsafe.Slice((*byte)(unsafe.Pointer(&out[0])), n*8), v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// strings copies the string elements out of one StringArray submessage into
|
||||
// a single backing byte arena and one backing []string, collapsing N per-string
|
||||
// allocations into 2. Strings point into the arena via unsafe.String; the arena
|
||||
// never reallocates (cap = upper bound), so the pointers stay valid for the life
|
||||
// of the returned slice. Validates UTF-8 per element when d.utf8 is set.
|
||||
//
|
||||
// The bool return is false if the StringArray carries any field other than the
|
||||
// `data` field (field 1) — the caller then bails to the official codec so the
|
||||
// unknown/future field is preserved exactly.
|
||||
func (d dec) strings(b []byte) ([]string, bool, error) {
|
||||
// pass 1: count elements (field 1) and total string bytes
|
||||
count, total := 0, 0
|
||||
p := b
|
||||
for len(p) > 0 {
|
||||
num, wtype, n := consumeTag(p)
|
||||
if n <= 0 {
|
||||
return nil, false, errMalformed
|
||||
}
|
||||
p = p[n:]
|
||||
if num != 1 || wtype != 2 {
|
||||
return nil, false, nil // unknown field present → caller falls back
|
||||
}
|
||||
v, n := consumeBytes(p)
|
||||
if n <= 0 {
|
||||
return nil, false, errMalformed
|
||||
}
|
||||
count++
|
||||
total += len(v)
|
||||
p = p[n:]
|
||||
}
|
||||
if count == 0 {
|
||||
return nil, true, nil
|
||||
}
|
||||
out := make([]string, 0, count)
|
||||
arena := make([]byte, 0, total)
|
||||
// pass 2: fill
|
||||
p = b
|
||||
for len(p) > 0 {
|
||||
_, _, n := consumeTag(p)
|
||||
p = p[n:]
|
||||
v, n := consumeBytes(p)
|
||||
if d.utf8 && !validUTF8(v) {
|
||||
return nil, false, errInvalidUTF8
|
||||
}
|
||||
if len(v) == 0 {
|
||||
out = append(out, "")
|
||||
} else {
|
||||
off := len(arena)
|
||||
arena = append(arena, v...) // never reallocates: cap == total
|
||||
out = append(out, unsafe.String(&arena[off], len(v)))
|
||||
}
|
||||
p = p[n:]
|
||||
}
|
||||
return out, true, nil
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
package fastpb
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
schemapb "github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
)
|
||||
|
||||
// TestWireTypeMismatchFallback: a known field carrying an unexpected wire type is
|
||||
// not something proto3 produces (only malformed / adversarial encoders do). fastpb
|
||||
// must hand such a message to the official codec and produce the exact same result
|
||||
// (and the same error behavior) as proto.Unmarshal, rather than blindly decoding
|
||||
// the wrong shape.
|
||||
func TestWireTypeMismatchFallback(t *testing.T) {
|
||||
// FieldData with field 5 (FieldId, a varint field) wrongly length-delimited,
|
||||
// and field 1 (Type, varint) wrongly length-delimited.
|
||||
t.Run("FieldData varint field as length-delimited", func(t *testing.T) {
|
||||
var fd []byte
|
||||
fd = protowire.AppendTag(fd, 2, protowire.BytesType) // FieldName (legit)
|
||||
fd = protowire.AppendString(fd, "f")
|
||||
fd = protowire.AppendTag(fd, 5, protowire.BytesType) // FieldId: should be varint(0)
|
||||
fd = protowire.AppendBytes(fd, []byte{0x01, 0x02})
|
||||
|
||||
want := &schemapb.FieldData{}
|
||||
wantErr := proto.Unmarshal(fd, want)
|
||||
got := &schemapb.FieldData{}
|
||||
gotErr := UnmarshalFieldData(fd, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want), "fallback must equal official decode")
|
||||
}
|
||||
})
|
||||
|
||||
// VectorField with field 1 (Dim, varint) wrongly length-delimited.
|
||||
t.Run("VectorField Dim as length-delimited", func(t *testing.T) {
|
||||
var vf []byte
|
||||
vf = protowire.AppendTag(vf, 1, protowire.BytesType) // Dim: should be varint(0)
|
||||
vf = protowire.AppendBytes(vf, []byte{0x07})
|
||||
|
||||
want := &schemapb.VectorField{}
|
||||
wantErr := proto.Unmarshal(vf, want)
|
||||
got := &schemapb.VectorField{}
|
||||
gotErr := unmarshalVectorField(vf, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want))
|
||||
}
|
||||
})
|
||||
|
||||
// SparseFloatArray with field 2 (dim, varint) wrongly length-delimited.
|
||||
t.Run("SparseFloatArray dim as length-delimited", func(t *testing.T) {
|
||||
var sa []byte
|
||||
sa = protowire.AppendTag(sa, 2, protowire.BytesType) // dim: should be varint(0)
|
||||
sa = protowire.AppendBytes(sa, []byte{0x05})
|
||||
|
||||
want := &schemapb.SparseFloatArray{}
|
||||
wantErr := proto.Unmarshal(sa, want)
|
||||
got := &schemapb.SparseFloatArray{}
|
||||
gotErr := unmarshalSparseFloatArray(sa, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want))
|
||||
}
|
||||
})
|
||||
|
||||
// Same mismatch nested inside an InsertRequest's fields_data (exercises the
|
||||
// untrusted ingress entry point + nested propagation).
|
||||
t.Run("InsertRequest with mismatched nested FieldData", func(t *testing.T) {
|
||||
var fd []byte
|
||||
fd = protowire.AppendTag(fd, 1, protowire.BytesType) // Type: should be varint(0)
|
||||
fd = protowire.AppendBytes(fd, []byte{0x09})
|
||||
|
||||
var wire []byte
|
||||
wire = protowire.AppendTag(wire, 3, protowire.BytesType) // collection_name
|
||||
wire = protowire.AppendString(wire, "c")
|
||||
wire = protowire.AppendTag(wire, 5, protowire.BytesType) // fields_data
|
||||
wire = protowire.AppendBytes(wire, fd)
|
||||
|
||||
want := &milvuspb.InsertRequest{}
|
||||
wantErr := proto.Unmarshal(wire, want)
|
||||
got := &milvuspb.InsertRequest{}
|
||||
gotErr := UnmarshalInsertRequest(wire, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestVarintArrayWireTypeMismatch: field 1 of a varint-family repeated field
|
||||
// (LongArray/IntArray/BoolArray) encoded with a fixed wire type is not valid
|
||||
// proto3 output — only a malformed or adversarial encoder produces it. The
|
||||
// varint decoders must hand such messages to the official codec (which treats
|
||||
// the mismatched field as unknown) instead of consuming the fixed-width bytes
|
||||
// as a varint and silently decoding garbage values.
|
||||
func TestVarintArrayWireTypeMismatch(t *testing.T) {
|
||||
t.Run("LongArray data as fixed64", func(t *testing.T) {
|
||||
var w []byte
|
||||
w = protowire.AppendTag(w, 1, protowire.Fixed64Type)
|
||||
w = protowire.AppendFixed64(w, 42)
|
||||
|
||||
want := &schemapb.LongArray{}
|
||||
wantErr := proto.Unmarshal(w, want)
|
||||
got := &schemapb.LongArray{}
|
||||
gotErr := decodePackedI64(w, &got.Data, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want), "mismatched wire type must not decode as data")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("IntArray data as fixed32", func(t *testing.T) {
|
||||
var w []byte
|
||||
w = protowire.AppendTag(w, 1, protowire.Fixed32Type)
|
||||
w = protowire.AppendFixed32(w, 7)
|
||||
|
||||
want := &schemapb.IntArray{}
|
||||
wantErr := proto.Unmarshal(w, want)
|
||||
got := &schemapb.IntArray{}
|
||||
gotErr := decodePackedI32(w, &got.Data, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want), "mismatched wire type must not decode as data")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("BoolArray data as fixed64", func(t *testing.T) {
|
||||
var w []byte
|
||||
w = protowire.AppendTag(w, 1, protowire.Fixed64Type)
|
||||
w = protowire.AppendFixed64(w, 1)
|
||||
|
||||
want := &schemapb.BoolArray{}
|
||||
wantErr := proto.Unmarshal(w, want)
|
||||
got := &schemapb.BoolArray{}
|
||||
gotErr := decodePackedBool(w, &got.Data, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want), "mismatched wire type must not decode as data")
|
||||
}
|
||||
})
|
||||
|
||||
// The adversarial shape: an alternating fixed32-tag/byte pattern in which
|
||||
// every byte re-parses as a valid field-1 tag or varint. The pre-fix decoders
|
||||
// consumed this WITHOUT error as [8,8,8,8] — silently diverging from the
|
||||
// official codec, which consumes the fixed32 payload into unknown fields and
|
||||
// decodes a different, shorter array. This is the case the wire-type check
|
||||
// exists for; the simple vectors above happen to be rescued by the num!=1
|
||||
// fallback, this one is not.
|
||||
t.Run("LongArray alternating fixed32 tags decode divergence", func(t *testing.T) {
|
||||
w := []byte{0x0D, 0x08, 0x0D, 0x08, 0x0D, 0x08, 0x0D, 0x08} // tag(1,fixed32) interleaved
|
||||
|
||||
want := &schemapb.LongArray{}
|
||||
wantErr := proto.Unmarshal(w, want)
|
||||
got := &schemapb.LongArray{}
|
||||
gotErr := decodePackedI64(w, &got.Data, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want),
|
||||
"want %v got %v: mismatched wire type must fall back, not decode as varints", want.Data, got.Data)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("IntArray alternating fixed32 tags decode divergence", func(t *testing.T) {
|
||||
w := []byte{0x0D, 0x08, 0x0D, 0x08, 0x0D, 0x08, 0x0D, 0x08}
|
||||
|
||||
want := &schemapb.IntArray{}
|
||||
wantErr := proto.Unmarshal(w, want)
|
||||
got := &schemapb.IntArray{}
|
||||
gotErr := decodePackedI32(w, &got.Data, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want),
|
||||
"want %v got %v: mismatched wire type must fall back, not decode as varints", want.Data, got.Data)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("BoolArray alternating fixed32 tags decode divergence", func(t *testing.T) {
|
||||
w := []byte{0x0D, 0x08, 0x0D, 0x08, 0x0D, 0x08, 0x0D, 0x08}
|
||||
|
||||
want := &schemapb.BoolArray{}
|
||||
wantErr := proto.Unmarshal(w, want)
|
||||
got := &schemapb.BoolArray{}
|
||||
gotErr := decodePackedBool(w, &got.Data, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want),
|
||||
"want %v got %v: mismatched wire type must fall back, not decode as varints", want.Data, got.Data)
|
||||
}
|
||||
})
|
||||
|
||||
// Mixed: a legitimate packed chunk followed by a mismatched wire type on the
|
||||
// same field — the whole message must fall back, keeping parity with the
|
||||
// official codec rather than keeping the half-decoded prefix.
|
||||
t.Run("LongArray packed then fixed64", func(t *testing.T) {
|
||||
var packed []byte
|
||||
packed = protowire.AppendVarint(packed, 1)
|
||||
packed = protowire.AppendVarint(packed, 2)
|
||||
var w []byte
|
||||
w = protowire.AppendTag(w, 1, protowire.BytesType)
|
||||
w = protowire.AppendBytes(w, packed)
|
||||
w = protowire.AppendTag(w, 1, protowire.Fixed64Type)
|
||||
w = protowire.AppendFixed64(w, 42)
|
||||
|
||||
want := &schemapb.LongArray{}
|
||||
wantErr := proto.Unmarshal(w, want)
|
||||
got := &schemapb.LongArray{}
|
||||
gotErr := decodePackedI64(w, &got.Data, got)
|
||||
|
||||
require.Equal(t, wantErr == nil, gotErr == nil, "error parity vs official (want=%v got=%v)", wantErr, gotErr)
|
||||
if wantErr == nil {
|
||||
assert.True(t, proto.Equal(got, want), "fallback must equal official decode")
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -22,6 +22,7 @@ import (
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/protoadapt"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/fastpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
||||
)
|
||||
|
||||
@@ -85,6 +86,12 @@ func (releaseCodec) Unmarshal(data mem.BufferSlice, v any) error {
|
||||
|
||||
buf := data.MaterializeToBuffer(mem.DefaultBufferPool())
|
||||
defer buf.Free()
|
||||
// Fast path for the top-level hot RPC messages supported by TryUnmarshal:
|
||||
// RetrieveResults, InsertRequest, and UpsertRequest. Unsupported message
|
||||
// types fall through to the official codec.
|
||||
if handled, err := fastpb.TryUnmarshal(v, buf.ReadOnlyData()); handled {
|
||||
return err
|
||||
}
|
||||
return proto.Unmarshal(buf.ReadOnlyData(), msg)
|
||||
}
|
||||
|
||||
|
||||
@@ -21,7 +21,10 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/grpc/mem"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
milvuspb "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
||||
)
|
||||
|
||||
@@ -156,6 +159,41 @@ func TestSearchResultsImplementsMsgPinnable(t *testing.T) {
|
||||
require.False(t, ok, "SearchRequest should not implement MsgPinnable")
|
||||
}
|
||||
|
||||
// TestReleaseCodecFastpbFastPath covers the fastpb.TryUnmarshal fast path in
|
||||
// Unmarshal: the two hot types decode through fastpb (round-trip equal to the
|
||||
// official codec), and a fast-path decode error is surfaced to the caller.
|
||||
func TestReleaseCodecFastpbFastPath(t *testing.T) {
|
||||
codec := releaseCodec{}
|
||||
|
||||
rr := &internalpb.RetrieveResults{ReqID: 42, ChannelIDsRetrieved: []string{"ch1"}}
|
||||
out, err := codec.Marshal(rr)
|
||||
require.NoError(t, err)
|
||||
gotRR := &internalpb.RetrieveResults{}
|
||||
require.NoError(t, codec.Unmarshal(out, gotRR))
|
||||
require.True(t, proto.Equal(rr, gotRR))
|
||||
|
||||
ir := &milvuspb.InsertRequest{CollectionName: "c", NumRows: 3}
|
||||
out, err = codec.Marshal(ir)
|
||||
require.NoError(t, err)
|
||||
gotIR := &milvuspb.InsertRequest{}
|
||||
require.NoError(t, codec.Unmarshal(out, gotIR))
|
||||
require.True(t, proto.Equal(ir, gotIR))
|
||||
|
||||
// UpsertRequest takes the fast path too; the upsert-only fields
|
||||
// (partial_update/namespace/field_ops) fold to the official codec.
|
||||
ns := "tenant-x"
|
||||
ur := &milvuspb.UpsertRequest{CollectionName: "c", NumRows: 3, PartialUpdate: true, Namespace: &ns}
|
||||
out, err = codec.Marshal(ur)
|
||||
require.NoError(t, err)
|
||||
gotUR := &milvuspb.UpsertRequest{}
|
||||
require.NoError(t, codec.Unmarshal(out, gotUR))
|
||||
require.True(t, proto.Equal(ur, gotUR))
|
||||
|
||||
// malformed wire (truncated varint tag) → fast path must return the error
|
||||
bad := mem.BufferSlice{mem.SliceBuffer([]byte{0x80})}
|
||||
require.Error(t, codec.Unmarshal(bad, &internalpb.RetrieveResults{}))
|
||||
}
|
||||
|
||||
func TestReleaseCodecMarshalNotPinnedPinnable(t *testing.T) {
|
||||
// A MsgPinnable message that is NOT pinned — Marshal should succeed without error.
|
||||
MsgPins.ResetForTest()
|
||||
|
||||
Reference in New Issue
Block a user