SRV selection

2 files changed, 41 insertions, 38 deletions

diff --git a/srv.go b/srv.go
index 79f63f0..7453793 100644
--- a/srv.go
+++ b/srv.go
@@ -1,10 +1,8 @@
 package dnsutils
 
 import (
-	"fmt"
 	"math/rand"
 	"net"
-	"os"
 	"sort"
 )
 
@@ -26,21 +24,23 @@ func (s SRVRecords) Less(i, j int) bool {
 
 // OrderedSRV returns a sorted int slice for the entries in the map. To use in the correct order:
 //
-// index, orderedSRV, err := OrderedSRV(addrs)
-// for _, i := range index {
+// count, orderedSRV, err := OrderedSRV(addrs)
+// i := 1
+// for  i <= count {
 //   srv := orderedSRV[i]
 //   // Do something such as dial this SRV. If fails move on the the next
+//   i += 1
 // }
-func OrderedSRV(service, proto, name string) ([]int, map[int]*net.SRV, error) {
+func OrderedSRV(service, proto, name string) (int, map[int]*net.SRV, error) {
 	_, addrs, err := net.LookupSRV(service, proto, name)
 	if err != nil {
-		return []int{}, make(map[int]*net.SRV), err
+		return 0, make(map[int]*net.SRV), err
 	}
 	index, os := orderSRV(addrs)
 	return index, os, nil
 }
 
-func orderSRV(addrs []*net.SRV) ([]int, map[int]*net.SRV) {
+func orderSRV(addrs []*net.SRV) (int, map[int]*net.SRV) {
 	// Initialise the ordered map
 	var o int
 	osrv := make(map[int]*net.SRV)
@@ -56,18 +56,17 @@ func orderSRV(addrs []*net.SRV) ([]int, map[int]*net.SRV) {
 	}
 	sort.Ints(priorities)
 
-	index := make([]int, 0)
+	var count int
 	sort.Ints(priorities)
 	for _, p := range priorities {
 		tos := weightedOrder(prioMap[p])
 		for i, s := range tos {
-			osrv[o+i] = s
-			index = append(index, o+i)
+			count = o + i
+			osrv[count] = s
 		}
 		o += len(tos)
 	}
-	sort.Ints(index)
-	return index, osrv
+	return count, osrv
 }
 
 func weightedOrder(srvs []*net.SRV) map[int]*net.SRV {
@@ -85,26 +84,24 @@ func weightedOrder(srvs []*net.SRV) map[int]*net.SRV {
 	l := len(srvs)
 	for l > 0 {
 		rw := rand.Intn(tw)
-		for i, s := range srvs {
-			// Greater the weight the more likely this will be zero or less
-			//fmt.Fprintf(os.Stderr, "rw: %d\n", rw)
-			rw = rw - int(s.Weight)
-			//fmt.Fprintf(os.Stderr, "tw: %d, rw: %d, w: %d\n", tw, rw, int(s.Weight))
-			if rw <= 0 {
-				// Put entry in position
-				osrv[o] = s
-				if len(srvs) > 1 {
-					// Remove the entry from the source slice by swapping with the last entry and truncating
-					srvs[len(srvs)-1], srvs[i] = srvs[i], srvs[len(srvs)-1]
-					srvs = srvs[:len(srvs)-1]
-					l = len(srvs)
-				} else {
-					l = 0
-				}
-				o += 1
-				tw = tw - int(s.Weight)
-				break
+		i := rand.Intn(l)
+		s := srvs[i]
+		// Greater the weight the more likely this will be zero or less
+		//fmt.Fprintf(os.Stderr, "rw: %d\n", rw)
+		rw = rw - int(s.Weight)
+		if rw <= 0 {
+			// Put entry in position
+			osrv[o] = s
+			if len(srvs) > 1 {
+				// Remove the entry from the source slice by swapping with the last entry and truncating
+				srvs[len(srvs)-1], srvs[i] = srvs[i], srvs[len(srvs)-1]
+				srvs = srvs[:len(srvs)-1]
+				l = len(srvs)
+			} else {
+				l = 0
 			}
+			o += 1
+			tw = tw - int(s.Weight)
 		}
 	}
 	return osrv
diff --git a/srv_test.go b/srv_test.go
index ff47019..95864fe 100644
--- a/srv_test.go
+++ b/srv_test.go
@@ -1,6 +1,7 @@
 package dnsutils
 
 import (
+	"github.com/stretchr/testify/assert"
 	"net"
 	"testing"
 )
@@ -16,7 +17,13 @@ func TestOrderSRV(t *testing.T) {
 		Target:   "t12",
 		Port:     1234,
 		Priority: 1,
-		Weight:   101,
+		Weight:   100,
+	}
+	srv13 := net.SRV{
+		Target:   "t13",
+		Port:     1234,
+		Priority: 1,
+		Weight:   100,
 	}
 	srv21 := net.SRV{
 		Target:   "t21",
@@ -26,11 +33,10 @@ func TestOrderSRV(t *testing.T) {
 	}
 
 	addrs := []*net.SRV{
-		&srv11, &srv21, &srv12,
-	}
-	index, orderedSRV := orderSRV(addrs)
-	for _, i := range index {
-		srv := orderedSRV[i]
-		t.Logf("PRIO: %d WEIGHT: %d TARGET: %s", srv.Priority, srv.Weight, srv.Target)
+		&srv11, &srv21, &srv12, &srv13,
 	}
+	count, orderedSRV := orderSRV(addrs)
+	assert.Equal(t, len(addrs), count, "Index not the expected size")
+	assert.Equal(t, len(addrs), len(orderedSRV), "orderedSRV not the expected size")
+	assert.Equal(t, uint16(2), orderedSRV[4].Priority, "Priority order not as expected")
 }