[vtadmin] cluster rpc pools

go/pools/rpc_pool.go

@@ -0,0 +1,91 @@
+/*
+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 pools
+
+import (
+	"context"
+	"time"
+)
+
+// RPCPool is a specialized version of the ResourcePool, for bounding concurrent
+// access to making RPC calls.
+//
+// Whether you use this, or a sync2.Semaphore to gate RPCs (or shared access to
+// any shared resource) depends on what timeout semantics you need. A sync2.Semaphore
+// takes a global timeout, which applies to calls to Acquire(); if you need to
+// respect a context's deadline, you can call AcquireContext(context.Context),
+// but that ignores the global timeout. Conversely, an RPCPool provides only
+// one method of acquisition, Acquire(context.Context), which always uses the
+// lower of the pool-global timeout or the context deadline.
+type RPCPool struct {
+	rp          *ResourcePool
+	waitTimeout time.Duration
+}
+
+// NewRPCPool returns an RPCPool with the given size and wait timeout. A zero
+// timeout will be ignored on calls to Acquire, and only the context deadline
+// will be used. If a logWait function is provided, it will be called whenever
+// a call to Acquire has to wait for a resource, but only after a successful
+// wait (meaning if we hit a timeout before a resource becomes available, it
+// will not be called).
+func NewRPCPool(size int, waitTimeout time.Duration, logWait func(time.Time)) *RPCPool {
+	return &RPCPool{
+		rp:          NewResourcePool(rpcResourceFactory, size, size, 0, size, logWait),
+		waitTimeout: waitTimeout,
+	}
+}
+
+// Acquire acquires one slot in the RPCPool. If a slot is not immediately
+// available, it will block until one becomes available or until a timeout is
+// reached. The lower of the context deadline and the pool's waitTimeout will
+// be used as the timeout, except when the pool's waitTimeout is zero, in which
+// case the context deadline will always serve as the overall timeout.
+//
+// It returns nil on successful acquisition, and an error if a timeout occurred
+// before a slot became available.
+//
+// Note: For every successful call to Acquire, the caller must make a
+// corresponding call to Release.
+func (pool *RPCPool) Acquire(ctx context.Context) error {
+	if pool.waitTimeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, pool.waitTimeout)
+		defer cancel()
+	}
+
+	_, err := pool.rp.Get(ctx)
+	return err
+}
+
+// Release frees a slot in the pool. It must only be called after a successful
+// call to Acquire.
+func (pool *RPCPool) Release() { pool.rp.Put(rpc) }
+
+// Close empties the pool, preventing further Acquire calls from succeeding.
+// It waits for all slots to be freed via Release.
+func (pool *RPCPool) Close() { pool.rp.Close() }
+
+type _rpc struct{}
+
+var rpc = &_rpc{}
+
+// Close implements Resource for _rpc.
+func (*_rpc) Close() {}
+
+// we only ever return the same rpc pointer. it's used as a sentinel and is
+// only used internally so using the same one over and over doesn't matter.
+func rpcResourceFactory(ctx context.Context) (Resource, error) { return rpc, nil }

go/pools/rpc_pool_test.go

@@ -0,0 +1,83 @@
+/*
+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 pools
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRPCPool(t *testing.T) {
+	t.Parallel()
+
+	pool := NewRPCPool(1, time.Millisecond*100, nil)
+
+	err := pool.Acquire(context.Background())
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*5)
+	err = pool.Acquire(ctx)
+	cancel()
+	assert.Error(t, err)
+
+	pool.Release()
+	err = pool.Acquire(context.Background())
+	require.NoError(t, err)
+
+	t.Run("context timeout exceeded", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*50)
+		defer cancel()
+
+		err := pool.Acquire(ctx)
+		assert.Error(t, err)
+
+		select {
+		case <-ctx.Done():
+		default:
+			assert.Fail(t, "calling context should be done after failed Acquire")
+		}
+	})
+
+	t.Run("pool timeout exceeded", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*500)
+		defer cancel()
+
+		err := pool.Acquire(ctx)
+		assert.Error(t, err)
+
+		select {
+		case <-ctx.Done():
+			assert.Fail(t, "calling context should not be done after failed Acquire")
+		default:
+		}
+	})
+}
+
+func BenchmarkRPCPoolAllocs(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		NewRPCPool(1000, 0, nil)
+		// p.Close()
+	}
+}