Add implementation on update / backfill for gke revenue labels in for…

…warding rules.

Update tests for the update / backfill

pkg/loadbalancers/fakes.go

@@ -51,3 +51,16 @@ func InsertForwardingRuleHook(ctx context.Context, key *meta.Key, obj *compute.F
 	}
 	return false, nil
 }
+
+func SetLabelsHook(ctx context.Context, key *meta.Key, req *compute.GlobalSetLabelsRequest, m *cloud.MockGlobalForwardingRules, options ...cloud.Option) error {
+	m.Lock.Lock()
+	defer m.Lock.Unlock()
+	var fr *compute.ForwardingRule
+	if obj, ok := m.Objects[*key]; ok {
+		fr = obj.ToGA()
+		fr.Labels = req.Labels
+		m.Objects[*key] = &cloud.MockGlobalForwardingRulesObj{Obj: fr}
+		return nil
+	}
+	return fmt.Errorf("Failed to find the forwarding rule with key %v", key)
+}

pkg/loadbalancers/forwarding_rules.go

@@ -50,6 +50,9 @@ const (
 	// addressAlreadyInUseMessageInternal is the error message string returned by the compute API
 	// when creating an internal forwarding rule that uses a conflicting IP address.
 	addressAlreadyInUseMessageInternal = "IP_IN_USE_BY_ANOTHER_RESOURCE"
+
+	// Label to be added to all forwarding rules to indicate the FR was created throughGKE
+	gkeLabel = "goog-gke-node"
 )
 
 func (l7 *L7) checkHttpForwardingRule() (err error) {
@@ -98,11 +101,13 @@ func (l7 *L7) checkForwardingRule(protocol namer.NamerProtocol, name, proxyLink,
 		return nil, err
 	}
 
+	fwLabels := map[string]string{gkeLabel: ""}
+
 	isL7ILB := utils.IsGCEL7ILBIngress(l7.runtimeInfo.Ingress)
 	isL7XLBRegional := utils.IsGCEL7XLBRegionalIngress(&l7.ingress)
 	tr := translator.NewTranslator(isL7ILB, isL7XLBRegional, l7.namer)
 	env := &translator.Env{VIP: ip, Network: l7.cloud.NetworkURL(), Subnetwork: l7.cloud.SubnetworkURL()}
-	fr := tr.ToCompositeForwardingRule(env, protocol, version, proxyLink, description, l7.runtimeInfo.StaticIPSubnet)
+	fr := tr.ToCompositeForwardingRule(env, protocol, version, proxyLink, description, l7.runtimeInfo.StaticIPSubnet, fwLabels)
 
 	existing, _ = composite.GetForwardingRule(l7.cloud, key, version, l7.logger)
 	if existing != nil && (fr.IPAddress != "" && existing.IPAddress != fr.IPAddress || existing.PortRange != fr.PortRange) {
@@ -162,6 +167,17 @@ func (l7 *L7) checkForwardingRule(protocol namer.NamerProtocol, name, proxyLink,
 			return nil, err
 		}
 	}
+	labels := existing.Labels
+	if labels == nil {
+		labels = make(map[string]string)
+	}
+	// Set GKE revenue label on the forwarding rule.
+	if _, ok := existing.Labels[gkeLabel]; !ok {
+		labels[gkeLabel] = ""
+		if err := composite.SetLabelsForwardingRule(l7.cloud, key, fr, labels, l7.logger); err != nil {
+			return nil, err
+		}
+	}
 	return existing, nil
 }
 

pkg/loadbalancers/loadbalancers_test.go

@@ -123,6 +123,7 @@ func newTestJig(t *testing.T) *testJig {
 		return nil
 	}
 	mockGCE.MockGlobalForwardingRules.InsertHook = InsertGlobalForwardingRuleHook
+	mockGCE.MockGlobalForwardingRules.SetLabelsHook = SetLabelsHook
 
 	ing := newIngress()
 	feNamer := namer_util.NewFrontendNamerFactory(namer, "", klog.TODO()).Namer(ing)
@@ -187,6 +188,72 @@ func TestCreateHTTPLoadBalancer(t *testing.T) {
 	verifyHTTPForwardingRuleAndProxyLinks(t, j, l7, "")
 }
 
