From d920e5c91583388bc71ee36998b4343d7a3d216d Mon Sep 17 00:00:00 2001
From: XuanYang-cn <xuan.yang@zilliz.com>
Date: Tue, 31 May 2022 18:02:03 +0800
Subject: [PATCH] Fix DataNode panic while allocating IDs (#17294)

See also: #17270

Signed-off-by: yangxuan <xuan.yang@zilliz.com>
---
 internal/datanode/allocator.go      | 18 +++++++++---------
 internal/datanode/allocator_test.go |  6 +++++-
 internal/datanode/mock_test.go      |  3 +--
 3 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/internal/datanode/allocator.go b/internal/datanode/allocator.go
index 989856f90914a..dc73fc212246d 100644
--- a/internal/datanode/allocator.go
+++ b/internal/datanode/allocator.go
@@ -58,15 +58,14 @@ func (alloc *allocator) allocID() (UniqueID, error) {
 		},
 		Count: 1,
 	})
-
-	if resp.Status.ErrorCode != commonpb.ErrorCode_Success {
-		return 0, errors.New(resp.Status.GetReason())
-	}
-
 	if err != nil {
 		return 0, err
 	}
 
+	if resp.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
+		return 0, errors.New(resp.GetStatus().GetReason())
+	}
+
 	return resp.ID, nil
 }
 
@@ -81,13 +80,14 @@ func (alloc *allocator) allocIDBatch(count uint32) (UniqueID, uint32, error) {
 		Count: count,
 	})
 
-	if resp.Status.ErrorCode != commonpb.ErrorCode_Success {
-		return 0, 0, errors.New(resp.Status.GetReason())
-	}
-
 	if err != nil {
 		return 0, 0, err
 	}
+
+	if resp.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
+		return 0, 0, errors.New(resp.GetStatus().GetReason())
+	}
+
 	return resp.GetID(), resp.GetCount(), nil
 }
 
diff --git a/internal/datanode/allocator_test.go b/internal/datanode/allocator_test.go
index c13335a2d477e..e69d41e0f8853 100644
--- a/internal/datanode/allocator_test.go
+++ b/internal/datanode/allocator_test.go
@@ -30,11 +30,15 @@ func TestAllocator_Basic(t *testing.T) {
 		ms.setID(666)
 		_, err := allocator.allocID()
 		assert.NoError(t, err)
+
+		ms.setID(-1)
+		_, err = allocator.allocID()
+		assert.Error(t, err)
 	})
 
 	t.Run("Test alloc ID batch", func(t *testing.T) {
 		// If id == 0, AllocID will return not successful status
-		// If id == -1, AllocID will return err
+		// If id == -1, AllocID will return err with nil status
 		ms.setID(666)
 		_, count, err := allocator.allocIDBatch(10)
 		assert.NoError(t, err)
diff --git a/internal/datanode/mock_test.go b/internal/datanode/mock_test.go
index 2f4820a28b9af..9ab3be0e50246 100644
--- a/internal/datanode/mock_test.go
+++ b/internal/datanode/mock_test.go
@@ -864,8 +864,7 @@ func (m *RootCoordFactory) AllocID(ctx context.Context, in *rootcoordpb.AllocIDR
 	}
 
 	if m.ID == -1 {
-		resp.Status.ErrorCode = commonpb.ErrorCode_Success
-		return resp, errors.New(resp.Status.GetReason())
+		return nil, errors.New(resp.Status.GetReason())
 	}
 
 	resp.ID = m.ID