balance.go

package main

import (
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"strings"
	"sync"
)

var quit = make(chan bool)

type Pool struct {
	size  int
	addr  *net.TCPAddr
	conns chan *net.TCPConn
}

func NewPool(server string, size int) *Pool {
	addr, err := net.ResolveTCPAddr("tcp", server)
	if err != nil {
		log.Fatalf("invalid tcp address: %s\n", err)
	}

	p := &Pool{
		size:  size,
		addr:  addr,
		conns: make(chan *net.TCPConn, size),
	}

	for x := 0; x < size; x++ {
		p.conns <- p.Dial()
	}

	return p
}

func (p *Pool) Dial() *net.TCPConn {
	c, err := net.DialTCP("tcp", nil, p.addr)
	if err != nil {
		log.Fatalf("failed to dial %s: %s\n", c.RemoteAddr(), err)
	}
	return c
}

func (p *Pool) Checkout() *net.TCPConn {
	return <-p.conns
}

func (p *Pool) Checkin(conn *net.TCPConn) {
	p.conns <- conn
}

func iocopy(dst, src *net.TCPConn, wg *sync.WaitGroup) {
	defer wg.Done()
	_, err := io.Copy(dst, src)
	log.Println("iocopy called")
	if err != nil {
		log.Fatal(err)
	}
}

func handleConnection(p *Pool, conn *net.TCPConn) {
	var wg sync.WaitGroup

	pconn := p.Checkout()

	wg.Add(2)
	go iocopy(pconn, conn, &wg)
	go iocopy(conn, pconn, &wg)
	wg.Wait()

	err := conn.Close()
	if err != nil {
		log.Fatal(err)
	}

	p.Checkin(pconn)

	log.Printf("connection from %v closed.\n", conn.RemoteAddr())
}

func parseEnviron() (map[string]string, error) {
	hosts := make(map[string]string)

	for _, v := range os.Environ() {
		parts := strings.SplitN(v, "=", 2)

		if !strings.HasPrefix(parts[0], "HEROKU_POSTGRESQL") {
			continue
		}

		url, err := url.Parse(parts[1])
		if err != nil {
			return nil, err
		}

		name, port, _ := net.SplitHostPort(url.Host)
		hosts[name] = port
	}

	return hosts, nil
}

func listen(hostname string, port string) {
	addr, _ := net.ResolveTCPAddr("tcp", net.JoinHostPort("localhost", port+"1"))

	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		log.Fatalf("failed to bind: %s\n", err)
	}

	log.Printf("listening on %s, balancing %s\n", port, hostname)

	p := NewPool(net.JoinHostPort(hostname, port), 1)

	for {
		conn, err := ln.AcceptTCP()
		if err != nil {
			log.Printf("failed to accept: %s\n", err)
			continue
		}
		go handleConnection(p, conn)
	}
}

func main() {
	/* runtime.GOMAXPROCS(runtime.NumCPU()) */
	hosts, _ := parseEnviron()
	for hostname, port := range hosts {
		go listen(hostname, port)
	}

	<-quit
}