Skip to content

Commit ca9985d

Browse files
pappzpappz
committed
Add options for manager to avoid race in tests
1 parent 0781908 commit ca9985d

File tree

4 files changed

+84
-28
lines changed

4 files changed

+84
-28
lines changed

‎client/internal/connect.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,9 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
253253
return wrapErr(err)
254254
}
255255

256-
relayManager := relayClient.NewManager(engineCtx, relayURLs, myPrivateKey.PublicKey().String(), engineConfig.MTU)
256+
relayManager := relayClient.NewManager(engineCtx, relayURLs, myPrivateKey.PublicKey().String(), &relayClient.ManagerOpts{
257+
MTU: engineConfig.MTU,
258+
})
257259
c.statusRecorder.SetRelayMgr(relayManager)
258260
if len(relayURLs) > 0 {
259261
if token != nil {

‎client/internal/engine_test.go‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ import (
1414

1515
"github.com/golang/mock/gomock"
1616
"github.com/google/uuid"
17-
"github.com/netbirdio/netbird/client/internal/stdnet"
1817
log "github.com/sirupsen/logrus"
1918
"github.com/stretchr/testify/assert"
2019
"github.com/stretchr/testify/require"
@@ -25,7 +24,10 @@ import (
2524
"google.golang.org/grpc"
2625
"google.golang.org/grpc/keepalive"
2726

27+
"github.com/netbirdio/netbird/client/internal/stdnet"
28+
2829
"github.com/netbirdio/management-integrations/integrations"
30+
2931
"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
3032
"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
3133
nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
@@ -227,7 +229,7 @@ func TestEngine_SSH(t *testing.T) {
227229
ctx, cancel := context.WithCancel(context.Background())
228230
defer cancel()
229231

230-
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU)
232+
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), &relayClient.ManagerOpts{MTU: iface.DefaultMTU})
231233
engine := NewEngine(
232234
ctx, cancel,
233235
&signal.MockClient{},
@@ -373,7 +375,7 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
373375
ctx, cancel := context.WithCancel(context.Background())
374376
defer cancel()
375377

376-
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU)
378+
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), &relayClient.ManagerOpts{MTU: iface.DefaultMTU})
377379
engine := NewEngine(
378380
ctx, cancel,
379381
&signal.MockClient{},
@@ -600,7 +602,7 @@ func TestEngine_Sync(t *testing.T) {
600602
}
601603
return nil
602604
}
603-
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU)
605+
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), &relayClient.ManagerOpts{MTU: iface.DefaultMTU})
604606
engine := NewEngine(ctx, cancel, &signal.MockClient{}, &mgmt.MockClient{SyncFunc: syncFunc}, relayMgr, &EngineConfig{
605607
WgIfaceName: "utun103",
606608
WgAddr: "100.64.0.1/24",
@@ -765,7 +767,7 @@ func TestEngine_UpdateNetworkMapWithRoutes(t *testing.T) {
765767
wgIfaceName := fmt.Sprintf("utun%d", 104+n)
766768
wgAddr := fmt.Sprintf("100.66.%d.1/24", n)
767769

768-
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU)
770+
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), &relayClient.ManagerOpts{MTU: iface.DefaultMTU})
769771
engine := NewEngine(ctx, cancel, &signal.MockClient{}, &mgmt.MockClient{}, relayMgr, &EngineConfig{
770772
WgIfaceName: wgIfaceName,
771773
WgAddr: wgAddr,
@@ -967,7 +969,7 @@ func TestEngine_UpdateNetworkMapWithDNSUpdate(t *testing.T) {
967969
wgIfaceName := fmt.Sprintf("utun%d", 104+n)
968970
wgAddr := fmt.Sprintf("100.66.%d.1/24", n)
969971

970-
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU)
972+
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), &relayClient.ManagerOpts{MTU: iface.DefaultMTU})
971973
engine := NewEngine(ctx, cancel, &signal.MockClient{}, &mgmt.MockClient{}, relayMgr, &EngineConfig{
972974
WgIfaceName: wgIfaceName,
973975
WgAddr: wgAddr,
@@ -1499,7 +1501,7 @@ func createEngine(ctx context.Context, cancel context.CancelFunc, setupKey strin
14991501
MTU: iface.DefaultMTU,
15001502
}
15011503

1502-
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU)
1504+
relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), &relayClient.ManagerOpts{MTU: iface.DefaultMTU})
15031505
e, err := NewEngine(ctx, cancel, signalClient, mgmtClient, relayMgr, conf, MobileDependency{}, peer.NewRecorder("https://mgm"), nil), nil
15041506
e.ctx = ctx
15051507
return e, err

‎shared/relay/client/manager.go‎

Lines changed: 55 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,15 @@ import (
1414
relayAuth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
1515
)
1616

