Simple TCP load balancer in Go.

balance.go

package main

import (
    "flag"
    "io"
    "log"
    "net"
    "strings"
)

type Backends struct {
    servers []string
    n       int
}

func (b *Backends) Choose() string {
    idx := b.n % len(b.servers)
    b.n++
    return b.servers[idx]
}

func (b *Backends) String() string {
    return strings.Join(b.servers, ", ")
}

var (
    bind     = flag.String("bind", "", "The address to bind on")
    balance  = flag.String("balance", "", "The backend servers to balance connections across, separated by commas")
    backends *Backends
)

func init() {
    flag.Parse()

    if *bind == "" {
        log.Fatalln("specify the address to listen on with -bind")
    }

    servers := strings.Split(*balance, ",")
    if len(servers) == 1 && servers[0] == "" {
        log.Fatalln("please specify backend servers with -backends")
    }

    backends = &Backends{servers: servers}
}

func copy(wc io.WriteCloser, r io.Reader) {
    defer wc.Close()
    io.Copy(wc, r)
}

func handleConnection(us net.Conn, server string) {
    ds, err := net.Dial("tcp", server)
    if err != nil {
        us.Close()
        log.Printf("failed to dial %s: %s", server, err)
        return
    }

    go copy(ds, us)
    go copy(us, ds)
}

func main() {
    ln, err := net.Listen("tcp", *bind)
    if err != nil {
        log.Fatalf("failed to bind: %s", err)
    }

    log.Printf("listening on %s, balancing %s", *bind, backends)

    for {
        conn, err := ln.Accept()
        if err != nil {
            log.Printf("failed to accept: %s", err)
            continue
        }
        go handleConnection(conn, backends.Choose())
    }
}