github.com/iron-io/functions@v0.0.0-20180820112432-d59d7d1c40b2/lb/lb.go (about) 1 package lb 2 3 import ( 4 "bytes" 5 "context" 6 "net/http" 7 "net/http/httputil" 8 "strconv" 9 "sync" 10 "sync/atomic" 11 12 "github.com/golang/groupcache/consistenthash" 13 ) 14 15 // ConsistentHashReverseProxy returns a new ReverseProxy that routes 16 // URLs to the scheme, host, and base path provided in by a consistent hash 17 // algorithm. If the target's path is "/base" and the incoming request was for 18 // "/dir", the target request will be for /base/dir. 19 // ConsistentHashReverseProxy does not rewrite the Host header. 20 func ConsistentHashReverseProxy(ctx context.Context, nodes []string) *httputil.ReverseProxy { 21 ch := consistenthash.New(len(nodes), nil) 22 ch.Add(nodes...) 23 24 bufPool := sync.Pool{ 25 New: func() interface{} { 26 return new(bytes.Buffer) 27 }, 28 } 29 30 var i int64 31 director := func(req *http.Request) { 32 buf := bufPool.Get().(*bytes.Buffer) 33 defer bufPool.Put(buf) 34 buf.Reset() 35 buf.WriteString(req.URL.Path) 36 buf.WriteString("??") 37 b := strconv.AppendInt(buf.Bytes(), atomic.AddInt64(&i, 1), 10) 38 39 target := ch.Get(string(b)) 40 req.URL.Scheme = "http" 41 req.URL.Host = target 42 if _, ok := req.Header["User-Agent"]; !ok { 43 // explicitly disable User-Agent so it's not set to default value 44 req.Header.Set("User-Agent", "") 45 } 46 } 47 48 return &httputil.ReverseProxy{ 49 Director: director, 50 Transport: NewRoundTripper(ctx, nodes), 51 } 52 }