17-
var (
18-
relayCleanupInterval = 60 * time.Second
19-
keepUnusedServerTime = 5 * time.Second
17+
const (
18+
defaultRelayCleanupInterval = 60 * time.Second
19+
defaultKeepUnusedServerTime = 5 * time.Second
20+
defaultMTU = 1280
21+
minMTU = 1280
22+
maxMTU = 65535
23+
)
2024

25+
var (
2126
ErrRelayClientNotConnected = fmt.Errorf("relay client not connected")
2227
)
2328

@@ -64,14 +69,55 @@ type Manager struct {
6469
onReconnectedListenerFn func()
6570
listenerLock sync.Mutex
6671

67-
mtu uint16
72+
cleanupInterval time.Duration
73+
unusedServerTime time.Duration
74+
mtu uint16
75+
}
76+
77+
// ManagerOpts contains optional configuration for Manager
78+
type ManagerOpts struct {
79+
// CleanupInterval is the interval for cleaning up unused relay connections.
80+
// If zero, defaults to defaultRelayCleanupInterval.
81+
CleanupInterval time.Duration
82+
// UnusedServerTime is the time to wait before closing unused relay connections.
83+
// If zero, defaults to defaultKeepUnusedServerTime.
84+
UnusedServerTime time.Duration
85+
// MTU is the maximum transmission unit for relay connections.
86+
// If zero, defaults to defaultMTU (1280).
87+
// Must be between minMTU (1280) and maxMTU (65535).
88+
MTU uint16
6889
}
6990

7091
// NewManager creates a new manager instance.
7192
// The serverURL address can be empty. In this case, the manager will not serve.
72-
func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uint16) *Manager {
93+
// Optional parameters can be configured using ManagerOpts. Pass nil to use default values.
94+
func NewManager(ctx context.Context, serverURLs []string, peerID string, opts *ManagerOpts) *Manager {
7395
tokenStore := &relayAuth.TokenStore{}
7496

97+
cleanupInterval := defaultRelayCleanupInterval
98+
unusedServerTime := defaultKeepUnusedServerTime
99+
mtu := uint16(defaultMTU)
100+
101+
if opts != nil {
102+
if opts.CleanupInterval > 0 {
103+
cleanupInterval = opts.CleanupInterval
104+
}
105+
if opts.UnusedServerTime > 0 {
106+
unusedServerTime = opts.UnusedServerTime
107+
}
108+
if opts.MTU > 0 {
109+
if opts.MTU < minMTU {
110+
log.Warnf("MTU %d is below minimum %d, using minimum", opts.MTU, minMTU)
111+
mtu = minMTU
112+
} else if opts.MTU > maxMTU {
113+
log.Warnf("MTU %d exceeds maximum %d, using maximum", opts.MTU, maxMTU)
114+
mtu = maxMTU
115+
} else {
116+
mtu = opts.MTU
117+
}
118+
}
119+
}
120+
75121
m := &Manager{
76122
ctx: ctx,
77123
peerID: peerID,
@@ -85,6 +131,8 @@ func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uin
85131
},
86132
relayClients: make(map[string]*RelayTrack),
87133
onDisconnectedListeners: make(map[string]*list.List),
134+
cleanupInterval: cleanupInterval,
135+
unusedServerTime: unusedServerTime,
88136
}
89137
m.serverPicker.ServerURLs.Store(serverURLs)
90138
m.reconnectGuard = NewGuard(m.serverPicker)
@@ -334,7 +382,7 @@ func (m *Manager) isForeignServer(address string) (bool, error) {
334382
}
335383

336384
func (m *Manager) startCleanupLoop() {
337-
ticker := time.NewTicker(relayCleanupInterval)
385+
ticker := time.NewTicker(m.cleanupInterval)
338386
defer ticker.Stop()
339387
for {
340388
select {
@@ -359,7 +407,7 @@ func (m *Manager) cleanUpUnusedRelays() {
359407
continue
360408
}
361409

362-
if time.Since(rt.created) <= keepUnusedServerTime {
410+
if time.Since(rt.created) <= m.unusedServerTime {
363411
rt.Unlock()
364412
continue
365413
}

‎shared/relay/client/manager_test.go‎

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import (
1616
func TestEmptyURL(t *testing.T) {
1717
ctx, cancel := context.WithCancel(context.Background())
1818
defer cancel()
19-
mgr := NewManager(ctx, nil, "alice", iface.DefaultMTU)
19+
mgr := NewManager(ctx, nil, "alice", &ManagerOpts{MTU: iface.DefaultMTU})
2020
err := mgr.Serve()
2121
if err == nil {
2222
t.Errorf("expected error, got nil")
@@ -91,12 +91,12 @@ func TestForeignConn(t *testing.T) {
9191

9292
mCtx, cancel := context.WithCancel(ctx)
9393
defer cancel()
94-
clientAlice := NewManager(mCtx, toURL(lstCfg1), "alice", iface.DefaultMTU)
94+
clientAlice := NewManager(mCtx, toURL(lstCfg1), "alice", &ManagerOpts{MTU: iface.DefaultMTU})
9595
if err := clientAlice.Serve(); err != nil {
9696
t.Fatalf("failed to serve manager: %s", err)
9797
}
9898

99-
clientBob := NewManager(mCtx, toURL(srvCfg2), "bob", iface.DefaultMTU)
99+
clientBob := NewManager(mCtx, toURL(srvCfg2), "bob", &ManagerOpts{MTU: iface.DefaultMTU})
100100
if err := clientBob.Serve(); err != nil {
101101
t.Fatalf("failed to serve manager: %s", err)
102102
}
@@ -198,12 +198,12 @@ func TestForeginConnClose(t *testing.T) {
198198
mCtx, cancel := context.WithCancel(ctx)
199199
defer cancel()
200200

201-
mgrBob := NewManager(mCtx, toURL(srvCfg2), "bob", iface.DefaultMTU)
201+
mgrBob := NewManager(mCtx, toURL(srvCfg2), "bob", &ManagerOpts{MTU: iface.DefaultMTU})
202202
if err := mgrBob.Serve(); err != nil {
203203
t.Fatalf("failed to serve manager: %s", err)
204204
}
205205

206-
mgr := NewManager(mCtx, toURL(srvCfg1), "alice", iface.DefaultMTU)
206+
mgr := NewManager(mCtx, toURL(srvCfg1), "alice", &ManagerOpts{MTU: iface.DefaultMTU})
207207
err = mgr.Serve()
208208
if err != nil {
209209
t.Fatalf("failed to serve manager: %s", err)
@@ -221,8 +221,8 @@ func TestForeginConnClose(t *testing.T) {
221221

222222
func TestForeignAutoClose(t *testing.T) {
223223
ctx := context.Background()
224-
relayCleanupInterval = 1 * time.Second
225-
keepUnusedServerTime = 2 * time.Second
224+
testCleanupInterval := 1 * time.Second
225+
testUnusedServerTime := 2 * time.Second
226226

227227
srvCfg1 := server.ListenerConfig{
228228
Address: "localhost:1234",
@@ -283,7 +283,11 @@ func TestForeignAutoClose(t *testing.T) {
283283
t.Log("connect to server 1.")
284284
mCtx, cancel := context.WithCancel(ctx)
285285
defer cancel()
286-
mgr := NewManager(mCtx, toURL(srvCfg1), idAlice, iface.DefaultMTU)
286+
mgr := NewManager(mCtx, toURL(srvCfg1), idAlice, &ManagerOpts{
287+
MTU: iface.DefaultMTU,
288+
CleanupInterval: testCleanupInterval,
289+
UnusedServerTime: testUnusedServerTime,
290+
})
287291
err = mgr.Serve()
288292
if err != nil {
289293
t.Fatalf("failed to serve manager: %s", err)
@@ -310,7 +314,7 @@ func TestForeignAutoClose(t *testing.T) {
310314
}
311315

312316
// Wait for cleanup to happen
313-
timeout := relayCleanupInterval + keepUnusedServerTime + 2*time.Second
317+
timeout := testCleanupInterval + testUnusedServerTime + 2*time.Second
314318
t.Logf("waiting for relay cleanup: %s", timeout)
315319

316320
select {
@@ -354,13 +358,13 @@ func TestAutoReconnect(t *testing.T) {
354358
mCtx, cancel := context.WithCancel(ctx)
355359
defer cancel()
356360

357-
clientBob := NewManager(mCtx, toURL(srvCfg), "bob", iface.DefaultMTU)
361+
clientBob := NewManager(mCtx, toURL(srvCfg), "bob", &ManagerOpts{MTU: iface.DefaultMTU})
358362
err = clientBob.Serve()
359363
if err != nil {
360364
t.Fatalf("failed to serve manager: %s", err)
361365
}
362366

363-
clientAlice := NewManager(mCtx, toURL(srvCfg), "alice", iface.DefaultMTU)
367+
clientAlice := NewManager(mCtx, toURL(srvCfg), "alice", &ManagerOpts{MTU: iface.DefaultMTU})
364368
err = clientAlice.Serve()
365369
if err != nil {
366370
t.Fatalf("failed to serve manager: %s", err)
@@ -429,12 +433,12 @@ func TestNotifierDoubleAdd(t *testing.T) {
429433
mCtx, cancel := context.WithCancel(ctx)
430434
defer cancel()
431435

432-
clientBob := NewManager(mCtx, toURL(listenerCfg1), "bob", iface.DefaultMTU)
436+
clientBob := NewManager(mCtx, toURL(listenerCfg1), "bob", &ManagerOpts{MTU: iface.DefaultMTU})
433437
if err = clientBob.Serve(); err != nil {
434438
t.Fatalf("failed to serve manager: %s", err)
435439
}
436440

437-
clientAlice := NewManager(mCtx, toURL(listenerCfg1), "alice", iface.DefaultMTU)
441+
clientAlice := NewManager(mCtx, toURL(listenerCfg1), "alice", &ManagerOpts{MTU: iface.DefaultMTU})
438442
if err = clientAlice.Serve(); err != nil {
439443
t.Fatalf("failed to serve manager: %s", err)
440444
}

0 commit comments

Comments
 (0)