Add sticky session support to rebalancing rr

roundrobin/rebalancer.go

@@ -48,6 +48,9 @@ type Rebalancer struct {
 	// creates new meters
 	newMeter NewMeterFn
 
+	// sticky session object
+	stickySession *StickySession
+
 	requestRewriteListener RequestRewriteListener
 }
 
@@ -80,6 +83,13 @@ func RebalancerErrorHandler(h utils.ErrorHandler) RebalancerOption {
 	}
 }
 
+func RebalancerStickySession(stickySession *StickySession) RebalancerOption {
+	return func(r *Rebalancer) error {
+		r.stickySession = stickySession
+		return nil
+	}
+}
+
 // RebalancerErrorHandler is a functional argument that sets error handler of the server
 func RebalancerRequestRewriteListener(rrl RequestRewriteListener) RebalancerOption {
 	return func(r *Rebalancer) error {
@@ -90,8 +100,9 @@ func RebalancerRequestRewriteListener(rrl RequestRewriteListener) RebalancerOpti
 
 func NewRebalancer(handler balancerHandler, opts ...RebalancerOption) (*Rebalancer, error) {
 	rb := &Rebalancer{
-		mtx:  &sync.Mutex{},
-		next: handler,
+		mtx:           &sync.Mutex{},
+		next:          handler,
+		stickySession: nil,
 	}
 	for _, o := range opts {
 		if err := o(rb); err != nil {
@@ -139,20 +150,42 @@ func (rb *Rebalancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 
 	pw := &utils.ProxyWriter{W: w}
 	start := rb.clock.UtcNow()
-	url, err := rb.next.NextServer()
-	if err != nil {
-		rb.errHandler.ServeHTTP(w, req, err)
-		return
-	}
-
-	if log.GetLevel() >= log.DebugLevel {
-		//log which backend URL we're sending this request to
-		log.WithFields(log.Fields{"Request": utils.DumpHttpRequest(req), "ForwardURL": url}).Debugf("vulcand/oxy/roundrobin/rebalancer: Forwarding this request to URL")
-	}
 
 	// make shallow copy of request before changing anything to avoid side effects
 	newReq := *req
-	newReq.URL = url
+	stuck := false
+
+	if rb.stickySession != nil {
+		cookieUrl, present, err := rb.stickySession.GetBackend(&newReq, rb.Servers())
+
+		if err != nil {
+			log.Infof("vulcand/oxy/roundrobin/rebalancer: error using server from cookie: %v", err)
+		}
+
+		if present {
+			newReq.URL = cookieUrl
+			stuck = true
+		}
+	}
+
+	if !stuck {
+		url, err := rb.next.NextServer()
+		if err != nil {
+			rb.errHandler.ServeHTTP(w, req, err)
+			return
+		}
+
+		if log.GetLevel() >= log.DebugLevel {
+			//log which backend URL we're sending this request to
+			log.WithFields(log.Fields{"Request": utils.DumpHttpRequest(req), "ForwardURL": url}).Debugf("vulcand/oxy/roundrobin/rebalancer: Forwarding this request to URL")
+		}
+
+		if rb.stickySession != nil {
+			rb.stickySession.StickBackend(url, &w)
+		}
+
+		newReq.URL = url
+	}
 
 	//Emit event to a listener if one exists
 	if rb.requestRewriteListener != nil {
@@ -161,7 +194,7 @@ func (rb *Rebalancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 
 	rb.next.Next().ServeHTTP(pw, &newReq)
 
-	rb.recordMetrics(url, pw.Code, rb.clock.UtcNow().Sub(start))
+	rb.recordMetrics(newReq.URL, pw.Code, rb.clock.UtcNow().Sub(start))
 	rb.adjustWeights()
 }
 
@@ -244,11 +277,11 @@ func (rb *Rebalancer) upsertServer(u *url.URL, weight int) error {
 	return nil
 }
 
-func (r *Rebalancer) findServer(u *url.URL) (*rbServer, int) {
-	if len(r.servers) == 0 {
+func (rb *Rebalancer) findServer(u *url.URL) (*rbServer, int) {
+	if len(rb.servers) == 0 {
 		return nil, -1
 	}
-	for i, s := range r.servers {
+	for i, s := range rb.servers {
 		if sameURL(u, s.url) {
 			return s, i
 		}
@@ -351,7 +384,7 @@ func (rb *Rebalancer) markServers() bool {
 }
 
 func (rb *Rebalancer) convergeWeights() bool {
-	// If we have previoulsy changed servers try to restore weights to the original state
+	// If we have previously changed servers try to restore weights to the original state
 	changed := false
 	for _, s := range rb.servers {
 		if s.origWeight == s.curWeight {

roundrobin/rebalancer_test.go

@@ -1,14 +1,14 @@
 package roundrobin
 
 import (
+	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"time"
 
 	"github.com/mailgun/timetools"
 	"github.com/vulcand/oxy/forward"
 	"github.com/vulcand/oxy/testutils"
-
 	. "gopkg.in/check.v1"
 )
 
@@ -339,6 +339,52 @@ func (s *RBSuite) TestRequestRewriteListener(c *C) {
 	c.Assert(rb.requestRewriteListener, NotNil)
 }
 
+func (s *RBSuite) TestRebalancerStickySession(c *C) {
+	a, b, x := testutils.NewResponder("a"), testutils.NewResponder("b"), testutils.NewResponder("x")
+	defer a.Close()
+	defer b.Close()
+	defer x.Close()
+
+	sticky := NewStickySession("test")
+	c.Assert(sticky, NotNil)
+
+	fwd, err := forward.New()
+	c.Assert(err, IsNil)
+
+	lb, err := New(fwd)
+	c.Assert(err, IsNil)
+
+	rb, err := NewRebalancer(lb, RebalancerStickySession(sticky))
+	c.Assert(err, IsNil)
+
+	rb.UpsertServer(testutils.ParseURI(a.URL))
+	rb.UpsertServer(testutils.ParseURI(b.URL))
+	rb.UpsertServer(testutils.ParseURI(x.URL))
+
+	proxy := httptest.NewServer(rb)
+	defer proxy.Close()
+
+	for i := 0; i < 10; i++ {
+		req, err := http.NewRequest(http.MethodGet, proxy.URL, nil)
+		c.Assert(err, IsNil)
+		req.AddCookie(&http.Cookie{Name: "test", Value: a.URL})
+
+		resp, err := http.DefaultClient.Do(req)
+		c.Assert(err, IsNil)
+
+		defer resp.Body.Close()
+		body, err := ioutil.ReadAll(resp.Body)
+
+		c.Assert(err, IsNil)
+		c.Assert(string(body), Equals, "a")
+	}
+
+	c.Assert(rb.RemoveServer(testutils.ParseURI(a.URL)), IsNil)
+	c.Assert(seq(c, proxy.URL, 3), DeepEquals, []string{"b", "x", "b"})
+	c.Assert(rb.RemoveServer(testutils.ParseURI(b.URL)), IsNil)
+	c.Assert(seq(c, proxy.URL, 3), DeepEquals, []string{"x", "x", "x"})
+}
+
 type testMeter struct {
 	rating   float64
 	notReady bool