From 0b166c762754b9e04e33f9ab432b14d2b44c344f Mon Sep 17 00:00:00 2001
From: Alex Ellis <AlexEllis2@gmail.com>
Date: Wed, 5 Apr 2017 08:21:33 +0100
Subject: [PATCH] Enable DNS-RR override for task lookup.

---
 gateway/handlers/proxy.go | 32 +++++++++++++++++++++++++-------
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/gateway/handlers/proxy.go b/gateway/handlers/proxy.go
index 1e47ffbd..ed3e04a4 100644
--- a/gateway/handlers/proxy.go
+++ b/gateway/handlers/proxy.go
@@ -5,11 +5,14 @@ import (
 	"context"
 	"fmt"
 	"io/ioutil"
+	"math/rand"
 	"net"
 	"net/http"
 	"strconv"
 	"time"
 
+	"os"
+
 	"github.com/Sirupsen/logrus"
 	"github.com/alexellis/faas/gateway/metrics"
 	"github.com/docker/docker/api/types"
@@ -111,6 +114,11 @@ func copyHeaders(destination *http.Header, source *http.Header) {
 	}
 }
 
+func randomInt(min, max int) int {
+	rand.Seed(time.Now().Unix())
+	return rand.Intn(max-min) + min
+}
+
 func invokeService(w http.ResponseWriter, r *http.Request, metrics metrics.MetricOptions, service string, requestBody []byte, logger *logrus.Logger, proxyClient *http.Client) {
 	stamp := strconv.FormatInt(time.Now().Unix(), 10)
 
@@ -121,15 +129,24 @@ func invokeService(w http.ResponseWriter, r *http.Request, metrics metrics.Metri
 		metrics.GatewayFunctionsHistogram.WithLabelValues(service).Observe(seconds)
 	}(time.Now())
 
-	watchdogPort := 8080
-	addr, lookupErr := net.LookupIP(service)
-	var url string
-	if len(addr) > 0 && lookupErr == nil {
-		url = fmt.Sprintf("http://%s:%d/", addr[0].String(), watchdogPort)
-	} else {
-		url = fmt.Sprintf("http://%s:%d/", service, watchdogPort)
+	var dnsrr bool
+	if os.Getenv("dnsrr") == "true" {
+		dnsrr = true
 	}
 
+	watchdogPort := 8080
+
+	addr := service
+	// Use DNS-RR via tasks.servicename if enabled as override, otherwise VIP.
+	if dnsrr {
+		entries, lookupErr := net.LookupIP(fmt.Sprintf("tasks.%s", service))
+		if lookupErr == nil && len(entries) > 0 {
+			index := randomInt(0, len(entries))
+			addr = entries[index].String()
+		}
+	}
+	url := fmt.Sprintf("http://%s:%d/", addr, watchdogPort)
+
 	contentType := r.Header.Get("Content-Type")
 	fmt.Printf("[%s] Forwarding request [%s] to: %s\n", stamp, contentType, url)
 
@@ -161,6 +178,7 @@ func invokeService(w http.ResponseWriter, r *http.Request, metrics metrics.Metri
 	clientHeader := w.Header()
 	copyHeaders(&clientHeader, &response.Header)
 
+	// TODO: copyHeaders removes the need for this line - test removal.
 	// Match header for strict services
 	w.Header().Set("Content-Type", r.Header.Get("Content-Type"))