[V2] Add support for MaglevHashEndpoints (#530)

Author: nikita-vanyasin

CHANGELOG.md

@@ -8,6 +8,8 @@
 - [V2] Add optional status code checks. Consistent return of response
 - [V2] JavaScript Transactions API
 - [V2] Async Client
+- [V2] Fix connection.NewRequestWithEndpoint()
+- [V2] Add support for MaglevHashEndpoints
 
 ## [1.6.0](https://github.com/arangodb/go-driver/tree/v1.6.0) (2023-05-30)
 - Add ErrArangoDatabaseNotFound and IsExternalStorageError helper to v2

v2/connection/connection_http_internal.go

@@ -128,34 +128,31 @@ func (j *httpConnection) SetEndpoint(e Endpoint) error {
 	return nil
 }
 
-func (j *httpConnection) NewRequestWithEndpoint(endpoint string, method string, urls ...string) (Request, error) {
-	return j.newRequestWithEndpoint(endpoint, method, urls...)
+func (j *httpConnection) NewRequestWithEndpoint(endpoint string, method string, urlParts ...string) (Request, error) {
+	return j.newRequestWithEndpoint(endpoint, method, urlParts...)
 }
 
-func (j *httpConnection) NewRequest(method string, urls ...string) (Request, error) {
-	return j.newRequest(method, urls...)
+func (j *httpConnection) NewRequest(method string, urlParts ...string) (Request, error) {
+	return j.newRequestWithEndpoint("", method, urlParts...)
 }
 
-func (j *httpConnection) newRequest(method string, urls ...string) (*httpRequest, error) {
-	return j.newRequestWithEndpoint("", method, urls...)
-}
+func (j *httpConnection) newRequestWithEndpoint(endpoint string, method string, urlParts ...string) (*httpRequest, error) {
+	urlPath := path.Join(urlParts...)
 
-func (j *httpConnection) newRequestWithEndpoint(endpoint string, method string, urls ...string) (*httpRequest, error) {
-	e, ok := j.endpoint.Get(endpoint)
-	if !ok {
-		return nil, errors.Errorf("Unable to resolve endpoint for %s", e)
+	e, err := j.endpoint.Get(endpoint, method, urlPath)
+	if err != nil {
+		return nil, errors.Errorf("Unable to resolve endpoint for %s", endpoint)
 	}
-	url, err := url.Parse(e)
+	u, err := url.Parse(e)
 	if err != nil {
 		return nil, err
 	}
-
-	url.Path = path.Join(url.Path, path.Join(urls...))
+	u.Path = path.Join(u.Path, urlPath)
 
 	r := &httpRequest{
 		method:   method,
-		url:      url,
-		endpoint: endpoint,
+		url:      u,
+		endpoint: e,
 	}
 
 	return r, nil

v2/connection/connection_http_internal_test.go

@@ -1,7 +1,7 @@
 //
 // DISCLAIMER
 //
-// Copyright 2021 ArangoDB GmbH, Cologne, Germany
+// Copyright 2021-2023 ArangoDB GmbH, Cologne, Germany
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -17,12 +17,12 @@
 //
 // Copyright holder is ArangoDB GmbH, Cologne, Germany
 //
-// Author Tomasz Mielech
-//
 
 package connection
 
 import (
+	"net/http"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -79,3 +79,40 @@ func Test_httpConnection_Decoder(t *testing.T) {
 		})
 	}
 }
+
+func Test_httpConnection_NewRequest(t *testing.T) {
+	eps := []string{
+		"https://a:8529", "https://a:8539", "https://b:8529",
+	}
+
+	c := httpConnection{
+		endpoint: NewRoundRobinEndpoints(eps),
+	}
+
+	j := 0
+	for i := 0; i < 10; i++ {
+		expectedEp := eps[j]
+		req, err := c.NewRequest(http.MethodGet, "_api/version")
+		require.NoError(t, err)
+		require.Equal(t, expectedEp, req.Endpoint())
+		require.True(t, strings.HasPrefix(req.URL(), expectedEp))
+		j++
+		if j >= len(eps) {
+			j = 0
+		}
+	}
+}
+
+func Test_httpConnection_NewRequestWithEndpoint(t *testing.T) {
+	c := httpConnection{
+		endpoint: NewRoundRobinEndpoints([]string{"https://a:8529", "https://a:8539", "https://b:8529"}),
+	}
+
+	for i := 0; i < 10; i++ {
+		ep := "https://a:8539"
+		req, err := c.NewRequestWithEndpoint(ep, http.MethodGet, "_api/version")
+		require.NoError(t, err)
+		require.Equal(t, ep, req.Endpoint())
+		require.True(t, strings.HasPrefix(req.URL(), ep))
+	}
+}

v2/connection/endpoint.go

@@ -1,7 +1,7 @@
 //
 // DISCLAIMER
 //
-// Copyright 2020 ArangoDB GmbH, Cologne, Germany
+// Copyright 2020-2023 ArangoDB GmbH, Cologne, Germany
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -17,55 +17,30 @@
 //
 // Copyright holder is ArangoDB GmbH, Cologne, Germany
 //
-// Author Adam Janikowski
-//
 
 package connection
 
-import "sync"
+import (
+	"strings"
+)
 
 type Endpoint interface {
-	// Get return one of endpoints if is valid, if no default one is returned
-	Get(endpoints ...string) (string, bool)
-
+	// Get returns provided endpoint if it is known, otherwise chooses one endpoint from existing list
+	// Endpoint implementation might use the Request method and path values to determine which endpoint to return
+	Get(endpoint, method, path string) (string, error)
+	// List returns known endpoints
 	List() []string
 }
 
-func NewEndpoints(e ...string) Endpoint {
-	return &endpoints{
-		endpoints: e,
-	}
-}
-
-type endpoints struct {
-	lock sync.Mutex
-
-	endpoints []string
-
-	index int
-}
-
-func (e *endpoints) List() []string {
-	return e.endpoints
-}
-
-func (e *endpoints) Get(endpoints ...string) (string, bool) {
-	e.lock.Lock()
-	defer e.lock.Unlock()
-
-	if len(e.endpoints) == 0 {
-		return "", false
-	}
-
-	//return e.endpoints[0], true
-
-	if e.index >= len(e.endpoints) {
-		e.index = 0
-	}
-
-	r := e.endpoints[e.index]
-
-	e.index++
-
-	return r, true
+var (
+	urlFixer = strings.NewReplacer(
+		"tcp://", "http://",
+		"ssl://", "https://",
+	)
+)
+
+// FixupEndpointURLScheme changes endpoint URL schemes used by arangod to ones used by go.
+// E.g. "tcp://localhost:8529" -> "http://localhost:8529"
+func FixupEndpointURLScheme(u string) string {
+	return urlFixer.Replace(u)
 }

v2/connection/endpoints.go

@@ -0,0 +1,117 @@
+//
+// DISCLAIMER
+//
+// Copyright 2020-2023 ArangoDB GmbH, Cologne, Germany
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Copyright holder is ArangoDB GmbH, Cologne, Germany
+//
+
+package connection
+
+import (
+	"fmt"
+	"math/big"
+	"sort"
+	"strings"
+
+	"github.com/kkdai/maglev"
+	"github.com/pkg/errors"
+)
+
+// RequestHashValueExtractor accepts request method and full request path and must return a value which will be used for hash calculation
+type RequestHashValueExtractor func(requestMethod, requestPath string) (string, error)
+
+// NewMaglevHashEndpoints returns Endpoint manager which runs round-robin
+func NewMaglevHashEndpoints(eps []string, extractor RequestHashValueExtractor) (Endpoint, error) {
+	// order of endpoints affects hashing result
+	sort.Strings(eps)
+
+	// lookupSize must be equal or greater than len(eps) and it must be a prime number
+	lookupSize := findNextPrime(uint64(len(eps)))
+
+	table, err := maglev.NewMaglev(eps, lookupSize)
+	if err != nil {
+		return nil, err
+	}
+
+	return &maglevHashEndpoints{
+		extractor:   extractor,
+		endpoints:   eps,
+		maglevTable: table,
+	}, nil
+}
+
+func findNextPrime(i uint64) uint64 {
+	bigInt := big.NewInt(0).SetUint64(i)
+
+	for {
+		if bigInt.ProbablyPrime(1) {
+			return bigInt.Uint64()
+		}
+		i++
+		bigInt.SetUint64(i)
+	}
+}
+
+type maglevHashEndpoints struct {
+	extractor   RequestHashValueExtractor
+	endpoints   []string
+	maglevTable *maglev.Maglev
+}
+
+func (e *maglevHashEndpoints) List() []string {
+	return e.endpoints
+}
+
+func (e *maglevHashEndpoints) Get(providedEp, requestMethod, requestPath string) (string, error) {
+	if len(e.endpoints) == 0 {
+		return "", errors.New("no endpoints known")
+	}
+
+	for _, known := range e.endpoints {
+		if known == providedEp {
+			return known, nil
+		}
+	}
+
+	val, err := e.extractor(requestMethod, requestPath)
+	if err != nil {
+		return "", errors.WithMessagef(err, "could not extract value for method '%s' path '%s'", requestMethod, requestPath)
+	}
+
+	r, err := e.maglevTable.Get(val)
+	if err != nil {
+		return r, errors.WithMessage(err, "failed to lookup Maglev table")
+	}
+
+	return r, nil
+}
+
+var _ RequestHashValueExtractor = RequestDBNameValueExtractor
+
+// RequestDBNameValueExtractor might be used as RequestHashValueExtractor to use DB name from URL for hashing
+// It fallbacks to requestMethod+requestPath concatenation in case if path does not contain DB name
+func RequestDBNameValueExtractor(requestMethod, requestPath string) (string, error) {
+	// most go-driver requests to ArangoDB are executed against `_db/<db-name>/xxxx/yyy/`) URL pattern
+	// we can try to extract db-name to load-balance requests between endpoints
+
+	parts := strings.Split(strings.Trim(strings.TrimSpace(requestPath), "/"), "/")
+	if len(parts) >= 3 {
+		if parts[0] == "_db" {
+			return parts[1], nil
+		}
+	}
+	return fmt.Sprintf("%s_%s", requestMethod, requestPath), nil
+}