├── .gitignore ├── Godeps ├── README.md ├── version.go ├── LICENSE ├── backends ├── backends.go ├── round_robin.go └── round_robin_test.go ├── shared.go ├── http.go ├── https.go ├── tcp.go └── balance.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 2 | *.o 3 | *.a 4 | *.so 5 | 6 | # Folders 7 | _obj 8 | _test 9 | 10 | # Architecture specific extensions/prefixes 11 | *.[568vq] 12 | [568vq].out 13 | 14 | *.cgo1.go 15 | *.cgo2.c 16 | _cgo_defun.c 17 | _cgo_gotypes.go 18 | _cgo_export.* 19 | 20 | _testmain.go 21 | 22 | *.exe 23 | -------------------------------------------------------------------------------- /Godeps: -------------------------------------------------------------------------------- 1 | { 2 | "ImportPath": "github.com/darkhelmet/balance", 3 | "GoVersion": "go1.2", 4 | "Deps": [ 5 | { 6 | "ImportPath": "github.com/gonuts/commander", 7 | "Rev": "6f9453dc9492c0be4a631daebfa5e9e8243a52c6" 8 | }, 9 | { 10 | "ImportPath": "github.com/gonuts/flag", 11 | "Rev": "741a6cbd37a30dedc93f817e7de6aaf0ca38a493" 12 | } 13 | ] 14 | } 15 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # balance 2 | 3 | Simple TCP/HTTP/HTTPS load balancer in Go 4 | 5 | ## Install 6 | 7 | go get github.com/darkhelmet/balance 8 | 9 | ## Usage 10 | 11 | # Simple tcp mode 12 | balance tcp -bind :4000 localhost:4001 localhost:4002 13 | 14 | # HTTP mode 15 | balance http -bind :4000 localhost:4001 localhost:4002 16 | 17 | # HTTPS mode 18 | balance https -bind :4000 -cert ssl.crt -key ssl.key localhost:4001 localhost:4002 19 | 20 | ## License 21 | 22 | GNU AGPL, see `LICENSE` file. 23 | -------------------------------------------------------------------------------- /version.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/gonuts/commander" 7 | ) 8 | 9 | const Version = "0.0.1" 10 | 11 | func version(cmd *commander.Command, args []string) error { 12 | fmt.Println(Version) 13 | return nil 14 | } 15 | 16 | func init() { 17 | fs := newFlagSet("tcp") 18 | 19 | cmd.Subcommands = append(cmd.Subcommands, &commander.Command{ 20 | UsageLine: "version", 21 | Short: "display version information", 22 | Flag: *fs, 23 | Run: version, 24 | }) 25 | } 26 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (C) 2013 Daniel Huckstep 2 | 3 | This program is free software: you can redistribute it and/or modify 4 | it under the terms of the GNU Affero General Public License as 5 | published by the Free Software Foundation, either version 3 of the 6 | License, or (at your option) any later version. 7 | 8 | This program is distributed in the hope that it will be useful, 9 | but WITHOUT ANY WARRANTY; without even the implied warranty of 10 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 11 | GNU Affero General Public License for more details. 12 | 13 | You should have received a copy of the GNU Affero General Public License 14 | along with this program. If not, see . 15 | -------------------------------------------------------------------------------- /backends/backends.go: -------------------------------------------------------------------------------- 1 | package backends 2 | 3 | import ( 4 | "log" 5 | ) 6 | 7 | type backend struct { 8 | hostname string 9 | } 10 | 11 | func (b *backend) String() string { 12 | return b.hostname 13 | } 14 | 15 | type Backend interface { 16 | String() string 17 | } 18 | 19 | type Backends interface { 20 | Choose() Backend 21 | Len() int 22 | Add(string) 23 | Remove(string) 24 | } 25 | 26 | type Factory func([]string) Backends 27 | 28 | var factories = make(map[string]Factory) 29 | 30 | func Build(algorithm string, specs []string) Backends { 31 | factory, found := factories[algorithm] 32 | if !found { 33 | log.Fatalf("balance algorithm %s not supported", algorithm) 34 | } 35 | return factory(specs) 36 | } 37 | -------------------------------------------------------------------------------- /shared.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "net" 5 | "net/http" 6 | "net/http/httputil" 7 | ) 8 | 9 | const ( 10 | colon = ":" 11 | XRealIP = "X-Real-IP" 12 | ) 13 | 14 | type NoBackend struct{} 15 | 16 | func RealIP(req *http.Request) string { 17 | host, _, _ := net.SplitHostPort(req.RemoteAddr) 18 | return host 19 | } 20 | 21 | type Proxy struct { 22 | *httputil.ReverseProxy 23 | } 24 | 25 | func (p *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) { 26 | defer func() { 27 | if err := recover(); err != nil { 28 | switch err.(type) { 29 | case NoBackend: 30 | rw.WriteHeader(503) 31 | req.Body.Close() 32 | default: 33 | panic(err) 34 | } 35 | } 36 | }() 37 | p.ReverseProxy.ServeHTTP(rw, req) 38 | } 39 | -------------------------------------------------------------------------------- /http.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "net/http" 6 | "net/http/httputil" 7 | 8 | BA "github.com/darkhelmet/balance/backends" 9 | "github.com/gonuts/commander" 10 | ) 11 | 12 | func httpBalance(bind string, backends BA.Backends) error { 13 | log.Println("using http balancing") 14 | proxy := &Proxy{ 15 | &httputil.ReverseProxy{Director: func(req *http.Request) { 16 | backend := backends.Choose() 17 | if backend == nil { 18 | log.Printf("no backend for client %s", req.RemoteAddr) 19 | panic(NoBackend{}) 20 | } 21 | req.URL.Scheme = "http" 22 | req.URL.Host = backend.String() 23 | req.Header.Add(XRealIP, RealIP(req)) 24 | }}, 25 | } 26 | http.DefaultTransport.(*http.Transport).MaxIdleConnsPerHost = 100 27 | log.Printf("listening on %s, balancing %d backends", bind, backends.Len()) 28 | return http.ListenAndServe(bind, proxy) 29 | } 30 | 31 | func init() { 32 | fs := newFlagSet("http") 33 | 34 | cmd.Subcommands = append(cmd.Subcommands, &commander.Command{ 35 | UsageLine: "http [options] []", 36 | Short: "performs http based load balancing", 37 | Flag: *fs, 38 | Run: balancer(httpBalance), 39 | }) 40 | } 41 | -------------------------------------------------------------------------------- /backends/round_robin.go: -------------------------------------------------------------------------------- 1 | package backends 2 | 3 | import ( 4 | "container/ring" 5 | "sync" 6 | ) 7 | 8 | type RoundRobin struct { 9 | r *ring.Ring 10 | l sync.RWMutex 11 | } 12 | 13 | func NewRoundRobin(strs []string) Backends { 14 | r := ring.New(len(strs)) 15 | for _, s := range strs { 16 | r.Value = &backend{s} 17 | r = r.Next() 18 | } 19 | return &RoundRobin{r: r} 20 | } 21 | 22 | func init() { 23 | factories["round-robin"] = NewRoundRobin 24 | } 25 | 26 | func (rr *RoundRobin) Len() int { 27 | rr.l.RLock() 28 | defer rr.l.RUnlock() 29 | return rr.r.Len() 30 | } 31 | 32 | func (rr *RoundRobin) Choose() Backend { 33 | rr.l.Lock() 34 | defer rr.l.Unlock() 35 | if rr.r == nil { 36 | return nil 37 | } 38 | n := rr.r.Value.(*backend) 39 | rr.r = rr.r.Next() 40 | return n 41 | } 42 | 43 | func (rr *RoundRobin) Add(s string) { 44 | rr.l.Lock() 45 | defer rr.l.Unlock() 46 | nr := &ring.Ring{Value: &backend{s}} 47 | if rr.r == nil { 48 | rr.r = nr 49 | } else { 50 | rr.r = rr.r.Link(nr).Next() 51 | } 52 | } 53 | 54 | func (rr *RoundRobin) Remove(s string) { 55 | rr.l.Lock() 56 | defer rr.l.Unlock() 57 | r := rr.r 58 | if rr.r.Len() == 1 { 59 | rr.r = ring.New(0) 60 | return 61 | } 62 | 63 | for i := rr.r.Len(); i > 0; i-- { 64 | r = r.Next() 65 | ba := r.Value.(*backend) 66 | if s == ba.String() { 67 | rr.r = r.Unlink(1) 68 | return 69 | } 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /https.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "net/http" 6 | "net/http/httputil" 7 | 8 | BA "github.com/darkhelmet/balance/backends" 9 | "github.com/gonuts/commander" 10 | ) 11 | 12 | var ( 13 | httpsOptions = struct { 14 | certFile, keyFile string 15 | }{} 16 | ) 17 | 18 | func httpsBalance(bind string, backends BA.Backends) error { 19 | if httpsOptions.certFile == "" || httpsOptions.keyFile == "" { 20 | log.Fatalln("specify both -cert and -key") 21 | } 22 | 23 | log.Println("using https balancing") 24 | 25 | proxy := &Proxy{ 26 | &httputil.ReverseProxy{Director: func(req *http.Request) { 27 | backend := backends.Choose() 28 | if backend == nil { 29 | log.Printf("no backend for client %s", req.RemoteAddr) 30 | panic(NoBackend{}) 31 | } 32 | req.URL.Scheme = "http" 33 | req.Header.Add("X-Forwarded-Proto", "https") 34 | req.URL.Host = backend.String() 35 | req.Header.Add(XRealIP, RealIP(req)) 36 | }}, 37 | } 38 | http.DefaultTransport.(*http.Transport).MaxIdleConnsPerHost = 100 39 | log.Printf("listening on %s, balancing %d backends", bind, backends.Len()) 40 | return http.ListenAndServeTLS(bind, httpsOptions.certFile, httpsOptions.keyFile, proxy) 41 | } 42 | 43 | func init() { 44 | fs := newFlagSet("https") 45 | fs.StringVar(&httpsOptions.certFile, "cert", "", "the SSL certificate file to use") 46 | fs.StringVar(&httpsOptions.keyFile, "key", "", "the SSL key file to use") 47 | 48 | cmd.Subcommands = append(cmd.Subcommands, &commander.Command{ 49 | UsageLine: "https [options] []", 50 | Short: "performs https based load balancing", 51 | Flag: *fs, 52 | Run: balancer(httpsBalance), 53 | }) 54 | } 55 | -------------------------------------------------------------------------------- /tcp.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "log" 7 | "net" 8 | 9 | BA "github.com/darkhelmet/balance/backends" 10 | "github.com/gonuts/commander" 11 | ) 12 | 13 | func copy(wc io.WriteCloser, r io.Reader) { 14 | defer wc.Close() 15 | io.Copy(wc, r) 16 | } 17 | 18 | func handleConnection(us net.Conn, backend BA.Backend) { 19 | if backend == nil { 20 | log.Printf("no backend available for connection from %s", us.RemoteAddr()) 21 | us.Close() 22 | return 23 | } 24 | 25 | ds, err := net.Dial("tcp", backend.String()) 26 | if err != nil { 27 | log.Printf("failed to dial %s: %s", backend, err) 28 | us.Close() 29 | return 30 | } 31 | 32 | // Ignore errors 33 | go copy(ds, us) 34 | go copy(us, ds) 35 | } 36 | 37 | func tcpBalance(bind string, backends BA.Backends) error { 38 | log.Println("using tcp balancing") 39 | ln, err := net.Listen("tcp", bind) 40 | if err != nil { 41 | return fmt.Errorf("failed to bind: %s", err) 42 | } 43 | 44 | log.Printf("listening on %s, balancing %d backends", bind, backends.Len()) 45 | 46 | for { 47 | conn, err := ln.Accept() 48 | if err != nil { 49 | log.Printf("failed to accept: %s", err) 50 | continue 51 | } 52 | go handleConnection(conn, backends.Choose()) 53 | } 54 | 55 | return err 56 | } 57 | 58 | func init() { 59 | fs := newFlagSet("tcp") 60 | 61 | cmd.Subcommands = append(cmd.Subcommands, &commander.Command{ 62 | UsageLine: "tcp [options] []", 63 | Short: "performs tcp based load balancing", 64 | Flag: *fs, 65 | Run: balancer(tcpBalance), 66 | }) 67 | } 68 | -------------------------------------------------------------------------------- /balance.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "os" 6 | 7 | BA "github.com/darkhelmet/balance/backends" 8 | "github.com/gonuts/commander" 9 | "github.com/gonuts/flag" 10 | ) 11 | 12 | var cmd = &commander.Command{ 13 | Short: "load balance tcp, http, and https connections to multiple backends", 14 | } 15 | 16 | func ensureBind(bindFlag *flag.Flag) string { 17 | if bindFlag == nil { 18 | log.Fatalln("bind flag not defined") 19 | } 20 | 21 | bind, ok := bindFlag.Value.Get().(string) 22 | if !ok { 23 | log.Fatalln("bind flag must be defined as a string") 24 | } 25 | 26 | if bind == "" { 27 | log.Fatalln("specify the address to listen on with -bind") 28 | } 29 | 30 | return bind 31 | } 32 | 33 | func buildBackends(balanceFlag *flag.Flag, backends []string) BA.Backends { 34 | if balanceFlag == nil { 35 | log.Fatalln("balance flag not defined") 36 | } 37 | 38 | balance, ok := balanceFlag.Value.Get().(string) 39 | if !ok { 40 | log.Fatalln("balance flag must be defined as a string") 41 | } 42 | 43 | if balance == "" { 44 | log.Fatalln("specify the balancing algorithm with -balance") 45 | } 46 | 47 | return BA.Build(balance, backends) 48 | } 49 | 50 | func newFlagSet(name string) *flag.FlagSet { 51 | fs := flag.NewFlagSet(name, flag.ExitOnError) 52 | fs.String("bind", "", "the address to listen on") 53 | fs.String("balance", "round-robin", "the balancing algorithm to use") 54 | return fs 55 | } 56 | 57 | func balancer(f func(string, BA.Backends) error) func(*commander.Command, []string) error { 58 | return func(cmd *commander.Command, args []string) error { 59 | bind := ensureBind(cmd.Flag.Lookup("bind")) 60 | backends := buildBackends(cmd.Flag.Lookup("balance"), args) 61 | return f(bind, backends) 62 | } 63 | } 64 | 65 | func main() { 66 | err := cmd.Dispatch(os.Args[1:]) 67 | if err != nil { 68 | log.Fatalf("%v\n", err) 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /backends/round_robin_test.go: -------------------------------------------------------------------------------- 1 | package backends_test 2 | 3 | import ( 4 | BA "github.com/darkhelmet/balance/backends" 5 | . "launchpad.net/gocheck" 6 | "testing" 7 | ) 8 | 9 | func Test(t *testing.T) { TestingT(t) } 10 | 11 | type S struct{} 12 | 13 | var ( 14 | _ = Suite(&S{}) 15 | a = "1.2.3.4" 16 | b = "2.3.4.5" 17 | ) 18 | 19 | func (s *S) TestLen(c *C) { 20 | r := BA.NewRoundRobin([]string{a, b}) 21 | c.Assert(r.Len(), Equals, 2) 22 | } 23 | 24 | func (s *S) TestChoose(c *C) { 25 | r := BA.NewRoundRobin([]string{a, b}) 26 | c.Assert(r.Choose().String(), Equals, a) 27 | c.Assert(r.Choose().String(), Equals, b) 28 | c.Assert(r.Choose().String(), Equals, a) 29 | c.Assert(r.Choose().String(), Equals, b) 30 | } 31 | 32 | func (s *S) TestChooseEmpty(c *C) { 33 | r := BA.NewRoundRobin([]string{}) 34 | c.Assert(r.Choose(), Equals, nil) 35 | } 36 | 37 | func (s *S) TestAdd(c *C) { 38 | r := BA.NewRoundRobin([]string{a}) 39 | c.Assert(r.Choose().String(), Equals, a) 40 | c.Assert(r.Choose().String(), Equals, a) 41 | r.Add(b) 42 | c.Assert(r.Choose().String(), Equals, b) 43 | c.Assert(r.Choose().String(), Equals, a) 44 | } 45 | 46 | func (s *S) TestAddEmpty(c *C) { 47 | r := BA.NewRoundRobin([]string{}) 48 | c.Assert(r.Len(), Equals, 0) 49 | r.Add(a) 50 | c.Assert(r.Len(), Equals, 1) 51 | c.Assert(r.Choose().String(), Equals, a) 52 | } 53 | 54 | func (s *S) TestRemove(c *C) { 55 | r := BA.NewRoundRobin([]string{a, b}) 56 | c.Assert(r.Len(), Equals, 2) 57 | c.Assert(r.Choose().String(), Equals, a) 58 | c.Assert(r.Choose().String(), Equals, b) 59 | r.Remove(b) 60 | c.Assert(r.Len(), Equals, 1) 61 | c.Assert(r.Choose().String(), Equals, a) 62 | c.Assert(r.Choose().String(), Equals, a) 63 | r.Remove(a) 64 | c.Assert(r.Len(), Equals, 0) 65 | } 66 | 67 | func (s *S) TestRemoveEmpty(c *C) { 68 | r := BA.NewRoundRobin([]string{}) 69 | c.Assert(r.Len(), Equals, 0) 70 | r.Remove(a) 71 | c.Assert(r.Len(), Equals, 0) 72 | r.Add(a) 73 | c.Assert(r.Len(), Equals, 1) 74 | r.Remove(a) 75 | c.Assert(r.Len(), Equals, 0) 76 | r.Remove(a) 77 | } 78 | --------------------------------------------------------------------------------