pool: fallback to random search when shard-aware port is missing

Author: Kulezi

transport/pool.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"math"
+	"math/bits"
 	"net"
 	"time"
 
@@ -126,10 +127,11 @@ func (p *ConnPool) closeAll() {
 }
 
 type PoolRefiller struct {
-	addr   string
-	pool   ConnPool
-	cfg    ConnConfig
-	active int
+	addr         string
+	pool         ConnPool
+	cfg          ConnConfig
+	active       int
+	fillRandomly bool
 }
 
 func (r *PoolRefiller) init(ctx context.Context, host string) error {
@@ -158,13 +160,17 @@ func (r *PoolRefiller) init(ctx context.Context, host string) error {
 		if v, ok := s.Options[ScyllaShardAwarePortSSL]; ok {
 			r.addr = net.JoinHostPort(host, v[0])
 		} else {
-			return fmt.Errorf("missing encrypted shard aware port information %v", s.Options)
+			r.cfg.Logger.Warnf("missing encrypted shard aware port information %v; falling back to random shard discovery", s.Options)
+			r.addr = net.JoinHostPort(host, r.cfg.DefaultPort)
+			r.fillRandomly = true
 		}
 	} else {
 		if v, ok := s.Options[ScyllaShardAwarePort]; ok {
 			r.addr = net.JoinHostPort(host, v[0])
 		} else {
-			return fmt.Errorf("missing shard aware port information %v", s.Options)
+			r.cfg.Logger.Warnf("missing shard aware port information %v; falling back to random shard discovery", s.Options)
+			r.addr = net.JoinHostPort(host, r.cfg.DefaultPort)
+			r.fillRandomly = true
 		}
 	}
 
@@ -226,7 +232,9 @@ func (r *PoolRefiller) fill(ctx context.Context) {
 	if !r.needsFilling() {
 		return
 	}
-
+	if r.fillRandomly {
+		r.fillRandom(ctx)
+	}
 	si := ShardInfo{
 		NrShards:  uint16(r.pool.nrShards),
 		MsbIgnore: r.pool.msbIgnore,
@@ -270,3 +278,41 @@ func (r *PoolRefiller) fill(ctx context.Context) {
 func (r *PoolRefiller) needsFilling() bool {
 	return r.active < r.pool.nrShards
 }
+
+func (r *PoolRefiller) fillRandom(ctx context.Context) {
+	// https://en.wikipedia.org/wiki/Coupon_collector%27s_problem
+	maxTries := r.pool.nrShards * bits.Len(uint(r.pool.nrShards))
+	r.cfg.Logger.Infof("trying random shards %d times", maxTries)
+	for try := 0; r.needsFilling() && try < maxTries; try++ {
+		span := startSpan()
+		conn, err := OpenConn(ctx, r.addr, nil, r.cfg)
+		span.stop()
+		if err != nil {
+			if r.pool.connObs != nil {
+				r.pool.connObs.OnConnect(ConnectEvent{ConnEvent: ConnEvent{Addr: r.addr, Shard: 0}, span: span, Err: err})
+			}
+			if conn != nil {
+				conn.Close()
+			}
+			continue
+		}
+
+		if r.pool.loadConn(conn.Shard()) != nil {
+			if r.pool.connObs != nil {
+				r.pool.connObs.OnConnect(ConnectEvent{
+					ConnEvent: ConnEvent{Addr: r.addr, Shard: uint16(conn.Shard())},
+					span:      span,
+					Err:       fmt.Errorf("shard already in pool")})
+			}
+			conn.Close()
+			continue
+		}
+
+		if r.pool.connObs != nil {
+			r.pool.connObs.OnConnect(ConnectEvent{ConnEvent: conn.Event(), span: span})
+		}
+		conn.setOnClose(r.onConnClose)
+		r.pool.storeConn(conn)
+		r.active++
+	}
+}