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  }