From d9f00faeec1b2b4f9bd8c4459e73c09543430ee2 Mon Sep 17 00:00:00 2001
From: jaime <yun.zhang@zilliz.com>
Date: Tue, 30 Jan 2024 15:57:31 +0800
Subject: [PATCH] Fix probe fail and exit while probe fail

Signed-off-by: jaime <yun.zhang@zilliz.com>
---
 states/etcd/common/collection.go | 12 ++++--------
 states/probe.go                  | 18 ++++++++++++++++--
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/states/etcd/common/collection.go b/states/etcd/common/collection.go
index f63e3a5..bf966ca 100644
--- a/states/etcd/common/collection.go
+++ b/states/etcd/common/collection.go
@@ -123,11 +123,7 @@ func GetCollectionByIDVersion(ctx context.Context, cli kv.MetaKV, basePath strin
 	// meta before database
 	prefix := path.Join(basePath, CollectionMetaPrefix, strconv.FormatInt(collID, 10))
 	val, err := cli.Load(ctx, prefix)
-	if err != nil {
-		fmt.Println("get error", err.Error())
-		return nil, err
-	}
-	if len(val) > 0 {
+	if err == nil && len(val) > 0 {
 		found = true
 		ck = prefix
 		cv = []byte(val)
@@ -135,16 +131,16 @@ func GetCollectionByIDVersion(ctx context.Context, cli kv.MetaKV, basePath strin
 
 	// with database, dbID unknown here
 	prefix = path.Join(basePath, DBCollectionMetaPrefix)
-	keys, _, _ := cli.LoadWithPrefix(ctx, prefix)
+	keys, values, _ := cli.LoadWithPrefix(ctx, prefix)
 	suffix := strconv.FormatInt(collID, 10)
-	for _, key := range keys {
+	for i, key := range keys {
 		if strings.HasSuffix(key, suffix) {
 			if found {
 				return nil, fmt.Errorf("multiple key found for collection %d: %s, %s", collID, ck, key)
 			}
 			found = true
 			ck = prefix
-			cv = []byte(val)
+			cv = []byte(values[i])
 		}
 	}
 	if !found {
diff --git a/states/probe.go b/states/probe.go
index d7d6f61..8082397 100644
--- a/states/probe.go
+++ b/states/probe.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"math"
 	"math/rand"
+	"os"
 	"path"
 	"strconv"
 	"strings"
@@ -87,10 +88,12 @@ func getProbeQueryCmd(cli kv.MetaKV, basePath string) *cobra.Command {
 				return
 			}
 
+			errCount := 0
 			for _, collection := range loaded {
 				fmt.Println("probing collection", collection.CollectionID)
 				req, err := getMockSearchRequest(ctx, cli, basePath, collection)
 				if err != nil {
+					errCount++
 					fmt.Println("failed to generated mock request", err.Error())
 					continue
 				}
@@ -101,14 +104,20 @@ func getProbeQueryCmd(cli kv.MetaKV, basePath string) *cobra.Command {
 				})
 				if err != nil {
 					fmt.Println("querycoord get shard leaders error", err.Error())
+					errCount++
+					continue
+				}
+				if leaders.GetStatus().GetErrorCode() != commonpbv2.ErrorCode_Success {
+					fmt.Printf("collection[%d] failed to get shard leader, error: %s\n", collection.CollectionID, leaders.GetStatus().GetReason())
+					errCount++
 					continue
 				}
 
 				for _, shard := range leaders.GetShards() {
-
 					for _, nodeID := range shard.GetNodeIds() {
 						qn, ok := qns[nodeID]
 						if !ok {
+							errCount++
 							fmt.Printf("Shard leader %d not online\n", nodeID)
 							continue
 						}
@@ -120,17 +129,22 @@ func getProbeQueryCmd(cli kv.MetaKV, basePath string) *cobra.Command {
 						cancel()
 						if err != nil {
 							fmt.Printf("Shard %s Leader[%d] failed to search with eventually consistency level, err: %s\n", shard.GetChannelName(), nodeID, err.Error())
+							errCount++
 							continue
 						}
 						if resp.GetStatus().GetErrorCode() != commonpbv2.ErrorCode_Success {
 							fmt.Printf("Shard %s Leader[%d] failed to search,error code: %s reason:%s\n", shard.GetChannelName(), nodeID, resp.GetStatus().GetErrorCode().String(), resp.GetStatus().GetReason())
+							errCount++
 							continue
 						}
 						fmt.Printf("Shard %s leader[%d] probe with search success.\n", shard.GetChannelName(), nodeID)
 					}
 				}
 			}
-
+			if errCount != 0 {
+				fmt.Printf("probe failed, hit %d errors", errCount)
+				os.Exit(-1)
+			}
 		},
 	}