Abort mux srflx gather on close

[JoTurk]

Description

Makes srflx gather abort-able and passes context in more code, this can deadlock under early close for full-ice agents. continues after #934

did refactor GetXORMappedAddr to make the linter happy.

[Copilot]

Pull request overview

This PR aims to prevent agent shutdown deadlocks under early Close() by making server-reflexive (srflx) gathering abortable: it propagates context.Context through STUN/XOR-mapped address discovery and ensures blocked mux writes can be aborted during close (continuing the work from #934).

Changes:

• Add context-aware write plumbing in UDPMuxDefault so a canceled context can abort a blocked WriteTo.
• Introduce GetXORMappedAddrContext for UniversalUDPMuxDefault and use it during srflx gathering when available.
• Extend agent close pre-stop behavior to abort srflx mux writes and add a regression test covering the deadlock scenario.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
udp_mux.go	Adds context-aware write path (writeToContext / startWriteContext) to support aborting blocked writes on cancellation.
udp_mux_universal.go	Adds GetXORMappedAddrContext, routes STUN writes through context-aware mux writing, and avoids time.After allocation by using a timer.
gather.go	Uses a helper to prefer context-aware XOR-mapped address lookup during srflx gathering with UDPMux.
agent.go	Replaces the close pre-stop hook to also abort srflx mux writes during shutdown.
agent_test.go	Adds a regression test that validates close aborts a blocked srflx gather write.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

❌ Patch coverage is 82.89474% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.21%. Comparing base (9651b02) to head (1e87163).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
udp_mux.go	68.96%	5 Missing and 4 partials ⚠️
gather.go	71.42%	1 Missing and 1 partial ⚠️
udp_mux_universal.go	95.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #940      +/-   ##
==========================================
- Coverage   88.33%   88.21%   -0.12%     
==========================================
  Files          46       46              
  Lines        6080     6129      +49     
==========================================
+ Hits         5371     5407      +36     
- Misses        487      497      +10     
- Partials      222      225       +3     

Flag	Coverage Δ
go	88.21% <82.89%> (-0.12%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

udp_mux_universal.go:263

• writeSTUN currently sends a new STUN BindingRequest even when there is already an entry in xorMappedMap for the same server. That contradicts the nearby comment (“we already sent a STUN request”) and can generate a burst of duplicate STUN requests when multiple goroutines call GetXORMappedAddr* concurrently for the same STUN server. Consider reusing the existing waitAddrReceived channel when the map entry already exists, and only sending a packet when creating a new entry; also delete/close the newly-created entry if the initial write fails so callers don’t end up waiting on a never-fulfilled pending entry.

	// If record present in the map, we already sent a STUN request,
	// just wait when waitAddrReceived will be closed
	addrMap, ok := m.xorMappedMap[serverAddr.String()]
	if !ok {
		addrMap = &xorMapped{

[JoTurk]

Hmm that's not possible without doing an extraction and a refactor to the actual code, I'll try to do it later with future cleanups.

[boks1971]

lgtm! but, I am not able to quite follow the abort write loop in writeToContext. Would be good to get more 👀 on this.

[JoTurk]

lgtm! but, I am not able to quite follow the abort write loop in writeToContext. Would be good to get more 👀 on this.

Thank you for the review, And sorry for not adding comments to explain it. I'll add more comments.

This path is only for stun-writes and doesn't affect normal writes, because we don't use cancellable contexts with non stun writes. So,

• WriteState is an atomic state, blocked, deadline or number of current writes.

for a good path:

• startWriteContext increments in-fight count unless blocked is set.
• finishWrite decrements the count.

For aborting:

• WriteToContext start a goroutine before calling writeTo
• if writeTo returns first the watcher is stopped.

For blocked aborting:

• if ctx.done fires while writeTo is still blocked.
• abortWrite checks if there at least one active writer.
• it sets state to blocked so no new write can enter.
• we call setWriteDeadline to unblock.
• the last writer clears the write deadline and resets writeState.