Add connection pool implementation (part 2)

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

connections/pool.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"math"
 	"math/rand"
+	"sync"
 	"time"
 
 	"github.com/nspcc-dev/neofs-api-go/pkg/client"
@@ -17,6 +18,8 @@ type PoolBuilderOptions struct {
 	NodeConnectionTimeout   time.Duration
 	NodeRequestTimeout      time.Duration
 	ClientRebalanceInterval time.Duration
+	weights                 []float64
+	connections             []*grpc.ClientConn
 }
 
 type PoolBuilder struct {
@@ -50,7 +53,9 @@ func (pb *PoolBuilder) Build(ctx context.Context, options *PoolBuilderOptions) (
 		}
 		cons[i] = con
 	}
-	return new(pb.weights, options.Key, cons)
+	options.weights = pb.weights
+	options.connections = cons
+	return new(ctx, options)
 }
 
 type Pool interface {
@@ -58,29 +63,70 @@ type Pool interface {
 }
 
 type pool struct {
-	generator *Generator
-	clients   []client.Client
+	lock    sync.RWMutex
+	sampler *Sampler
+	clients []client.Client
+	healthy []bool
 }
 
-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))
+func new(ctx context.Context, options *PoolBuilderOptions) (Pool, error) {
+	n := len(options.weights)
+	clients := make([]client.Client, n)
+	healthy := make([]bool, n)
+	for i, con := range options.connections {
+		c, err := client.New(client.WithDefaultPrivateKey(options.Key), client.WithGRPCConnection(con))
 		if err != nil {
 			return nil, err
 		}
 		clients[i] = c
+		healthy[i] = true
 	}
 	source := rand.NewSource(time.Now().UnixNano())
-	return &pool{
-		generator: NewGenerator(weights, source),
-		clients:   clients,
-	}, nil
+	pool := &pool{
+		sampler: NewSampler(options.weights, source),
+		clients: clients,
+		healthy: healthy,
+	}
+	go func() {
+		ticker := time.NewTimer(options.ClientRebalanceInterval)
+		for range ticker.C {
+			ok := true
+			for i, client := range pool.clients {
+				func() {
+					tctx, c := context.WithTimeout(ctx, options.NodeRequestTimeout)
+					defer c()
+					if _, err := client.EndpointInfo(tctx); err != nil {
+						ok = false
+					}
+					pool.lock.Lock()
+					pool.healthy[i] = ok
+					pool.lock.Unlock()
+				}()
+			}
+			ticker.Reset(options.ClientRebalanceInterval)
+		}
+	}()
+	return pool, nil
 }
 
 func (p *pool) Client() client.Client {
+	p.lock.RLock()
+	defer p.lock.RUnlock()
 	if len(p.clients) == 1 {
-		return p.clients[0]
+		if p.healthy[0] {
+			return p.clients[0]
+		}
+		return nil
 	}
-	return p.clients[p.generator.Next()]
+	var i *int = nil
+	for k := 0; k < 10; k++ {
+		i_ := p.sampler.Next()
+		if p.healthy[i_] {
+			i = &i_
+		}
+	}
+	if i != nil {
+		return p.clients[*i]
+	}
+	return nil
 }