[grpctmclient] Add support for (bounded) per-host connection reuse

go/vt/grpcclient/client.go

@@ -19,6 +19,7 @@ limitations under the License.
 package grpcclient
 
 import (
+	"context"
 	"flag"
 	"time"
 
@@ -58,6 +59,16 @@ func RegisterGRPCDialOptions(grpcDialOptionsFunc func(opts []grpc.DialOption) ([
 // failFast is a non-optional parameter because callers are required to specify
 // what that should be.
 func Dial(target string, failFast FailFast, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
+	return DialContext(context.Background(), target, failFast, opts...)
+}
+
+// DialContext creates a grpc connection to the given target. Setup steps are
+// covered by the context deadline, and, if WithBlock is specified in the dial
+// options, connection establishment steps are covered by the context as well.
+//
+// failFast is a non-optional parameter because callers are required to specify
+// what that should be.
+func DialContext(ctx context.Context, target string, failFast FailFast, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
 	grpccommon.EnableTracingOpt()
 	newopts := []grpc.DialOption{
 		grpc.WithDefaultCallOptions(
@@ -98,7 +109,7 @@ func Dial(target string, failFast FailFast, opts ...grpc.DialOption) (*grpc.Clie
 
 	newopts = append(newopts, interceptors()...)
 
-	return grpc.Dial(target, newopts...)
+	return grpc.DialContext(ctx, target, newopts...)
 }
 
 func interceptors() []grpc.DialOption {

go/vt/vttablet/grpctmclient/cached_client.go

@@ -0,0 +1,331 @@
+/*
+Copyright 2021 The Vitess Authors.
+
+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.
+*/
+
+package grpctmclient
+
+import (
+	"context"
+	"flag"
+	"io"
+	"sort"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc"
+
+	"vitess.io/vitess/go/netutil"
+	"vitess.io/vitess/go/stats"
+	"vitess.io/vitess/go/sync2"
+	"vitess.io/vitess/go/vt/grpcclient"
+	"vitess.io/vitess/go/vt/vttablet/tmclient"
+
+	tabletmanagerservicepb "vitess.io/vitess/go/vt/proto/tabletmanagerservice"
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+)
+
+var (
+	defaultPoolCapacity = flag.Int("tablet_manager_grpc_connpool_size", 100, "number of tablets to keep tmclient connections open to")
+)
+
+func init() {
+	tmclient.RegisterTabletManagerClientFactory("grpc-cached", func() tmclient.TabletManagerClient {
+		return NewCachedConnClient(*defaultPoolCapacity)
+	})
+}
+
+// closeFunc allows a standalone function to implement io.Closer, similar to
+// how http.HandlerFunc allows standalone functions to implement http.Handler.
+type closeFunc func() error
+
+func (fn closeFunc) Close() error {
+	return fn()
+}
+
+var _ io.Closer = (*closeFunc)(nil)
+
+type cachedConn struct {
+	tabletmanagerservicepb.TabletManagerClient
+	cc *grpc.ClientConn
+
+	addr           string
+	lastAccessTime time.Time
+	refs           int
+}
+
+type cachedConnDialer struct {
+	m            sync.Mutex
+	conns        map[string]*cachedConn
+	evict        []*cachedConn
+	evictSorted  bool
+	connWaitSema *sync2.Semaphore
+	capacity     int
+}
+
+var dialerStats = struct {
+	ConnReuse    *stats.Gauge
+	ConnNew      *stats.Gauge
+	DialTimeouts *stats.Gauge
+	DialTimings  *stats.Timings
+}{
+	ConnReuse:    stats.NewGauge("tabletmanagerclient_cachedconn_reuse", "number of times a call to dial() was able to reuse an existing connection"),
+	ConnNew:      stats.NewGauge("tabletmanagerclient_cachedconn_new", "number of times a call to dial() resulted in a dialing a new grpc clientconn"),
+	DialTimeouts: stats.NewGauge("tabletmanagerclient_cachedconn_dial_timeouts", "number of context timeouts during dial()"),
+	DialTimings:  stats.NewTimings("tabletmanagerclient_cachedconn_dial_timings", "timings for various dial paths", "path", "cache_fast", "sema_fast", "sema_poll"),
+}
+
+// NewCachedConnClient returns a grpc Client that caches connections to the
+// different tablets.
+func NewCachedConnClient(capacity int) *Client {
+	dialer := &cachedConnDialer{
+		conns:        make(map[string]*cachedConn, capacity),
+		evict:        make([]*cachedConn, 0, capacity),
+		connWaitSema: sync2.NewSemaphore(capacity, 0),
+		capacity:     capacity,
+	}
+	return &Client{dialer}
+}
+
+var _ dialer = (*cachedConnDialer)(nil)
+
+func (dialer *cachedConnDialer) sortEvictionsLocked() {
+	if !dialer.evictSorted {
+		sort.Slice(dialer.evict, func(i, j int) bool {
+			left, right := dialer.evict[i], dialer.evict[j]
+			if left.refs == right.refs {
+				return right.lastAccessTime.After(left.lastAccessTime)
+			}
+			return right.refs > left.refs
+		})
+		dialer.evictSorted = true
+	}
+}
+
+func (dialer *cachedConnDialer) dial(ctx context.Context, tablet *topodatapb.Tablet) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
+	start := time.Now()
+	addr := getTabletAddr(tablet)
+
+	if client, closer, found, err := dialer.tryFromCache(addr, &dialer.m); found {
+		dialerStats.DialTimings.Add("cache_fast", time.Since(start))
+		return client, closer, err
+	}
+
+	if dialer.connWaitSema.TryAcquire() {
+		defer func() {
+			dialerStats.DialTimings.Add("sema_fast", time.Since(start))
+		}()
+
+		// Check if another goroutine managed to dial a conn for the same addr
+		// while we were waiting for the write lock. This is identical to the
+		// read-lock section above, except we release the connWaitSema if we
+		// are able to use the cache, allowing another goroutine to dial a new
+		// conn instead.
+		if client, closer, found, err := dialer.tryFromCache(addr, &dialer.m); found {
+			dialer.connWaitSema.Release()
+			return client, closer, err
+		}
+		return dialer.newdial(ctx, addr)
+	}
+
+	defer func() {
+		dialerStats.DialTimings.Add("sema_poll", time.Since(start))
+	}()
+
+	for {
+		select {
+		case <-ctx.Done():
+			dialerStats.DialTimeouts.Add(1)
+			return nil, nil, ctx.Err()
+		default:
+			if client, closer, found, err := dialer.pollOnce(ctx, addr); found {
+				return client, closer, err
+			}
+		}
+	}
+}
+
+// tryFromCache tries to get a connection from the cache, performing a redial
+// on that connection if it exists. It returns a TabletManagerClient impl, an
+// io.Closer, a flag to indicate whether a connection was found in the cache,
+// and an error, which is always nil.
+//
+// In addition to the addr being dialed, tryFromCache takes a sync.Locker which,
+// if not nil, will be used to wrap the lookup and redial in that lock. This
+// function can be called in situations where the conns map is locked
+// externally (like in pollOnce), so we do not want to manage the locks here. In
+// other cases (like in the cache_fast path of dial()), we pass in the dialer.m
+// to ensure we have a lock on the cache for the duration of the call.
+func (dialer *cachedConnDialer) tryFromCache(addr string, locker sync.Locker) (client tabletmanagerservicepb.TabletManagerClient, closer io.Closer, found bool, err error) {
+	if locker != nil {
+		locker.Lock()
+		defer locker.Unlock()
+	}
+
+	if conn, ok := dialer.conns[addr]; ok {
+		client, closer, err := dialer.redialLocked(conn)
+		return client, closer, ok, err
+	}
+
+	return nil, nil, false, nil
+}
+
+// pollOnce is called on each iteration of the polling loop in dial(). It:
+// - locks the conns cache for writes
+// - attempts to get a connection from the cache. If found, redial() it and exit.
+// - peeks at the head of the eviction queue. if the peeked conn has no refs, it
+//   is unused, and can be evicted to make room for the new connection to addr.
+//   If the peeked conn has refs, exit.
+// - pops the conn we just peeked from the queue, deletes it from the cache, and
+//   close the underlying ClientConn for that conn.
+// - attempt a newdial. if the newdial fails, it will release a slot on the
+//   connWaitSema, so another dial() call can successfully acquire it to dial
+//   a new conn. if the newdial succeeds, we will have evicted one conn, but
+//   added another, so the net change is 0, and no changes to the connWaitSema
+//   are made.
+//
+// It returns a TabletManagerClient impl, an io.Closer, a flag to indicate
+// whether the dial() poll loop should exit, and an error.
+func (dialer *cachedConnDialer) pollOnce(ctx context.Context, addr string) (client tabletmanagerservicepb.TabletManagerClient, closer io.Closer, found bool, err error) {
+	dialer.m.Lock()
+
+	if client, closer, found, err := dialer.tryFromCache(addr, nil); found {
+		dialer.m.Unlock()
+		return client, closer, found, err
+	}
+
+	dialer.sortEvictionsLocked()
+
+	conn := dialer.evict[0]
+	if conn.refs != 0 {
+		dialer.m.Unlock()
+		return nil, nil, false, nil
+	}
+
+	dialer.evict = dialer.evict[1:]
+	delete(dialer.conns, conn.addr)
+	conn.cc.Close()
+	dialer.m.Unlock()
+
+	client, closer, err = dialer.newdial(ctx, addr)
+	return client, closer, true, err
+}
+
+// newdial creates a new cached connection, and updates the cache and eviction
+// queue accordingly. If newdial fails to create the underlying
+// gRPC connection, it will make a call to Release the connWaitSema for other
+// newdial calls.
+//
+// It returns the three-tuple of client-interface, closer, and error that the
+// main dial func returns.
+func (dialer *cachedConnDialer) newdial(ctx context.Context, addr string) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
+	opt, err := grpcclient.SecureDialOption(*cert, *key, *ca, *name)
+	if err != nil {
+		dialer.connWaitSema.Release()
+		return nil, nil, err
+	}
+
+	cc, err := grpcclient.DialContext(ctx, addr, grpcclient.FailFast(false), opt)
+	if err != nil {
+		dialer.connWaitSema.Release()
+		return nil, nil, err
+	}
+
+	dialer.m.Lock()
+	defer dialer.m.Unlock()
+
+	if conn, existing := dialer.conns[addr]; existing {
+		// race condition: some other goroutine has dialed our tablet before we have;
+		// this is not great, but shouldn't happen often (if at all), so we're going to
+		// close this connection and reuse the existing one. by doing this, we can keep
+		// the actual Dial out of the global lock and significantly increase throughput
+		cc.Close()
+		dialer.connWaitSema.Release()
+		return dialer.redialLocked(conn)
+	}
+
+	dialerStats.ConnNew.Add(1)
+
+	conn := &cachedConn{
+		TabletManagerClient: tabletmanagerservicepb.NewTabletManagerClient(cc),
+		cc:                  cc,
+		lastAccessTime:      time.Now(),
+		refs:                1,
+		addr:                addr,
+	}
+
+	// NOTE: we deliberately do not set dialer.evictSorted=false here. Since
+	// cachedConns are evicted from the front of the queue, and we are appending
+	// to the end, if there is already a second evictable connection, it will be
+	// at the front of the queue, so we can speed up the edge case where we need
+	// to evict multiple connections in a row.
+	dialer.evict = append(dialer.evict, conn)
+	dialer.conns[addr] = conn
+
+	return dialer.connWithCloser(conn)
+}
+
+// redialLocked takes an already-dialed connection in the cache does all the
+// work of lending that connection out to one more caller. It returns the
+// three-tuple of client-interface, closer, and error that the main dial func
+// returns.
+func (dialer *cachedConnDialer) redialLocked(conn *cachedConn) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
+	dialerStats.ConnReuse.Add(1)
+	conn.lastAccessTime = time.Now()
+	conn.refs++
+	dialer.evictSorted = false
+	return dialer.connWithCloser(conn)
+}
+
+// connWithCloser returns the three-tuple expected by the main dial func, where
+// the closer handles the correct state management for updating the conns place
+// in the eviction queue.
+func (dialer *cachedConnDialer) connWithCloser(conn *cachedConn) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
+	return conn, closeFunc(func() error {
+		dialer.m.Lock()
+		defer dialer.m.Unlock()
+		conn.refs--
+		dialer.evictSorted = false
+		return nil
+	}), nil
+}
+
+// Close closes all currently cached connections, ***regardless of whether
+// those connections are in use***. Calling Close therefore will fail any RPCs
+// using currently lent-out connections, and, furthermore, will invalidate the
+// io.Closer that was returned for that connection from dialer.dial(). When
+// calling those io.Closers, they will still lock the dialer's mutex, and then
+// perform needless operations that will slow down dial throughput, but not
+// actually impact the correctness of the internal state of the dialer.
+//
+// As a result, while it is safe to reuse a cachedConnDialer after calling Close,
+// it will be less performant than getting a new one, either by calling
+// tmclient.TabletManagerClient() with
+// TabletManagerProtocol set to "grpc-cached", or by calling
+// grpctmclient.NewCachedConnClient directly.
+func (dialer *cachedConnDialer) Close() {
+	dialer.m.Lock()
+	defer dialer.m.Unlock()
+
+	for _, conn := range dialer.evict {
+		conn.cc.Close()
+		delete(dialer.conns, conn.addr)
+		dialer.connWaitSema.Release()
+	}
+	dialer.evict = make([]*cachedConn, 0, dialer.capacity)
+}
+
+func getTabletAddr(tablet *topodatapb.Tablet) string {
+	return netutil.JoinHostPort(tablet.Hostname, int32(tablet.PortMap["grpc"]))
+}