From dc52abf4ce05ae5c2235d418549ae636b3db5f15 Mon Sep 17 00:00:00 2001
From: Owen Marshall <owen.marshall@ge.com>
Date: Fri, 13 May 2016 10:22:11 -0400
Subject: [PATCH] Add sticky session support to Traefik.

This change adds sticky session support, by using the new
oxy/rr/StickySession feature.

To use it, set traefik.backend.sticky to true.

This is currently only implemented in the wrr load balancer, and against
the Marathon backend, but lifting it should be very doable.

In the wrr load balancer, a cookie called _TRAEFIK_SERVERNAME will be
set with the backend to use.  If the cookie is altered to an invalid
backend server, or the server is removed from the load balancer, the
next server will be used instead.

Otherwise, the cookie will be checked in Oxy's rr on access and if valid
the connection will be wired through to it.
---
 provider/marathon.go      |  8 ++++++++
 provider/marathon_test.go |  2 +-
 server.go                 | 10 ++++++++++
 templates/marathon.tmpl   |  1 +
 types/types.go            |  1 +
 5 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/provider/marathon.go b/provider/marathon.go
index ecf736c78..2dc7c27bc 100644
--- a/provider/marathon.go
+++ b/provider/marathon.go
@@ -138,6 +138,7 @@ func (provider *Marathon) loadMarathonConfig() *types.Configuration {
 		"getMaxConnAmount":            provider.getMaxConnAmount,
 		"getLoadBalancerMethod":       provider.getLoadBalancerMethod,
 		"getCircuitBreakerExpression": provider.getCircuitBreakerExpression,
+		"getSticky":                   provider.getSticky,
 	}
 
 	applications, err := provider.marathonClient.Applications(nil)
@@ -347,6 +348,13 @@ func (provider *Marathon) getProtocol(task marathon.Task, applications []maratho
 	return "http"
 }
 
+func (provider *Marathon) getSticky(application marathon.Application) string {
+	if sticky, err := provider.getLabel(application, "traefik.backend.sticky"); err == nil {
+		return sticky
+	}
+	return "false"
+}
+
 func (provider *Marathon) getPassHostHeader(application marathon.Application) string {
 	if passHostHeader, err := provider.getLabel(application, "traefik.frontend.passHostHeader"); err == nil {
 		return passHostHeader
diff --git a/provider/marathon_test.go b/provider/marathon_test.go
index cb4edea27..4211f1d26 100644
--- a/provider/marathon_test.go
+++ b/provider/marathon_test.go
@@ -107,7 +107,7 @@ func TestMarathonLoadConfig(t *testing.T) {
 						},
 					},
 					CircuitBreaker: nil,
-					LoadBalancer:   nil,
+					LoadBalancer:   &types.LoadBalancer{Sticky: false},
 				},
 			},
 		},
diff --git a/server.go b/server.go
index 0d6cf7fb6..ab51614c5 100644
--- a/server.go
+++ b/server.go
@@ -537,12 +537,16 @@ func (server *Server) loadConfig(configurations configs, globalConfiguration Glo
 							log.Errorf("Skipping frontend %s...", frontendName)
 							continue frontend
 						}
+
 						lbMethod, err := types.NewLoadBalancerMethod(configuration.Backends[frontend.Backend].LoadBalancer)
 						if err != nil {
 							log.Errorf("Error loading load balancer method '%+v' for frontend %s: %v", configuration.Backends[frontend.Backend].LoadBalancer, frontendName, err)
 							log.Errorf("Skipping frontend %s...", frontendName)
 							continue frontend
 						}
+
+						stickysession := configuration.Backends[frontend.Backend].LoadBalancer.Sticky
+
 						switch lbMethod {
 						case types.Drr:
 							log.Debugf("Creating load-balancer drr")
@@ -565,6 +569,12 @@ func (server *Server) loadConfig(configurations configs, globalConfiguration Glo
 							}
 						case types.Wrr:
 							log.Debugf("Creating load-balancer wrr")
+							if stickysession {
+								cookiename := "_TRAEFIK_SERVERNAME"
+								log.Debugf("... setting to sticky session with cookie named %v", cookiename)
+								sticky := roundrobin.NewStickySession(cookiename)
+								rr, _ = roundrobin.New(saveBackend, roundrobin.EnableStickySession(sticky))
+							}
 							lb = rr
 							for serverName, server := range configuration.Backends[frontend.Backend].Servers {
 								url, err := url.Parse(server.URL)
diff --git a/templates/marathon.tmpl b/templates/marathon.tmpl
index fbd2bf801..04a6912c2 100644
--- a/templates/marathon.tmpl
+++ b/templates/marathon.tmpl
@@ -14,6 +14,7 @@
 {{ if hasLoadBalancerLabels . }}
       [backends.backend{{getFrontendBackend . }}.loadbalancer]
         method = "{{getLoadBalancerMethod . }}"
+        sticky = {{getSticky .}}
 {{end}}
 {{ if hasCircuitBreakerLabels . }}
       [backends.backend{{getFrontendBackend . }}.circuitbreaker]
diff --git a/types/types.go b/types/types.go
index 8b4d2a264..d5a0f0ead 100644
--- a/types/types.go
+++ b/types/types.go
@@ -24,6 +24,7 @@ type MaxConn struct {
 // LoadBalancer holds load balancing configuration.
 type LoadBalancer struct {
 	Method string `json:"method,omitempty"`
+	Sticky bool   `json:"sticky,omitempty"`
 }
 
 // CircuitBreaker holds circuit breaker configuration.