[client] Update RaceDial to accept context for improved cancellation

[pappz]

Describe your changes

The relay RaceDial.Dial() created its abort context from context.Background(), making it immune to parent context cancellation.

During shutdown, this caused the client to block for ~27 seconds, waiting for the WebSocket dial timeout to expire instead of exiting immediately.

Fix:
Propagate the parent context into RaceDial.Dial(ctx) so that when the engine cancels its context on shutdown, all in-flight dial attempts are aborted immediately.

Issue ticket number and link

Stack

Checklist

• Is it a bug fix
• Is a typo/documentation fix
• Is a feature enhancement
• It is a refactor
• Created tests that fail without the change (if possible)

By submitting this pull request, you confirm that you have read and agree to the terms of the Contributor License Agreement.

Documentation

Select exactly one:

• I added/updated documentation for this change
• Documentation is not needed for this change (explain why)

Docs PR URL (required if "docs added" is checked)

Paste the PR link from https://github.com/netbirdio/docs here:

https://github.com/netbirdio/docs/pull/__

Summary by CodeRabbit

• Refactor
  • Enhanced relay client connection establishment to properly propagate context and timeout signals, enabling more responsive connection handling and cancellation behavior when establishing relay connections.

[coderabbitai]

📝 Walkthrough

Walkthrough

The relay client's race dialer has been updated to accept and propagate a context parameter through its Dial method. The context is now passed from the caller through the client connection establishment to the internal dial operation, enabling context-aware cancellation and deadline handling.

Changes

Cohort / File(s)	Summary
Relay Client Context Propagation shared/relay/client/client.go	Passes caller-provided context to rd.Dial() for context-aware connection establishment.
Race Dialer Signature Update shared/relay/client/dialer/race_dialer.go	Updated Dial() method signature to accept ctx context.Context parameter; internal cancellation context now derives from caller-provided context instead of context.Background().
Test Updates shared/relay/client/dialer/race_dialer_test.go	Updated all test invocations of rd.Dial() to pass context.Background() as parameter, aligning with the new context-aware method signature.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Suggested reviewers

• mlsmaycon
• lixmal

Poem

🐰 A context hops through dialers deep,
Where cancellations safely keep,
No more background, now it flows,
Through racing connections, as context goes!

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: updating RaceDial to accept context for improved cancellation, which directly relates to the core purpose of the changeset.
Description check	✅ Passed	The PR description clearly explains the problem, solution, and impact. Bug fix is properly checked. Required sections are completed with appropriate context.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/relay-race-dialer-context-propagation

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[sonarqubecloud]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[coderabbitai]

🧹 Nitpick comments (1)

shared/relay/client/dialer/race_dialer_test.go (1)

76-252: Add a regression test for parent-context cancellation propagation.

Current tests update the signature but don’t verify the new cancellation behavior this PR is meant to deliver.

Proposed test addition

+func TestRaceDialParentContextCancellation(t *testing.T) {
+	logger := logrus.NewEntry(logrus.New())
+	serverURL := "test.server.com"
+
+	mockDialer := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			<-ctx.Done()
+			return nil, ctx.Err()
+		},
+		protocolStr: "proto1",
+	}
+
+	rd := NewRaceDial(logger, 30*time.Second, serverURL, mockDialer)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	start := time.Now()
+	conn, err := rd.Dial(ctx)
+	if err == nil {
+		t.Fatalf("expected cancellation error, got nil")
+	}
+	if conn != nil {
+		t.Fatalf("expected nil conn, got %v", conn)
+	}
+	if time.Since(start) > 500*time.Millisecond {
+		t.Fatalf("dial did not abort promptly after parent context cancellation")
+	}
+}

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@shared/relay/client/dialer/race_dialer_test.go` around lines 76 - 252, Add a
new unit test (e.g., TestRaceDialParentContextCancellation) that verifies
parent-context cancellation is propagated to dialers: create a parent context
with cancel, make one or more MockDialer implementations whose dialFunc blocks
on <-ctx.Done() and sends ctx.Err() to a channel, call rd.Dial in a goroutine
using the parent context, call cancel() before any dialer returns, then assert
Dial returns an error (and nil conn) and that each mock dialer received
context.Canceled via the channel; reference NewRaceDial, rd.Dial, and the
MockDialer dialFunc/mock2err-style channels used in other tests to locate where
to add this test.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@shared/relay/client/dialer/race_dialer_test.go`:
- Around line 76-252: Add a new unit test (e.g.,
TestRaceDialParentContextCancellation) that verifies parent-context cancellation
is propagated to dialers: create a parent context with cancel, make one or more
MockDialer implementations whose dialFunc blocks on <-ctx.Done() and sends
ctx.Err() to a channel, call rd.Dial in a goroutine using the parent context,
call cancel() before any dialer returns, then assert Dial returns an error (and
nil conn) and that each mock dialer received context.Canceled via the channel;
reference NewRaceDial, rd.Dial, and the MockDialer dialFunc/mock2err-style
channels used in other tests to locate where to add this test.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 56d98e9a-e717-49e7-a9c7-cb98eed40458

📥 Commits

Reviewing files that changed from the base of the PR and between cf86b9a and 0816c48.

📒 Files selected for processing (3)

• shared/relay/client/client.go
• shared/relay/client/dialer/race_dialer.go
• shared/relay/client/dialer/race_dialer_test.go