Fix DataNode panic while allocating IDs (#17294)

See also: #17270
Signed-off-by: Nyangxuan <xuan.yang@zilliz.com>

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
 }
 

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)

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