Demonstration of the two's choice method of balancing servers in a reverse proxy

twoschoice.go

package main

import (
	"flag"
	"log"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"
)

func main() {
	port := flag.String("port", "8080", "port to listen on")
	serverSetAddresses := flag.String("servers", "localhost:8080", "server set")
	flag.Parse()

	var serverSet []*url.URL
	for _, address := range strings.Split(*serverSetAddresses, ",") {
		u, err := parseToUrl(address)
		if err != nil {
			log.Fatal(err)
		}
		serverSet = append(serverSet, u)
	}

	proxy := buildReverseProxy(serverSet)
	log.Fatal(http.ListenAndServe(":"+*port, proxy))
}

func buildReverseProxy(serverSet []*url.URL) *httputil.ReverseProxy {
	requestCount := make(map[*url.URL]int)

	director := func(req *http.Request) {
		var choice *url.URL
		option1, option2 := chooseTwo(serverSet)
		if requestCount[option1] < requestCount[option2] {
			choice = option1
		} else {
			choice = option2
		}

		req.URL.Scheme = choice.Scheme
		req.URL.Host = choice.Host
		requestCount[choice]++
	}

	return &httputil.ReverseProxy{Director: director}
}

func chooseTwo(serverSet []*url.URL) (*url.URL, *url.URL) {
	s := rand.NewSource(time.Now().Unix())
	r := rand.New(s)

	c1 := r.Intn(len(serverSet))
	c2 := c1
	for c2 == c1 {
		c2 = r.Intn(len(serverSet))
	}

	return serverSet[c1], serverSet[c2]
}

func parseToUrl(address string) (*url.URL, error) {
	if !strings.HasPrefix(address, "http") {
		address = "http://" + address
	}

	toUrl, err := url.Parse(address)
	if err != nil {
		return nil, err
	}

	return toUrl, nil
}