[Misc] enhance headless service sync with update logic

Signed-off-by: Omer Aplatony <omerap12@gmail.com>

Author: omerap12

pkg/controller/stormservice/sync.go

@@ -24,6 +24,7 @@ import (
 	apps "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	apiequality "k8s.io/apimachinery/pkg/api/equality"
+	errors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/klog/v2"
 	"k8s.io/utils/integer"
@@ -84,16 +85,29 @@ func (r *StormServiceReconciler) syncHeadlessService(ctx context.Context, servic
 	}
 
 	headlessService := &corev1.Service{}
-	if err := r.Client.Get(ctx, client.ObjectKey{Name: service.Name, Namespace: service.Namespace}, headlessService); client.IgnoreNotFound(err) != nil {
-		return err
-	} else if err != nil {
-		// not found pg, need to create
-		if err = r.Client.Create(ctx, expectedService); err == nil {
-			r.EventRecorder.Eventf(service, corev1.EventTypeNormal, HeadlessServiceEventType, "Headless Service(discovery) %s synced", service.Name)
+	err := r.Client.Get(ctx, client.ObjectKey{Name: service.Name, Namespace: service.Namespace}, headlessService)
+	if err != nil {
+		if errors.IsNotFound(err) {
+			// service doesn't exist, create it
+			if createErr := r.Client.Create(ctx, expectedService); createErr != nil {
+				return fmt.Errorf("failed to create headless service: %w", createErr)
+			}
+			r.EventRecorder.Eventf(service, corev1.EventTypeNormal, HeadlessServiceEventType, "Headless Service(discovery) %s created", service.Name)
+			return nil
 		}
-		return err
+		return err // Return other errors immediately
 	}
 
+	if !isServiceEqual(headlessService, expectedService) {
+		headlessService.Spec = expectedService.Spec
+		if err := r.Client.Update(ctx, headlessService); err != nil {
+			return fmt.Errorf("failed to update headless service: %w", err)
+		}
+		r.EventRecorder.Eventf(service, corev1.EventTypeNormal, "Updated", "Headless Service updated")
+	}
+
+	r.EventRecorder.Eventf(service, corev1.EventTypeNormal, HeadlessServiceEventType, "Headless Service(discovery) %s updated", service.Name)
+
 	return nil
 }
 

pkg/controller/stormservice/sync_test.go

@@ -16,7 +16,21 @@ limitations under the License.
 
 package stormservice
 
-import "testing"
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/tools/record"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+
+	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
+	"github.com/vllm-project/aibrix/pkg/controller/constants"
+)
 
 func TestCalculateReplicas(t *testing.T) {
 	type args struct {
@@ -112,3 +126,104 @@ func TestCalculateReplicas(t *testing.T) {
 		}
 	}
 }
+
+func TestSyncHeadlessService(t *testing.T) {
+	scheme := runtime.NewScheme()
+	_ = corev1.AddToScheme(scheme)
+	_ = orchestrationv1alpha1.AddToScheme(scheme)
+
+	tests := []struct {
+		name            string
+		stormService    *orchestrationv1alpha1.StormService
+		existingService *corev1.Service
+		wantError       bool
+	}{
+		{
+			name: "create new headless service",
+			stormService: &orchestrationv1alpha1.StormService{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-storm",
+					Namespace: "default",
+					Labels: map[string]string{
+						"app": "test",
+					},
+				},
+			},
+			existingService: nil,
+			wantError:       false,
+		},
+		{
+			name: "service already exists",
+			stormService: &orchestrationv1alpha1.StormService{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-storm",
+					Namespace: "default",
+				},
+			},
+			existingService: &corev1.Service{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-storm",
+					Namespace: "default",
+				},
+				Spec: corev1.ServiceSpec{
+					Type:      corev1.ServiceTypeClusterIP,
+					ClusterIP: corev1.ClusterIPNone,
+					Selector:  map[string]string{}, // empty selector that should be updated
+				},
+			},
+			wantError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var objs []client.Object
+			if tt.existingService != nil {
+				objs = append(objs, tt.existingService)
+			}
+
+			fakeClient := fake.NewClientBuilder().
+				WithScheme(scheme).
+				WithObjects(objs...).
+				Build()
+
+			r := &StormServiceReconciler{
+				Client:        fakeClient,
+				EventRecorder: &record.FakeRecorder{},
+			}
+
+			err := r.syncHeadlessService(context.TODO(), tt.stormService)
+
+			if (err != nil) != tt.wantError {
+				t.Errorf("syncHeadlessService() error = %v, wantError %v", err, tt.wantError)
+				return
+			}
+
+			// Check if service was created/updated
+			service := &corev1.Service{}
+			err = fakeClient.Get(context.TODO(), client.ObjectKey{
+				Name:      tt.stormService.Name,
+				Namespace: tt.stormService.Namespace,
+			}, service)
+
+			if err != nil {
+				t.Errorf("Failed to get service: %v", err)
+				return
+			}
+
+			// Verify service properties
+			if service.Spec.ClusterIP != corev1.ClusterIPNone {
+				t.Errorf("Expected ClusterIP to be None, got %s", service.Spec.ClusterIP)
+			}
+
+			expectedSelector := map[string]string{constants.StormServiceNameLabelKey: tt.stormService.Name}
+			if !reflect.DeepEqual(service.Spec.Selector, expectedSelector) {
+				t.Errorf("Expected selector %v, got %v", expectedSelector, service.Spec.Selector)
+			}
+
+			if service.Spec.Type != corev1.ServiceTypeClusterIP {
+				t.Errorf("Expected service type ClusterIP, got %v", service.Spec.Type)
+			}
+		})
+	}
+}

pkg/controller/stormservice/utils.go

@@ -19,6 +19,8 @@ package stormservice
 import (
 	"sort"
 
+	corev1 "k8s.io/api/core/v1"
+	apiequality "k8s.io/apimachinery/pkg/api/equality"
 	intstrutil "k8s.io/apimachinery/pkg/util/intstr"
 
 	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
@@ -224,3 +226,10 @@ func sortRoleSetByRevision(roleSets []*orchestrationv1alpha1.RoleSet, updatedRev
 		}
 	})
 }
+
+// isServiceEqual compares two Kubernetes Service objects for equality
+func isServiceEqual(a, b *corev1.Service) bool {
+	return a.Spec.Type == b.Spec.Type &&
+		apiequality.Semantic.DeepEqual(a.Spec.Selector, b.Spec.Selector) &&
+		a.Spec.ClusterIP == b.Spec.ClusterIP
+}