+func TestUpdateHTTPLoadBalancer(t *testing.T) {
+	for _, tc := range []struct {
+		desc   string
+		labels map[string]string
+	}{
+		{
+			desc:   "No label added in the forwarding rule",
+			labels: nil,
+		},
+		{
+			desc: "Other label added in the forwarding rule",
+			labels: map[string]string{
+				"foo": "bar",
+			},
+		},
+		{
+			desc: "GKE label added in the forwarding rule",
+			labels: map[string]string{
+				"goog-gke-node": "",
+			},
+		},
+	} {
+		j := newTestJig(t)
+
+		key := &meta.Key{
+			Name:   "k8s-fw-namespace1-test--uid1",
+			Zone:   "",
+			Region: "",
+		}
+
+		fr := &composite.ForwardingRule{
+			Version:    "ga",
+			IPProtocol: "TCP",
+			Name:       "k8s-fw-namespace1-test--uid1",
+			PortRange:  "80-80",
+			Target:     "https://www.googleapis.com/compute/v1/projects/test-project/global/targetHttpProxies/k8s-tp-namespace1-test--uid1",
+			Labels:     tc.labels,
+		}
+
+		if err := composite.CreateForwardingRule(j.fakeGCE, key, fr, j.pool.logger); err != nil {
+			t.Fatalf("Failed to create forwarding rule, err: %v", err)
+		}
+
+		gceUrlMap := utils.NewGCEURLMap(klog.TODO())
+		gceUrlMap.DefaultBackend = &utils.ServicePort{NodePort: 31234, BackendNamer: j.namer}
+		lbInfo := &L7RuntimeInfo{
+			AllowHTTP: true,
+			UrlMap:    gceUrlMap,
+			Ingress:   newIngress(),
+		}
+		l7, err := j.pool.Ensure(lbInfo)
+		if err != nil || l7 == nil {
+			t.Fatalf("Expected l7 not created, err: %v", err)
+		}
+		fr = getForwardingRule(t, j, l7)
+		if _, ok := fr.Labels["goog-gke-node"]; !ok {
+			t.Errorf("desc: %q, no GKE revenue label found in test case", tc.desc)
+		}
+		for key := range tc.labels {
+			if _, ok := fr.Labels[key]; !ok {
+				t.Errorf("desc: %q, previous label %q not found", key, tc.desc)
+			}
+		}
+	}
+}
+
 func TestCreateHTTPILBLoadBalancer(t *testing.T) {
 	// This should NOT create the forwarding rule and target proxy
 	// associated with the HTTPS branch of this loadbalancer.
@@ -370,6 +437,13 @@ func verifyHTTPSForwardingRuleAndProxyLinks(t *testing.T, j *testJig, l7 *L7) {
 	if fws.Description == "" {
 		t.Errorf("fws.Description not set; expected it to be")
 	}
+	if len(fws.Labels) == 0 {
+		t.Errorf("fws.Labels are empty; expected gke label")
+	} else {
+		if _, ok := fws.Labels[gkeLabel]; !ok {
+			t.Errorf("fws.Labels are non-empty but does not contain the gke label")
+		}
+	}
 }
 
 func verifyHTTPForwardingRuleAndProxyLinks(t *testing.T, j *testJig, l7 *L7, ip string) {
@@ -407,6 +481,31 @@ func verifyHTTPForwardingRuleAndProxyLinks(t *testing.T, j *testJig, l7 *L7, ip
 			t.Fatalf("fws.IPAddress = %q, want %q", fws.IPAddress, ip)
 		}
 	}
+	if len(fws.Labels) == 0 {
+		t.Errorf("fws.Labels are empty; expected revenue label")
+	} else {
+		if _, ok := fws.Labels[gkeLabel]; !ok {
+			t.Errorf("fws.Labels are non-empty but does not contain the revenue label")
+		}
+	}
+}
+
+func getForwardingRule(t *testing.T, j *testJig, l7 *L7) *composite.ForwardingRule {
+	t.Helper()
+
+	versions := l7.Versions()
+	key, err := composite.CreateKey(j.fakeGCE, l7.namer.UrlMap(), l7.scope)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	fwName := l7.namer.ForwardingRule(namer_util.HTTPProtocol)
+	key.Name = fwName
+	fws, err := composite.GetForwardingRule(j.fakeGCE, key, versions.ForwardingRule, klog.TODO())
+	if err != nil {
+		t.Fatalf("j.fakeGCE.GetGlobalForwardingRule(%q) = _, %v, want nil", fwName, err)
+	}
+	return fws
 }
 
 // Tests that a certificate is created from the provided Key/Cert combo

pkg/translator/translator.go

@@ -262,7 +262,7 @@ const (
 )
 
 // ToCompositeForwardingRule returns a composite.ForwardingRule of type HTTP or HTTPS.
-func (t *Translator) ToCompositeForwardingRule(env *Env, protocol namer.NamerProtocol, version meta.Version, proxyLink, description, fwSubnet string) *composite.ForwardingRule {
+func (t *Translator) ToCompositeForwardingRule(env *Env, protocol namer.NamerProtocol, version meta.Version, proxyLink, description, fwSubnet string, labels map[string]string) *composite.ForwardingRule {
 	var portRange string
 	if protocol == namer.HTTPProtocol {
 		portRange = httpDefaultPortRange
@@ -278,6 +278,7 @@ func (t *Translator) ToCompositeForwardingRule(env *Env, protocol namer.NamerPro
 		IPProtocol:  "TCP",
 		Description: description,
 		Version:     version,
+		Labels:      labels,
 	}
 
 	if t.IsL7ILB {

pkg/translator/translator_test.go

@@ -388,6 +388,7 @@ func TestToForwardingRule(t *testing.T) {
 		isL7XLBRegional bool
 		protocol        namer_util.NamerProtocol
 		ipSubnet        string
+		labels          map[string]string
 		want            *composite.ForwardingRule
 	}{
 		{
@@ -500,13 +501,28 @@ func TestToForwardingRule(t *testing.T) {
 				Subnetwork:          "different-subnet",
 			},
 		},
+		{
+			desc:     "http-xlb with gke label",
+			protocol: namer_util.HTTPProtocol,
+			labels:   map[string]string{"goog-gke-node": ""},
+			want: &composite.ForwardingRule{
+				Name:        "foo-fr",
+				IPAddress:   vip,
+				Target:      proxyLink,
+				PortRange:   httpDefaultPortRange,
+				IPProtocol:  "TCP",
+				Description: description,
+				Version:     version,
+				Labels:      map[string]string{"goog-gke-node": ""},
+			},
+		},
 	}
 
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			tr := NewTranslator(tc.isL7ILB, tc.isL7XLBRegional, &testNamer{"foo"})
 			env := &Env{VIP: vip, Network: network, Subnetwork: subnetwork}
-			got := tr.ToCompositeForwardingRule(env, tc.protocol, version, proxyLink, description, tc.ipSubnet)
+			got := tr.ToCompositeForwardingRule(env, tc.protocol, version, proxyLink, description, tc.ipSubnet, tc.labels)
 			if diff := cmp.Diff(tc.want, got); diff != "" {
 				t.Fatalf("Got diff for ForwardingRule (-want +got):\n%s", diff)
 			}