fix: improve client pool concurrency safety and performance (#67)

Author: MattDevy

elastictransport/client_pool_benchmark_test.go

@@ -0,0 +1,180 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you 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.
+
+//go:build !integration
+// +build !integration
+
+package elastictransport
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"testing"
+)
+
+// benchTransport returns a minimal response without hitting the network.
+// Each call returns a fresh *http.Response so parallel goroutines don't share
+// mutable state.
+type benchTransport struct{}
+
+func (t *benchTransport) RoundTrip(*http.Request) (*http.Response, error) {
+	return &http.Response{
+		StatusCode:    200,
+		Body:          io.NopCloser(strings.NewReader("")),
+		ContentLength: 0,
+		Header:        http.Header{},
+	}, nil
+}
+
+// benchmarkPerform exercises Client.Perform under varying parallelism factors.
+// The sub-benchmarks vary the value passed to b.SetParallelism, which scales
+// parallel work relative to GOMAXPROCS rather than creating an exact number
+// of goroutines.
+func benchmarkPerform(b *testing.B, name string, client *Client) {
+	b.Helper()
+	b.Run(name, func(b *testing.B) {
+		for _, p := range []int{1, 10, 100, 1000} {
+			b.Run(fmt.Sprintf("parallelism-%d", p), func(b *testing.B) {
+				b.SetParallelism(p)
+				b.RunParallel(func(pb *testing.PB) {
+					for pb.Next() {
+						req, _ := http.NewRequest("GET", "/", nil)
+						res, err := client.Perform(req)
+						if err != nil {
+							b.Fatalf("Perform: %s", err)
+						}
+						if err := res.Body.Close(); err != nil {
+							b.Fatalf("Close response body: %s", err)
+						}
+					}
+				})
+			})
+		}
+	})
+}
+
+func BenchmarkClientPoolAccess(b *testing.B) {
+	b.ReportAllocs()
+
+	b.Run("SingleConnectionPool", func(b *testing.B) {
+		u, _ := url.Parse("http://localhost:9200")
+		tp, err := New(Config{
+			URLs:         []*url.URL{u},
+			Transport:    &benchTransport{},
+			DisableRetry: true,
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+		benchmarkPerform(b, "Perform", tp)
+	})
+
+	b.Run("StatusConnectionPool", func(b *testing.B) {
+		var urls []*url.URL
+		for i := 0; i < 10; i++ {
+			u, _ := url.Parse(fmt.Sprintf("http://node%d:9200", i))
+			urls = append(urls, u)
+		}
+		tp, err := New(Config{
+			URLs:         urls,
+			Transport:    &benchTransport{},
+			DisableRetry: true,
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+		benchmarkPerform(b, "Perform", tp)
+	})
+
+	b.Run("CustomPool/SynchronizedWrapper", func(b *testing.B) {
+		var urls []*url.URL
+		for i := 0; i < 10; i++ {
+			u, _ := url.Parse(fmt.Sprintf("http://node%d:9200", i))
+			urls = append(urls, u)
+		}
+		tp, err := New(Config{
+			URLs: urls,
+			ConnectionPoolFunc: func(conns []*Connection, sel Selector) ConnectionPool {
+				return &benchCustomPool{conns: conns}
+			},
+			Transport:    &benchTransport{},
+			DisableRetry: true,
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+		benchmarkPerform(b, "Perform", tp)
+	})
+
+	b.Run("CustomPool/ConcurrentSafeOptIn", func(b *testing.B) {
+		var urls []*url.URL
+		for i := 0; i < 10; i++ {
+			u, _ := url.Parse(fmt.Sprintf("http://node%d:9200", i))
+			urls = append(urls, u)
+		}
+		tp, err := New(Config{
+			URLs: urls,
+			ConnectionPoolFunc: func(conns []*Connection, sel Selector) ConnectionPool {
+				return &benchConcurrentSafeCustomPool{
+					benchCustomPool: benchCustomPool{conns: conns},
+				}
+			},
+			Transport:    &benchTransport{},
+			DisableRetry: true,
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+		benchmarkPerform(b, "Perform", tp)
+	})
+}
+
+type benchCustomPool struct {
+	mu    sync.Mutex
+	conns []*Connection
+	curr  int
+}
+
+func (p *benchCustomPool) Next() (*Connection, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	c := p.conns[p.curr%len(p.conns)]
+	p.curr++
+	return c, nil
+}
+
+func (p *benchCustomPool) OnSuccess(*Connection) error { return nil }
+func (p *benchCustomPool) OnFailure(*Connection) error { return nil }
+func (p *benchCustomPool) URLs() []*url.URL {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	out := make([]*url.URL, len(p.conns))
+	for i, c := range p.conns {
+		out[i] = c.URL
+	}
+	return out
+}
+
+type benchConcurrentSafeCustomPool struct {
+	benchCustomPool
+}
+
+func (p *benchConcurrentSafeCustomPool) ConcurrentSafe() {}

elastictransport/connection.go

@@ -39,6 +39,12 @@ var (
 	_ ConnectionPool          = (*statusConnectionPool)(nil)
 	_ UpdatableConnectionPool = (*statusConnectionPool)(nil)
 	_ CloseableConnectionPool = (*statusConnectionPool)(nil)
+	_ ConnectionPool          = (*synchronizedPool)(nil)
+	_ CloseableConnectionPool = (*synchronizedPool)(nil)
+	_ connectionable          = (*synchronizedPool)(nil)
+	_ ConnectionPool          = (*synchronizedUpdatablePool)(nil)
+	_ CloseableConnectionPool = (*synchronizedUpdatablePool)(nil)
+	_ UpdatableConnectionPool = (*synchronizedUpdatablePool)(nil)
 	_ Selector                = (*roundRobinSelector)(nil)
 )
 
@@ -59,11 +65,107 @@ type UpdatableConnectionPool interface {
 	Update([]*Connection) error // Update injects newly found nodes in the cluster.
 }
 
+// ConcurrentSafeConnectionPool marks a custom pool as safe for concurrent use
+// by the transport. Pools returned by ConnectionPoolFunc that implement this
+// interface are not wrapped by synchronizedPool.
+type ConcurrentSafeConnectionPool interface {
+	ConnectionPool
+	ConcurrentSafe()
+}
+
 // CloseableConnectionPool defines the interface for the connection pool that can be closed.
 type CloseableConnectionPool interface {
 	Close(context.Context) error
 }
 
+// synchronizedPool wraps a ConnectionPool with a mutex to serialize all method
+// calls. Used to wrap user-provided pools from ConnectionPoolFunc that may not
+// be safe for concurrent use.
+//
+// Use newSynchronizedPool to construct; it returns a synchronizedUpdatablePool
+// when the inner pool implements UpdatableConnectionPool.
+type synchronizedPool struct {
+	mu   sync.Mutex
+	pool ConnectionPool
+}
+
+// synchronizedUpdatablePool extends synchronizedPool for inner pools that
+// support in-place updates. Discovery will prefer Update() over full pool
+// replacement when this interface is present.
+type synchronizedUpdatablePool struct {
+	synchronizedPool
+}
+
+// newSynchronizedPool wraps pool in a synchronized adapter by default.
+// ConcurrentSafeConnectionPool implementations are returned as-is.
+// If pool implements UpdatableConnectionPool the returned value will too, so
+// discovery can update the pool in place rather than replacing it.
+func newSynchronizedPool(pool ConnectionPool) ConnectionPool {
+	if _, ok := pool.(ConcurrentSafeConnectionPool); ok {
+		return pool
+	}
+
+	if _, ok := pool.(UpdatableConnectionPool); ok {
+		return &synchronizedUpdatablePool{synchronizedPool{pool: pool}}
+	}
+	return &synchronizedPool{pool: pool}
+}
+
+func (sp *synchronizedPool) Next() (*Connection, error) {
+	sp.mu.Lock()
+	defer sp.mu.Unlock()
+	return sp.pool.Next()
+}
+
+func (sp *synchronizedPool) OnSuccess(c *Connection) error {
+	sp.mu.Lock()
+	defer sp.mu.Unlock()
+	return sp.pool.OnSuccess(c)
+}
+
+func (sp *synchronizedPool) OnFailure(c *Connection) error {
+	sp.mu.Lock()
+	defer sp.mu.Unlock()
+	return sp.pool.OnFailure(c)
+}
+
+func (sp *synchronizedPool) URLs() []*url.URL {
+	sp.mu.Lock()
+	defer sp.mu.Unlock()
+	return sp.pool.URLs()
+}
+
+func (sp *synchronizedPool) Close(ctx context.Context) error {
+	sp.mu.Lock()
+	cp, ok := sp.pool.(CloseableConnectionPool)
+	sp.mu.Unlock()
+	if !ok {
+		return nil
+	}
+	// Avoid holding the wrapper lock during close; in-flight callbacks also take
+	// this lock and may need to complete before the inner pool can fully close.
+	return cp.Close(ctx)
+}
+
+func (sp *synchronizedPool) connections() []*Connection {
+	sp.mu.Lock()
+	defer sp.mu.Unlock()
+	if cp, ok := sp.pool.(connectionable); ok {
+		return cp.connections()
+	}
+	return nil
+}
+
+func (sp *synchronizedUpdatablePool) Update(conns []*Connection) error {
+	sp.mu.Lock()
+	defer sp.mu.Unlock()
+	up, ok := sp.pool.(UpdatableConnectionPool)
+	if !ok {
+		return fmt.Errorf("inner pool %T does not implement UpdatableConnectionPool", sp.pool)
+	}
+	return up.Update(conns)
+}
+
 // Connection represents a connection to a node.
 type Connection struct {
 	sync.Mutex
@@ -240,8 +342,10 @@ func (cp *statusConnectionPool) OnFailure(c *Connection) error {
 }
 
 // Update merges the existing live and dead connections with the latest nodes discovered from the cluster.
-// ConnectionPool must be locked before calling.
 func (cp *statusConnectionPool) Update(connections []*Connection) error {
+	cp.Lock()
+	defer cp.Unlock()
+
 	if len(connections) == 0 {
 		return errors.New("no connections provided, connection pool left untouched")
 	}
@@ -354,7 +458,9 @@ func (cp *statusConnectionPool) isClosed() bool {
 }
 
 func (cp *statusConnectionPool) connections() []*Connection {
-	var conns []*Connection
+	cp.Lock()
+	defer cp.Unlock()
+	conns := make([]*Connection, 0, len(cp.live)+len(cp.dead))
 	conns = append(conns, cp.live...)
 	conns = append(conns, cp.dead...)
 	return conns