enhance: Add storage v1 binlog schema healthz check

states/healthz/check_item.go

@@ -26,6 +26,7 @@ func init() {
 	allCheckItems["DYNAMIC_ATTR_MISMATCH"] = newDynamicAttrMismatch()
 	allCheckItems["LOADED_INDEX_MISSING"] = newLoadedIndexMissing()
 	allCheckItems["MIXED_BINLOGS"] = newMixedBinlogs()
+	allCheckItems["STORAGE_V1_BINLOG_SCHEMA_MISMATCH"] = newStorageV1BinlogSchemaMismatch()
 	allCheckItems["GRANT_ALIAS_CHECK"] = newGrantAliasCheck()
 
 	defaultList = []string{

states/healthz/storage_v1_binlog_schema_mismatch.go

@@ -0,0 +1,103 @@
+package healthz
+
+import (
+	"context"
+	"fmt"
+	"sort"
+
+	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/states/etcd/common"
+	metakv "github.com/milvus-io/birdwatcher/states/kv"
+)
+
+type StorageV1BinlogSchemaMismatch struct {
+	checkItemBase
+}
+
+func newStorageV1BinlogSchemaMismatch() *StorageV1BinlogSchemaMismatch {
+	return &StorageV1BinlogSchemaMismatch{
+		checkItemBase: checkItemBase{
+			name:        "STORAGE_V1_BINLOG_SCHEMA_MISMATCH",
+			description: `Checks whether storage v1 segments (storage version 0 or 1) have insert binlogs for every field in the collection schema.`,
+		},
+	}
+}
+
+func (i *StorageV1BinlogSchemaMismatch) Check(ctx context.Context, client metakv.MetaKV, basePath string) ([]*HealthzCheckReport, error) {
+	collections, err := common.ListCollections(ctx, client, basePath)
+	if err != nil {
+		return nil, err
+	}
+
+	collectionByID := make(map[int64]*models.Collection, len(collections))
+	for _, collection := range collections {
+		collectionByID[collection.GetProto().GetID()] = collection
+	}
+
+	segments, err := common.ListSegments(ctx, client, basePath)
+	if err != nil {
+		return nil, err
+	}
+
+	var results []*HealthzCheckReport
+	for _, segment := range segments {
+		if segment.GetStorageVersion() != 0 && segment.GetStorageVersion() != 1 {
+			continue
+		}
+
+		collection, ok := collectionByID[segment.GetCollectionID()]
+		if !ok {
+			continue
+		}
+
+		schemaFieldIDs := make(map[int64]struct{})
+		for _, field := range collection.GetProto().GetSchema().GetFields() {
+			schemaFieldIDs[field.GetFieldID()] = struct{}{}
+		}
+
+		binlogFieldIDs := make(map[int64]struct{})
+		for _, fieldBinlog := range segment.GetBinlogs() {
+			binlogFieldIDs[fieldBinlog.FieldID] = struct{}{}
+		}
+
+		missingFieldIDs := missingIDs(schemaFieldIDs, binlogFieldIDs)
+		if len(missingFieldIDs) == 0 {
+			continue
+		}
+
+		results = append(results, &HealthzCheckReport{
+			Item: i.Name(),
+			Msg:  fmt.Sprintf("Segment %d storage v1 binlogs are missing collection schema fields", segment.GetID()),
+			Extra: map[string]any{
+				"segment_id":        segment.GetID(),
+				"collection_id":     segment.GetCollectionID(),
+				"storage_version":   segment.GetStorageVersion(),
+				"missing_field_ids": missingFieldIDs,
+				"schema_field_ids":  sortedIDs(schemaFieldIDs),
+				"binlog_field_ids":  sortedIDs(binlogFieldIDs),
+			},
+		})
+	}
+
+	return results, nil
+}
+
+func missingIDs(expected, actual map[int64]struct{}) []int64 {
+	missing := make([]int64, 0)
+	for id := range expected {
+		if _, ok := actual[id]; !ok {
+			missing = append(missing, id)
+		}
+	}
+	sort.Slice(missing, func(i, j int) bool { return missing[i] < missing[j] })
+	return missing
+}
+
+func sortedIDs(ids map[int64]struct{}) []int64 {
+	result := make([]int64, 0, len(ids))
+	for id := range ids {
+		result = append(result, id)
+	}
+	sort.Slice(result, func(i, j int) bool { return result[i] < result[j] })
+	return result
+}

states/healthz/storage_v1_binlog_schema_mismatch_test.go

@@ -0,0 +1,119 @@
+package healthz
+
+import (
+	"context"
+	"fmt"
+	"path"
+	"sort"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+
+	"github.com/milvus-io/birdwatcher/states/etcd/common"
+	metakv "github.com/milvus-io/birdwatcher/states/kv"
+	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
+	"github.com/milvus-io/milvus/pkg/v2/proto/etcdpb"
+)
+
+type healthzCheckKV struct {
+	metakv.MetaKV
+	data map[string]string
+}
+
+func (k *healthzCheckKV) LoadWithPrefix(_ context.Context, prefix string, _ ...metakv.LoadOption) ([]string, []string, error) {
+	keys := make([]string, 0)
+	for key := range k.data {
+		if strings.HasPrefix(key, prefix) {
+			keys = append(keys, key)
+		}
+	}
+	sort.Strings(keys)
+
+	values := make([]string, 0, len(keys))
+	for _, key := range keys {
+		values = append(values, k.data[key])
+	}
+	return keys, values, nil
+}
+
+func TestStorageV1BinlogSchemaMismatchRegistered(t *testing.T) {
+	item, ok := GetHealthzCheckItem("STORAGE_V1_BINLOG_SCHEMA_MISMATCH")
+	require.True(t, ok)
+	require.Equal(t, "STORAGE_V1_BINLOG_SCHEMA_MISMATCH", item.Name())
+	require.NotEmpty(t, item.Description())
+}
+
+func TestStorageV1BinlogSchemaMismatchCheck(t *testing.T) {
+	const (
+		basePath     = "root"
+		collectionID = int64(100)
+		partitionID  = int64(10)
+	)
+
+	cli := &healthzCheckKV{data: map[string]string{
+		path.Join(basePath, common.CollectionMetaPrefix, "100"): healthzProtoValue(t, &etcdpb.CollectionInfo{
+			ID:     collectionID,
+			Schema: &schemapb.CollectionSchema{Name: "test"},
+		}),
+		path.Join(basePath, common.FieldMetaPrefix, "100", "100"): healthzProtoValue(t, &schemapb.FieldSchema{FieldID: 100, Name: "pk"}),
+		path.Join(basePath, common.FieldMetaPrefix, "100", "101"): healthzProtoValue(t, &schemapb.FieldSchema{FieldID: 101, Name: "value"}),
+		path.Join(basePath, common.FieldMetaPrefix, "100", "102"): healthzProtoValue(t, &schemapb.FieldSchema{FieldID: 102, Name: "vector"}),
+
+		segmentKey(basePath, collectionID, partitionID, 1):     healthzProtoValue(t, &datapb.SegmentInfo{ID: 1, CollectionID: collectionID, PartitionID: partitionID, StorageVersion: 0}),
+		binlogKey(basePath, collectionID, partitionID, 1, 0):   healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 0}),
+		binlogKey(basePath, collectionID, partitionID, 1, 1):   healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 1}),
+		binlogKey(basePath, collectionID, partitionID, 1, 100): healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 100}),
+		binlogKey(basePath, collectionID, partitionID, 1, 102): healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 102}),
+		binlogKey(basePath, collectionID, partitionID, 1, 999): healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 999}),
+
+		segmentKey(basePath, collectionID, partitionID, 2):     healthzProtoValue(t, &datapb.SegmentInfo{ID: 2, CollectionID: collectionID, PartitionID: partitionID, StorageVersion: 1}),
+		binlogKey(basePath, collectionID, partitionID, 2, 100): healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 100}),
+		binlogKey(basePath, collectionID, partitionID, 2, 101): healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 101}),
+		binlogKey(basePath, collectionID, partitionID, 2, 999): healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 999}),
+
+		segmentKey(basePath, collectionID, partitionID, 3):     healthzProtoValue(t, &datapb.SegmentInfo{ID: 3, CollectionID: collectionID, PartitionID: partitionID, StorageVersion: 2}),
+		binlogKey(basePath, collectionID, partitionID, 3, 100): healthzProtoValue(t, &datapb.FieldBinlog{FieldID: 100}),
+	}}
+
+	item := newStorageV1BinlogSchemaMismatch()
+	reports, err := item.Check(context.Background(), cli, basePath)
+	require.NoError(t, err)
+	require.Len(t, reports, 2)
+
+	require.Equal(t, item.Name(), reports[0].Item)
+	require.EqualValues(t, 1, reports[0].Extra["segment_id"])
+	require.EqualValues(t, collectionID, reports[0].Extra["collection_id"])
+	require.EqualValues(t, 0, reports[0].Extra["storage_version"])
+	require.Equal(t, []int64{101}, reports[0].Extra["missing_field_ids"])
+	require.Equal(t, []int64{100, 101, 102}, reports[0].Extra["schema_field_ids"])
+	require.Equal(t, []int64{0, 1, 100, 102, 999}, reports[0].Extra["binlog_field_ids"])
+
+	require.Equal(t, item.Name(), reports[1].Item)
+	require.EqualValues(t, 2, reports[1].Extra["segment_id"])
+	require.EqualValues(t, collectionID, reports[1].Extra["collection_id"])
+	require.EqualValues(t, 1, reports[1].Extra["storage_version"])
+	require.Equal(t, []int64{102}, reports[1].Extra["missing_field_ids"])
+	require.Equal(t, []int64{100, 101, 999}, reports[1].Extra["binlog_field_ids"])
+}
+
+func segmentKey(basePath string, collectionID, partitionID, segmentID int64) string {
+	return path.Join(basePath, common.DCPrefix, common.SegmentMetaPrefix, formatID(collectionID), formatID(partitionID), formatID(segmentID))
+}
+
+func binlogKey(basePath string, collectionID, partitionID, segmentID, fieldID int64) string {
+	return path.Join(basePath, common.DCPrefix, "binlog", formatID(collectionID), formatID(partitionID), formatID(segmentID), formatID(fieldID))
+}
+
+func formatID(id int64) string {
+	return fmt.Sprintf("%d", id)
+}
+
+func healthzProtoValue(t *testing.T, value proto.Message) string {
+	t.Helper()
+	bs, err := proto.Marshal(value)
+	require.NoError(t, err)
+	return string(bs)
+}