Sticky Sessions: Tolerate ClickHouse Session ID Mechanism

main_test.go

@@ -284,6 +284,24 @@ func TestServe(t *testing.T) {
 			},
 			startHTTP,
 		},
+		{
+			"http POST request with session id",
+			"testdata/http-session-id.yml",
+			func(t *testing.T) {
+				req, err := http.NewRequest("POST", "http://127.0.0.1:9090/?query_id=45395792-a432-4b92-8cc9-536c14e1e1a9&extremes=0&session_id=default-session-id233", bytes.NewBufferString("SELECT * FROM system.numbers LIMIT 10"))
+				req.Header.Set("Content-Type", "application/x-www-form-urlencoded;") // This makes it work
+
+				checkErr(t, err)
+				resp, err := http.DefaultClient.Do(req)
+				checkErr(t, err)
+
+				if resp.StatusCode != http.StatusOK || resp.Header.Get("X-Clickhouse-Server-Session-Id") == "" {
+					t.Fatalf("unexpected status code: %d; expected: %d", resp.StatusCode, http.StatusOK)
+				}
+				resp.Body.Close()
+			},
+			startHTTP,
+		},
 		{
 			"http request",
 			"testdata/http.yml",

proxy.go

@@ -51,7 +51,6 @@ func newReverseProxy() *reverseProxy {
 
 func (rp *reverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	startTime := time.Now()
-
 	s, status, err := rp.getScope(req)
 	if err != nil {
 		q := getQuerySnippet(req)
@@ -99,6 +98,11 @@ func (rp *reverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		ReadCloser: req.Body,
 	}
 
+	// publish session_id if needed
+	if s.sessionId != "" {
+		rw.Header().Set("X-ClickHouse-Server-Session-Id", s.sessionId)
+	}
+
 	if s.user.cache == nil {
 		rp.proxyRequest(s, srw, srw, req)
 	} else {
@@ -110,9 +114,9 @@ func (rp *reverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	q := getQuerySnippet(req)
 	if srw.statusCode == http.StatusOK {
 		requestSuccess.With(s.labels).Inc()
-		log.Debugf("%s: request success; query: %q; URL: %q", s, q, req.URL.String())
+		log.Debugf("%s: request success; query: %q; Method: %s; URL: %q", s, q, req.Method, req.URL.String())
 	} else {
-		log.Debugf("%s: request failure: non-200 status code %d; query: %q; URL: %q", s, srw.statusCode, q, req.URL.String())
+		log.Debugf("%s: request failure: non-200 status code %d; query: %q; Method: %s; URL: %q", s, srw.statusCode, q, req.Method, req.URL.String())
 	}
 
 	statusCodes.With(
@@ -435,7 +439,7 @@ func (rp *reverseProxy) applyConfig(cfg *config.Config) error {
 	return nil
 }
 
-// refreshCacheMetrics refresehs cacheSize and cacheItems metrics.
+// refreshCacheMetrics refreshes cacheSize and cacheItems metrics.
 func (rp *reverseProxy) refreshCacheMetrics() {
 	rp.lock.RLock()
 	defer rp.lock.RUnlock()
@@ -452,7 +456,7 @@ func (rp *reverseProxy) refreshCacheMetrics() {
 
 func (rp *reverseProxy) getScope(req *http.Request) (*scope, int, error) {
 	name, password := getAuth(req)
-
+	sessionId := getSessionId(req)
 	var (
 		u  *user
 		c  *cluster
@@ -489,6 +493,6 @@ func (rp *reverseProxy) getScope(req *http.Request) (*scope, int, error) {
 		return nil, http.StatusForbidden, fmt.Errorf("cluster user %q is not allowed to access", cu.name)
 	}
 
-	s := newScope(req, u, c, cu)
+	s := newScope(req, u, c, cu, sessionId)
 	return s, 0, nil
 }

scope.go

@@ -19,6 +19,8 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 )
 
+// var route = make(map[int]*host)
+
 type scopeID uint64
 
 func (sid scopeID) String() string {
@@ -40,6 +42,8 @@ type scope struct {
 	user        *user
 	clusterUser *clusterUser
 
+	sessionId string
+
 	remoteAddr string
 	localAddr  string
 
@@ -49,9 +53,12 @@ type scope struct {
 	labels prometheus.Labels
 }
 
-func newScope(req *http.Request, u *user, c *cluster, cu *clusterUser) *scope {
-	h := c.getHost()
+func newScope(req *http.Request, u *user, c *cluster, cu *clusterUser, sessionId string) *scope {
 
+	h := c.getHost()
+	if sessionId != "" {
+		h = c.getHostSticky(sessionId)
+	}
 	var localAddr string
 	if addr, ok := req.Context().Value(http.LocalAddrContextKey).(net.Addr); ok {
 		localAddr = addr.String()
@@ -63,6 +70,7 @@ func newScope(req *http.Request, u *user, c *cluster, cu *clusterUser) *scope {
 		cluster:     c,
 		user:        u,
 		clusterUser: cu,
+		sessionId:   sessionId,
 
 		remoteAddr: req.RemoteAddr,
 		localAddr:  localAddr,
@@ -305,6 +313,8 @@ var allowedParams = []string{
 	"extremes",
 	// what to do if the volume of the result exceeds one of the limits
 	"result_overflow_mode",
+	// session stickiness
+	"session_id",
 }
 
 // This regexp must match params needed to describe a way to use external data
@@ -810,6 +820,44 @@ func (c *cluster) getReplica() *replica {
 	return r
 }
 
+// getHostSticky returns host by stickiness from replica.
+//
+// Always returns non-nil.
+func (r *replica) getHostSticky(sessionId string) *host {
+	idx := atomic.AddUint32(&r.nextHostIdx, 1)
+	n := uint32(len(r.hosts))
+	if n == 1 {
+		return r.hosts[0]
+	}
+
+	idx %= n
+	h := r.hosts[idx]
+
+	// Scan all the hosts for the least loaded host.
+	for i := uint32(1); i < n; i++ {
+		tmpIdx := (idx + i) % n
+
+		// handling sticky session
+		if sessionId != "" {
+			sessionId := hash(sessionId)
+			tmpIdx = (sessionId) % n
+			tmpHSticky := r.hosts[tmpIdx]
+			log.Debugf("Sticky server candidate is: %s", tmpHSticky.addr)
+			if !tmpHSticky.isActive() {
+				log.Debugf("Sticky session server has been picked up, but it is not available")
+				continue
+			}
+			log.Debugf("Sticky session server is: %s, session_id: %d, server_idx: %d, max nodes in pool: %d", tmpHSticky.addr, sessionId, tmpIdx, n)
+			return tmpHSticky
+		}
+	}
+
+	// The returned host may be inactive. This is OK,
+	// since this means all the hosts are inactive,
+	// so let's try proxying the request to any host.
+	return h
+}
+
 // getHost returns least loaded + round-robin host from replica.
 //
 // Always returns non-nil.
@@ -856,6 +904,14 @@ func (r *replica) getHost() *host {
 	return h
 }
 
+// getHostSticky returns host based on stickiness from cluster.
+//
+// Always returns non-nil.
+func (c *cluster) getHostSticky(sessionId string) *host {
+	r := c.getReplica()
+	return r.getHostSticky(sessionId)
+}
+
 // getHost returns least loaded + round-robin host from cluster.
 //
 // Always returns non-nil.

testdata/http-session-id.yml

@@ -0,0 +1,15 @@
+log_debug: true
+server:
+  http:
+      listen_addr: ":9090"
+      allowed_networks: ["127.0.0.1/24"]
+
+users:
+  - name: "default"
+    to_cluster: "default"
+    to_user: "default"
+
+clusters:
+  - name: "default"
+    nodes:
+      - 127.0.0.1:8124

utils.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"compress/gzip"
 	"fmt"
+	"hash/fnv"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -43,6 +44,13 @@ func getAuth(req *http.Request) (string, string) {
 	return "default", ""
 }
 
+// getSessionId retrieves session id
+func getSessionId(req *http.Request) string {
+	params := req.URL.Query()
+	sessionId := params.Get("session_id")
+	return sessionId
+}
+
 // getQuerySnippet returns query snippet.
 //
 // getQuerySnippet must be called only for error reporting.
@@ -57,6 +65,12 @@ func getQuerySnippet(req *http.Request) string {
 	return query + body
 }
 
+func hash(s string) uint32 {
+	h := fnv.New32a()
+	h.Write([]byte(s))
+	return h.Sum32()
+}
+
 func getQuerySnippetFromBody(req *http.Request) string {
 	if req.Body == nil {
 		return ""