Add session affinity to custom load balancing

rootfs/etc/nginx/lua/balancer.lua

@@ -5,6 +5,7 @@ local util = require("util")
 local lrucache = require("resty.lrucache")
 local resty_lock = require("resty.lock")
 local ewma = require("balancer.ewma")
+local sticky = require("sticky")
 
 -- measured in seconds
 -- for an Nginx worker to pick up the new list of upstream peers
@@ -48,6 +49,15 @@ end
 local function balance()
   local backend = get_current_backend()
   local lb_alg = get_current_lb_alg()
+  local is_sticky = sticky.is_sticky(backend)
+
+  if is_sticky then
+    local endpoint = sticky.get_endpoint(backend)
+    if endpoint ~= nil then
+      return endpoint.address, endpoint.port
+    end
+    lb_alg = "round_robin"
+  end
 
   if lb_alg == "ip_hash" then
     -- TODO(elvinefendi) implement me
@@ -77,6 +87,9 @@ local function balance()
   if forcible then
     ngx.log(ngx.WARN, "round_robin_state:set valid items forcibly overwritten")
   end
+  if is_sticky then
+    sticky.set_endpoint(endpoint, backend)
+  end
   round_robin_lock:unlock(backend.name .. ROUND_ROBIN_LOCK_KEY)
 
   return endpoint.address, endpoint.port

rootfs/etc/nginx/lua/sticky.lua

@@ -0,0 +1,133 @@
+local json = require('cjson')
+local str = require("resty.string")
+local sha1_crypto = require("resty.sha1")
+local md5_crypto = require("resty.md5")
+
+local sticky_sessions = ngx.shared.sticky_sessions
+
+local DEFAULT_SESSION_COOKIE_NAME = "route"
+local DEFAULT_SESSION_COOKIE_HASH = "md5"
+-- Currently STICKY_TIMEOUT never expires
+local STICKY_TIMEOUT = 0
+
+local _M = {}
+
+local function md5_digest(raw)
+  local md5 = md5_crypto:new()
+  if not md5 then
+    return nil, "md5: failed to create object"
+  end
+  local ok = md5:update(raw)
+  if not ok then
+    return nil, "md5: failed to add data"
+  end
+  local digest = md5:final()
+  if digest == nil then
+    return nil, "md5: failed to create digest"
+  end
+  return str.to_hex(digest), nil
+end
+
+local function sha1_digest(raw)
+  local sha1 = sha1_crypto:new()
+  if not sha1 then
+    return nil, "sha1: failed to create object"
+  end
+  local ok = sha1:update(raw)
+  if not ok then
+    return nil, "sha1: failed to add data"
+  end
+  local digest = sha1:final()
+  if digest == nil then
+    return nil, "sha1: failed to create digest"
+  end
+  return str.to_hex(digest), nil
+end
+
+local function get_cookie_name(backend)
+  local name = backend["sessionAffinityConfig"]["cookieSessionAffinity"]["name"]
+  return name or DEFAULT_SESSION_COOKIE_NAME
+end
+
+local function is_valid_endpoint(backend, address, port)
+  for _, ep in ipairs(backend.endpoints) do
+    if ep.address == address and ep.port == port then
+      return true
+    end
+  end
+  return false
+end
+
+function _M.is_sticky(backend)
+  return backend["sessionAffinityConfig"]["name"] == "cookie"
+end
+
+function _M.get_endpoint(backend)
+  local cookie_name = get_cookie_name(backend)
+  local cookie_key = "cookie_" .. cookie_name
+  local endpoint_key = ngx.var[cookie_key]
+  if endpoint_key == nil then
+    ngx.log(ngx.INFO, string.format(
+      "[backend=%s, affinity=cookie] cookie \"%s\" is not set for this request",
+      backend.name,
+      cookie_name
+    ))
+    return nil
+  end
+
+  local endpoint_string = sticky_sessions:get(endpoint_key)
+  if endpoint_string == nil then
+    ngx.log(ngx.INFO, string.format("[backend=%s, affinity=cookie] no endpoint assigned", backend.name))
+    return nil
+  end
+
+  local endpoint = json.decode(endpoint_string)
+  local valid = is_valid_endpoint(backend, endpoint.address, endpoint.port)
+  if not valid then
+    ngx.log(ngx.INFO, string.format("[backend=%s, affinity=cookie] assigned endpoint is no longer valid", backend.name))
+    sticky_sessions:delete(endpoint_key)
+    return nil
+  end
+  return endpoint
+end
+
+function _M.set_endpoint(endpoint, backend)
+  local cookie_name = get_cookie_name(backend)
+  local encrypted, err
+  local endpoint_string = json.encode(endpoint)
+  local hash = backend["sessionAffinityConfig"]["cookieSessionAffinity"]["hash"]
+
+  if hash == "sha1" then
+    encrypted, err = sha1_digest(endpoint_string)
+  else
+    if hash ~= DEFAULT_SESSION_COOKIE_HASH then
+      ngx.log(ngx.WARN, string.format(
+        "[backend=%s, affinity=cookie] session-cookie-hash \"%s\" is not valid, defaulting to %s",
+        backend.name,
+        hash,
+        DEFAULT_SESSION_COOKIE_HASH
+      ))
+    end
+    encrypted, err = md5_digest(endpoint_string)
+  end
+  if err ~= nil then
+    ngx.log(ngx.WARN, string.format("[backend=%s, affinity=cookie] failed to assign endpoint: %s", backend.name, err))
+    return
+  end
+
+  ngx.log(ngx.INFO, string.format("[backend=%s, affinity=cookie] assigning a new endpoint", backend.name))
+  ngx.header["Set-Cookie"] = cookie_name .. "=" .. encrypted .. ";"
+  local success, forcible
+  success, err, forcible = sticky_sessions:set(encrypted, endpoint_string, STICKY_TIMEOUT)
+  if not success then
+    ngx.log(ngx.WARN, string.format("[backend=%s, affinity=cookie] failed to assign endpoint: %s", backend.name, err))
+  end
+  if forcible then
+    ngx.log(ngx.WARN, string.format(
+      "[backend=%s, affinity=cookie] sticky_sessions shared dict is full; endpoint forcibly overwritten",
+      backend.name
+    ))
+  end
+end
+
+return _M