diff --git a/pkg/controllers/ragengine_status.go b/pkg/controllers/ragengine_status.go index 1ebc088c9..fb5ae2463 100644 --- a/pkg/controllers/ragengine_status.go +++ b/pkg/controllers/ragengine_status.go @@ -7,39 +7,12 @@ import ( "context" kaitov1alpha1 "github.com/azure/kaito/api/v1alpha1" - "k8s.io/apimachinery/pkg/api/errors" - apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/client-go/util/retry" "k8s.io/klog/v2" "sigs.k8s.io/controller-runtime/pkg/client" ) -func (c *RAGEngineReconciler) updateRAGEngineStatus(ctx context.Context, name *client.ObjectKey, condition *metav1.Condition, workerNodes []string) error { - return retry.OnError(retry.DefaultRetry, - func(err error) bool { - return apierrors.IsServiceUnavailable(err) || apierrors.IsServerTimeout(err) || apierrors.IsTooManyRequests(err) - }, - func() error { - // Read the latest version to avoid update conflict. - ragObj := &kaitov1alpha1.RAGEngine{} - if err := c.Client.Get(ctx, *name, ragObj); err != nil { - if !errors.IsNotFound(err) { - return err - } - return nil - } - if condition != nil { - meta.SetStatusCondition(&ragObj.Status.Conditions, *condition) - } - if workerNodes != nil { - ragObj.Status.WorkerNodes = workerNodes - } - return c.Client.Status().Update(ctx, ragObj) - }) -} - func (c *RAGEngineReconciler) updateStatusConditionIfNotMatch(ctx context.Context, ragObj *kaitov1alpha1.RAGEngine, cType kaitov1alpha1.ConditionType, cStatus metav1.ConditionStatus, cReason, cMessage string) error { if curCondition := meta.FindStatusCondition(ragObj.Status.Conditions, string(cType)); curCondition != nil { @@ -56,5 +29,5 @@ func (c *RAGEngineReconciler) updateStatusConditionIfNotMatch(ctx context.Contex ObservedGeneration: ragObj.GetGeneration(), Message: cMessage, } - return c.updateRAGEngineStatus(ctx, &client.ObjectKey{Name: ragObj.Name, Namespace: ragObj.Namespace}, &cObj, nil) + return updateObjStatus(ctx, c.Client, &client.ObjectKey{Name: ragObj.Name, Namespace: ragObj.Namespace}, "RAGEngine", &cObj, nil) } diff --git a/pkg/controllers/ragengine_status_test.go b/pkg/controllers/ragengine_status_test.go index 00a7bc447..1440998c0 100644 --- a/pkg/controllers/ragengine_status_test.go +++ b/pkg/controllers/ragengine_status_test.go @@ -38,7 +38,7 @@ func TestUpdateRAGEngineStatus(t *testing.T) { mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.RAGEngine{}), mock.Anything).Return(nil) mockClient.StatusMock.On("Update", mock.IsType(context.Background()), mock.IsType(&kaitov1alpha1.RAGEngine{}), mock.Anything).Return(nil) - err := reconciler.updateRAGEngineStatus(ctx, &client.ObjectKey{Name: ragengine.Name, Namespace: ragengine.Namespace}, &condition, workerNodes) + err := updateObjStatus(ctx, reconciler.Client, &client.ObjectKey{Name: ragengine.Name, Namespace: ragengine.Namespace}, "RAGEngine", &condition, workerNodes) assert.Nil(t, err) }) @@ -60,7 +60,7 @@ func TestUpdateRAGEngineStatus(t *testing.T) { mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.RAGEngine{}), mock.Anything).Return(errors.New("Get operation failed")) - err := reconciler.updateRAGEngineStatus(ctx, &client.ObjectKey{Name: ragengine.Name, Namespace: ragengine.Namespace}, &condition, workerNodes) + err := updateObjStatus(ctx, reconciler.Client, &client.ObjectKey{Name: ragengine.Name, Namespace: ragengine.Namespace}, "RAGEngine", &condition, workerNodes) assert.NotNil(t, err) }) @@ -82,7 +82,7 @@ func TestUpdateRAGEngineStatus(t *testing.T) { mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.RAGEngine{}), mock.Anything).Return(apierrors.NewNotFound(schema.GroupResource{}, "ragengine")) - err := reconciler.updateRAGEngineStatus(ctx, &client.ObjectKey{Name: ragengine.Name, Namespace: ragengine.Namespace}, &condition, workerNodes) + err := updateObjStatus(ctx, reconciler.Client, &client.ObjectKey{Name: ragengine.Name, Namespace: ragengine.Namespace}, "RAGEngine", &condition, workerNodes) assert.Nil(t, err) }) } diff --git a/pkg/controllers/workspace_status.go b/pkg/controllers/workspace_status.go index b972ceaa9..f79972d89 100644 --- a/pkg/controllers/workspace_status.go +++ b/pkg/controllers/workspace_status.go @@ -5,6 +5,7 @@ package controllers import ( "context" + "fmt" "reflect" "sort" @@ -20,27 +21,43 @@ import ( "sigs.k8s.io/controller-runtime/pkg/client" ) -func (c *WorkspaceReconciler) updateWorkspaceStatus(ctx context.Context, name *client.ObjectKey, condition *metav1.Condition, workerNodes []string) error { +func updateObjStatus(ctx context.Context, c client.Client, name *client.ObjectKey, objType string, condition *metav1.Condition, workerNodes []string) error { return retry.OnError(retry.DefaultRetry, func(err error) bool { return apierrors.IsServiceUnavailable(err) || apierrors.IsServerTimeout(err) || apierrors.IsTooManyRequests(err) }, func() error { - // Read the latest version to avoid update conflict. - wObj := &kaitov1alpha1.Workspace{} - if err := c.Client.Get(ctx, *name, wObj); err != nil { + var obj client.Object + var conditions *[]metav1.Condition + var workerNodesField *[]string + switch objType { + case "Workspace": + ragObj := &kaitov1alpha1.Workspace{} + obj = ragObj + conditions = &ragObj.Status.Conditions + workerNodesField = &ragObj.Status.WorkerNodes + case "RAGEngine": + wObj := &kaitov1alpha1.RAGEngine{} + obj = wObj + conditions = &wObj.Status.Conditions + workerNodesField = &wObj.Status.WorkerNodes + default: + return fmt.Errorf("unsupported object type: %s", objType) + } + + if err := c.Get(ctx, *name, obj); err != nil { if !errors.IsNotFound(err) { return err } return nil } if condition != nil { - meta.SetStatusCondition(&wObj.Status.Conditions, *condition) + meta.SetStatusCondition(conditions, *condition) } if workerNodes != nil { - wObj.Status.WorkerNodes = workerNodes + *workerNodesField = workerNodes } - return c.Client.Status().Update(ctx, wObj) + return c.Status().Update(ctx, obj) }) } @@ -60,7 +77,7 @@ func (c *WorkspaceReconciler) updateStatusConditionIfNotMatch(ctx context.Contex ObservedGeneration: wObj.GetGeneration(), Message: cMessage, } - return c.updateWorkspaceStatus(ctx, &client.ObjectKey{Name: wObj.Name, Namespace: wObj.Namespace}, &cObj, nil) + return updateObjStatus(ctx, c.Client, &client.ObjectKey{Name: wObj.Name, Namespace: wObj.Namespace}, "Workspace", &cObj, nil) } func (c *WorkspaceReconciler) updateStatusNodeListIfNotMatch(ctx context.Context, wObj *kaitov1alpha1.Workspace, validNodeList []*corev1.Node) error { @@ -73,5 +90,5 @@ func (c *WorkspaceReconciler) updateStatusNodeListIfNotMatch(ctx context.Context return nil } klog.InfoS("updateStatusNodeList", "workspace", klog.KObj(wObj)) - return c.updateWorkspaceStatus(ctx, &client.ObjectKey{Name: wObj.Name, Namespace: wObj.Namespace}, nil, nodeNameList) + return updateObjStatus(ctx, c.Client, &client.ObjectKey{Name: wObj.Name, Namespace: wObj.Namespace}, "Workspace", nil, nodeNameList) } diff --git a/pkg/controllers/workspace_status_test.go b/pkg/controllers/workspace_status_test.go index d90d2b603..b2adbf864 100644 --- a/pkg/controllers/workspace_status_test.go +++ b/pkg/controllers/workspace_status_test.go @@ -38,7 +38,7 @@ func TestUpdateWorkspaceStatus(t *testing.T) { mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(nil) mockClient.StatusMock.On("Update", mock.IsType(context.Background()), mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(nil) - err := reconciler.updateWorkspaceStatus(ctx, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, &condition, workerNodes) + err := updateObjStatus(ctx, reconciler.Client, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, "Workspace", &condition, workerNodes) assert.Nil(t, err) }) @@ -60,7 +60,7 @@ func TestUpdateWorkspaceStatus(t *testing.T) { mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(errors.New("Get operation failed")) - err := reconciler.updateWorkspaceStatus(ctx, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, &condition, workerNodes) + err := updateObjStatus(ctx, reconciler.Client, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, "Workspace", &condition, workerNodes) assert.NotNil(t, err) }) @@ -82,7 +82,7 @@ func TestUpdateWorkspaceStatus(t *testing.T) { mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(apierrors.NewNotFound(schema.GroupResource{}, "workspace")) - err := reconciler.updateWorkspaceStatus(ctx, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, &condition, workerNodes) + err := updateObjStatus(ctx, reconciler.Client, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, "Workspace", &condition, workerNodes) assert.Nil(t, err) }) }