Merge pull request #45 from k8gb-io/random

[RRR] random round robin
k8gb-io · Feb 6, 2023 · 2480915 · 2480915
2 parents e0fe69a + 1e8cd1c
commit 2480915
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 8 deletions.
diff --git a/service/wrr/wrr.go b/service/wrr/wrr.go
@@ -20,7 +20,11 @@ Generated by GoLic, for more details see: https://github.com/AbsaOSS/golic
 
 import (
 	"context"
+	"fmt"
 	"net"
+	"strings"
+
+	"github.com/coredns/coredns/plugin/loadbalance"
 
 	"github.com/AbsaOSS/k8s_crd/common/k8sctrl"
 	"github.com/AbsaOSS/k8s_crd/common/netutils"
@@ -53,6 +57,15 @@ func (wrr *WeightRoundRobin) ServeDNS(_ context.Context, w dns.ResponseWriter, r
 	clientIP = netutils.ExtractEdnsSubnet(r)
 	indexKey := netutils.StripClosingDot(state.QName())
 	var ep = k8sctrl.Resources.DNSEndpoint.Lookup(indexKey, clientIP)
+	// weights are not defined, labels doesnt exists
+	if len(ep.Labels) == 1 && strings.ToUpper(ep.Labels["strategy"]) == "ROUNDROBIN" {
+		if err := (&loadbalance.RoundRobinResponseWriter{ResponseWriter: w}).WriteMsg(r); err != nil {
+			return dns.RcodeServerFailure, fmt.Errorf("[random] %s", err)
+		}
+		_, ansIP, _ := netutils.ParseAnswerSection(r.Answer)
+		log.Infof("[random]%s", ansIP)
+		return dns.RcodeSuccess, nil
+	}
 	g, err := parseGroups(ep.Labels)
 	if err != nil {
 		log.Errorf("error parsing labels (%s)", err)
@@ -65,6 +78,7 @@ func (wrr *WeightRoundRobin) ServeDNS(_ context.Context, w dns.ResponseWriter, r
 	// or when DNSEndpoint is not properly adjusted)
 	_, ansIP, _ := netutils.ParseAnswerSection(r.Answer)
 	if !g.equalIPs(ansIP) {
+		log.Infof("[skipping]%s", ansIP)
 		log.Debugf("DNSEndpoint labels does not match with DNS answer. DNSEndpoint might not be initialised yet. ep: %v; answers: %v", g.asSlice(), ansIP)
 		return dns.RcodeSuccess, nil
 	}

diff --git a/service/wrr/wrr_test.go b/service/wrr/wrr_test.go
@@ -155,7 +155,9 @@ func TestWeightRoundRobin(t *testing.T) {
 					test.A("alpha.cloud.example.com.		300	IN	A			10.240.0.1"),
 				},
 			},
-			writer:        newFakeWriter(ctrl, func(w *mocks.MockResponseWriter) {}),
+			writer: newFakeWriter(ctrl, func(w *mocks.MockResponseWriter) {
+				w.EXPECT().WriteMsg(gomock.Any()).Return(nil).Times(1)
+			}),
 			expectedError: false,
 			lookup: func(indexKey string, clientIP net.IP) (result k8sctrl.LocalDNSEndpoint) {
 				return k8sctrl.LocalDNSEndpoint{
@@ -465,13 +467,14 @@ func TestWeightRoundRobin(t *testing.T) {
 		},
 	}
 
-	for _, unit := range tests {
-		t.Run(unit.name, func(t *testing.T) {
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
 			wrr := NewWeightRoundRobin()
-			k8sctrl.Resources.DNSEndpoint.Lookup = unit.lookup
-			code, err := wrr.ServeDNS(context.TODO(), unit.writer.w, unit.msg)
-			assert.Equal(t, unit.rcode, code)
-			assert.Equal(t, unit.expectedError, err != nil)
+			test.msg.Question = append(test.msg.Question, dns.Question{Qtype: dns.TypeA})
+			k8sctrl.Resources.DNSEndpoint.Lookup = test.lookup
+			code, err := wrr.ServeDNS(context.TODO(), test.writer.w, test.msg)
+			assert.Equal(t, test.rcode, code)
+			assert.Equal(t, test.expectedError, err != nil)
 		})
 
 	}

diff --git a/terratest/test/wrr_test.go b/terratest/test/wrr_test.go
@@ -98,7 +98,7 @@ func TestWeightRoundRobinCornerCases(t *testing.T) {
 			"weight-no-labels.example.org",
 			[]string{"172.18.0.3", "172.18.0.4"},
 			func(ips []string) bool {
-				return same(ips, []string{"172.18.0.4", "172.18.0.3"})
+				return containsValues(ips, []string{"172.18.0.4", "172.18.0.3"})
 			},
 		},
 		{