From f63cf7c51426a043b0cfed5538fb2713c479a42d Mon Sep 17 00:00:00 2001
From: PhilippPlotnikov <70917163+PhilippDeveloper@users.noreply.github.com>
Date: Fri, 11 Mar 2022 16:23:53 +0500
Subject: [PATCH] Implement traefil support

---
 rollout/trafficrouting/traefik/traefik.go | 83 ++++++++++++++++++++---
 1 file changed, 72 insertions(+), 11 deletions(-)

diff --git a/rollout/trafficrouting/traefik/traefik.go b/rollout/trafficrouting/traefik/traefik.go
index 1af8c24d0b..6488af84ff 100644
--- a/rollout/trafficrouting/traefik/traefik.go
+++ b/rollout/trafficrouting/traefik/traefik.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"strings"
 
+	"github.com/pkg/errors"
+
 	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
 	"github.com/argoproj/argo-rollouts/utils/defaults"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -71,36 +73,95 @@ func (r *Reconciler) SetWeight(desiredWeight int32, additionalDestinations ...v1
 	rollout := r.Rollout
 	traefikServiceName := rollout.Spec.Strategy.Canary.TrafficRouting.Traefik.Service
 	traefikService, err := r.Client.Get(ctx, traefikServiceName, metav1.GetOptions{})
-	desiredTraefikService := traefikService.DeepCopy()
-	canaryServiceName := rollout.Spec.Strategy.Canary.CanaryService
 	if err != nil {
 		return err
 	}
+	canaryServiceName := rollout.Spec.Strategy.Canary.CanaryService
+	stableServiceName := rollout.Spec.Strategy.Canary.StableService
 	services, isFound, err := unstructured.NestedSlice(traefikService.Object, "spec", "weighted", "services")
 	if err != nil || !isFound {
 		return err
 	}
 	for _, service := range services {
-		serviceTest, ok := service.(map[string]interface{})
+		typedService, ok := service.(map[string]interface{})
 		if !ok {
-			continue
+			return errors.New("Failed type assertion setting weight for traefik service")
+		}
+		serviceName, isFound, err := unstructured.NestedString(typedService, "name")
+		if err != nil || !isFound {
+			return err
 		}
-		serviceName, _, _ := unstructured.NestedString(serviceTest, "name")
 		if serviceName == canaryServiceName {
-			unstructured.SetNestedField(serviceTest, int64(desiredWeight), "weight")
-			break
+			err := unstructured.SetNestedField(typedService, int64(desiredWeight), "weight")
+			if err != nil {
+				return err
+			}
+			continue
+		}
+		if serviceName == stableServiceName {
+			err := unstructured.SetNestedField(typedService, int64(100-desiredWeight), "weight")
+			if err != nil {
+				return err
+			}
+			continue
 		}
 	}
-	err = unstructured.SetNestedSlice(desiredTraefikService.Object, services, "spec", "weighted", "services")
+	err = unstructured.SetNestedSlice(traefikService.Object, services, "spec", "weighted", "services")
 	if err != nil {
 		return err
 	}
-	r.Client.Create(ctx, desiredTraefikService, metav1.CreateOptions{})
-	return nil
+	_, err = r.Client.Update(ctx, traefikService, metav1.UpdateOptions{})
+	return err
 }
 
 func (r *Reconciler) VerifyWeight(desiredWeight int32, additionalDestinations ...v1alpha1.WeightDestination) (*bool, error) {
-	return nil, nil
+	ctx := context.TODO()
+	verifyingStatus := false
+	rollout := r.Rollout
+	traefikServiceName := rollout.Spec.Strategy.Canary.TrafficRouting.Traefik.Service
+	traefikService, err := r.Client.Get(ctx, traefikServiceName, metav1.GetOptions{})
+	if err != nil {
+		return &verifyingStatus, err
+	}
+	canaryServiceName := rollout.Spec.Strategy.Canary.CanaryService
+	stableServiceName := rollout.Spec.Strategy.Canary.StableService
+	services, isFound, err := unstructured.NestedSlice(traefikService.Object, "spec", "weighted", "services")
+	if err != nil || !isFound {
+		return &verifyingStatus, err
+	}
+	for _, service := range services {
+		typedService, ok := service.(map[string]interface{})
+		if !ok {
+			return &verifyingStatus, errors.New("Failed type assertion setting weight for traefik service")
+		}
+		serviceName, isFound, err := unstructured.NestedString(typedService, "name")
+		if err != nil || !isFound {
+			return &verifyingStatus, err
+		}
+		if serviceName == canaryServiceName {
+			verifyingStatus = false
+			weight, isFound, err := unstructured.NestedInt64(typedService, "weight")
+			if err != nil || !isFound {
+				return &verifyingStatus, err
+			}
+			verifyingStatus = weight == int64(desiredWeight)
+			if !verifyingStatus {
+				return &verifyingStatus, errors.New("Traefik service weight for canary service is not right")
+			}
+		}
+		if serviceName == stableServiceName {
+			verifyingStatus = false
+			weight, isFound, err := unstructured.NestedInt64(typedService, "weight")
+			if err != nil || !isFound {
+				return &verifyingStatus, err
+			}
+			verifyingStatus = weight == int64(100-desiredWeight)
+			if !verifyingStatus {
+				return &verifyingStatus, errors.New("Traefik service weight for canary service is not right")
+			}
+		}
+	}
+	return &verifyingStatus, nil
 }
 
 func (r *Reconciler) Type() string {