Add connection pool implementation (part 1)

Signed-off-by: Pavel Korotkov <pavel@nspcc.ru>

connections/generator.go

@@ -0,0 +1,76 @@
+package connections
+
+import "math/rand"
+
+// https://www.keithschwarz.com/darts-dice-coins/
+type Generator struct {
+	randomGenerator *rand.Rand
+	probabilities   []float64
+	alias           []int
+}
+
+type workList []int
+
+func (wl *workList) push(e int) {
+	*wl = append(*wl, e)
+}
+
+func (wl *workList) pop() int {
+	l := len(*wl) - 1
+	n := (*wl)[l]
+	*wl = (*wl)[:l]
+	return n
+}
+
+func NewGenerator(probabilities []float64, source rand.Source) *Generator {
+	generator := &Generator{}
+	var (
+		small workList
+		large workList
+	)
+	n := len(probabilities)
+	generator.randomGenerator = rand.New(source)
+	generator.probabilities = make([]float64, n)
+	generator.alias = make([]int, n)
+	// Compute scaled probabilities.
+	p := make([]float64, n)
+	for i := 0; i < n; i++ {
+		p[i] = probabilities[i] * float64(n)
+	}
+	for i, pi := range p {
+		if pi < 1 {
+			small = append(small, i)
+		} else {
+			large = append(large, i)
+		}
+	}
+	for len(large) > 0 && len(small) > 0 {
+		l, g := small.pop(), large.pop()
+		generator.probabilities[l] = p[l]
+		generator.alias[l] = g
+		p[g] = (p[g] + p[l]) - 1
+		if p[g] < 1 {
+			small.push(g)
+		} else {
+			large.push(g)
+		}
+	}
+	for len(large) > 0 {
+		g := large.pop()
+		generator.probabilities[g] = 1
+	}
+	for len(small) > 0 {
+		l := small.pop()
+		generator.probabilities[l] = 1
+	}
+	return generator
+}
+
+func (g *Generator) Next() int {
+	n := len(g.alias)
+	i := g.randomGenerator.Intn(n)
+	if g.randomGenerator.Float64() < g.probabilities[i] {
+		return i
+	}
+	return g.alias[i]
+}

connections/pool.go

@@ -0,0 +1,86 @@
+package connections
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"errors"
+	"math"
+	"math/rand"
+	"time"
+
+	"github.com/nspcc-dev/neofs-api-go/pkg/client"
+	"google.golang.org/grpc"
+)
+
+type PoolBuilderOptions struct {
+	Key                     *ecdsa.PrivateKey
+	NodeConnectionTimeout   time.Duration
+	NodeRequestTimeout      time.Duration
+	ClientRebalanceInterval time.Duration
+}
+
+type PoolBuilder struct {
+	addresses []string
+	weights   []float64
+}
+
+func (pb *PoolBuilder) AddNode(address string, weight float64) *PoolBuilder {
+	pb.addresses = append(pb.addresses, address)
+	pb.weights = append(pb.weights, weight)
+	return pb
+}
+
+func (pb *PoolBuilder) Build(ctx context.Context, options *PoolBuilderOptions) (Pool, error) {
+	totalWeight := 0.0
+	for _, w := range pb.weights {
+		totalWeight += w
+	}
+	if math.Abs(totalWeight-1.0) >= 1e-4 {
+		return nil, errors.New("total weight must be equal to unity")
+	}
+	var cons = make([]*grpc.ClientConn, len(pb.addresses))
+	for i, address := range pb.addresses {
+		con, err := func() (*grpc.ClientConn, error) {
+			toctx, c := context.WithTimeout(ctx, options.NodeConnectionTimeout)
+			defer c()
+			return grpc.DialContext(toctx, address, grpc.WithInsecure(), grpc.WithBlock())
+		}()
+		if err != nil {
+			return nil, err
+		}
+		cons[i] = con
+	}
+	return new(pb.weights, options.Key, cons)
+}
+
+type Pool interface {
+	Client() client.Client
+}
+
+type pool struct {
+	generator *Generator
+	clients   []client.Client
+}
+
+func new(weights []float64, key *ecdsa.PrivateKey, connections []*grpc.ClientConn) (Pool, error) {
+	clients := make([]client.Client, len(weights))
+	for i, con := range connections {
+		c, err := client.New(client.WithDefaultPrivateKey(key), client.WithGRPCConnection(con))
+		if err != nil {
+			return nil, err
+		}
+		clients[i] = c
+	}
+	source := rand.NewSource(time.Now().UnixNano())
+	return &pool{
+		generator: NewGenerator(weights, source),
+		clients:   clients,
+	}, nil
+}
+
+func (p *pool) Client() client.Client {
+	if len(p.clients) == 1 {
+		return p.clients[0]
+	}
+	return p.clients[p.generator.Next()]
+}