@@ -3,8 +3,8 @@ package transport
3
3
import (
4
4
"context"
5
5
"fmt"
6
- "log"
7
6
"math"
7
+ "math/bits"
8
8
"net"
9
9
"time"
10
10
@@ -126,10 +126,11 @@ func (p *ConnPool) closeAll() {
126
126
}
127
127
128
128
type PoolRefiller struct {
129
- addr string
130
- pool ConnPool
131
- cfg ConnConfig
132
- active int
129
+ addr string
130
+ pool ConnPool
131
+ cfg ConnConfig
132
+ active int
133
+ fillRandomly bool
133
134
}
134
135
135
136
func (r * PoolRefiller ) init (ctx context.Context , host string ) error {
@@ -158,13 +159,17 @@ func (r *PoolRefiller) init(ctx context.Context, host string) error {
158
159
if v , ok := s .Options [ScyllaShardAwarePortSSL ]; ok {
159
160
r .addr = net .JoinHostPort (host , v [0 ])
160
161
} else {
161
- return fmt .Errorf ("missing encrypted shard aware port information %v" , s .Options )
162
+ r .cfg .Logger .Warnf ("missing encrypted shard aware port information %v; falling back to random shard discovery" , s .Options )
163
+ r .addr = net .JoinHostPort (host , r .cfg .DefaultPort )
164
+ r .fillRandomly = true
162
165
}
163
166
} else {
164
167
if v , ok := s .Options [ScyllaShardAwarePort ]; ok {
165
168
r .addr = net .JoinHostPort (host , v [0 ])
166
169
} else {
167
- return fmt .Errorf ("missing shard aware port information %v" , s .Options )
170
+ r .cfg .Logger .Warnf ("missing shard aware port information %v; falling back to random shard discovery" , s .Options )
171
+ r .addr = net .JoinHostPort (host , r .cfg .DefaultPort )
172
+ r .fillRandomly = true
168
173
}
169
174
}
170
175
@@ -222,17 +227,39 @@ func (r *PoolRefiller) loop(ctx context.Context) {
222
227
}
223
228
}
224
229
230
+ // storeShard assumes conn is non-nil.
231
+ func (r * PoolRefiller ) storeShard (conn * Conn , span span ) bool {
232
+ if r .pool .loadConn (conn .Shard ()) != nil {
233
+ if r .pool .connObs != nil {
234
+ r .pool .connObs .OnConnect (ConnectEvent {
235
+ ConnEvent : conn .Event (),
236
+ span : span ,
237
+ Err : fmt .Errorf ("shard already in pool" ),
238
+ })
239
+ }
240
+ conn .Close ()
241
+ return false
242
+ }
243
+
244
+ if r .pool .connObs != nil {
245
+ r .pool .connObs .OnConnect (ConnectEvent {ConnEvent : conn .Event (), span : span })
246
+ }
247
+ conn .setOnClose (r .onConnClose )
248
+ r .pool .storeConn (conn )
249
+ r .active ++
250
+ return true
251
+ }
252
+
225
253
func (r * PoolRefiller ) fill (ctx context.Context ) {
226
- if ! r . needsFilling () {
227
- return
254
+ if r . fillRandomly {
255
+ r . fillRandom ( ctx )
228
256
}
229
257
230
258
si := ShardInfo {
231
259
NrShards : uint16 (r .pool .nrShards ),
232
260
MsbIgnore : r .pool .msbIgnore ,
233
261
}
234
-
235
- for i := 0 ; i < r .pool .nrShards ; i ++ {
262
+ for i := 0 ; r .needsFilling () && i < r .pool .nrShards ; i ++ {
236
263
if r .pool .loadConn (i ) != nil {
237
264
continue
238
265
}
@@ -245,23 +272,14 @@ func (r *PoolRefiller) fill(ctx context.Context) {
245
272
if r .pool .connObs != nil {
246
273
r .pool .connObs .OnConnect (ConnectEvent {ConnEvent : ConnEvent {Addr : r .addr , Shard : si .Shard }, span : span , Err : err })
247
274
}
248
- if conn != nil {
249
- conn .Close ()
250
- }
251
275
continue
252
276
}
253
- if r .pool .connObs != nil {
254
- r .pool .connObs .OnConnect (ConnectEvent {ConnEvent : conn .Event (), span : span })
255
- }
256
277
278
+ r .storeShard (conn , span )
257
279
if conn .Shard () != i {
258
- log .Fatalf ("opened conn to wrong shard: expected %d got %d" , i , conn .Shard ())
259
- }
260
- conn .setOnClose (r .onConnClose )
261
- r .pool .storeConn (conn )
262
- r .active ++
263
-
264
- if ! r .needsFilling () {
280
+ r .cfg .Logger .Warnf ("opened conn to wrong shard: expected %d got %d; falling back to random discovery" , i , conn .Shard ())
281
+ r .fillRandomly = true
282
+ r .fillRandom (ctx )
265
283
return
266
284
}
267
285
}
@@ -270,3 +288,21 @@ func (r *PoolRefiller) fill(ctx context.Context) {
270
288
func (r * PoolRefiller ) needsFilling () bool {
271
289
return r .active < r .pool .nrShards
272
290
}
291
+
292
+ func (r * PoolRefiller ) fillRandom (ctx context.Context ) {
293
+ // https://en.wikipedia.org/wiki/Coupon_collector%27s_problem
294
+ maxTries := r .pool .nrShards * bits .Len (uint (r .pool .nrShards ))
295
+ for try := 0 ; r .needsFilling () && try < maxTries ; try ++ {
296
+ span := startSpan ()
297
+ conn , err := OpenConn (ctx , r .addr , nil , r .cfg )
298
+ span .stop ()
299
+ if err != nil {
300
+ if r .pool .connObs != nil {
301
+ r .pool .connObs .OnConnect (ConnectEvent {ConnEvent : ConnEvent {Addr : r .addr , Shard : UnknownShard }, span : span , Err : err })
302
+ }
303
+ continue
304
+ }
305
+
306
+ r .storeShard (conn , span )
307
+ }
308
+ }
0 commit comments