From c6358f2247a5f76ce7536b3d5d457569cfe32bdb Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Fri, 12 Mar 2021 11:34:49 -0800
Subject: [PATCH] net/netcheck: add a few more STUN retries for prior DERP home

For #1310, maybe.

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 net/netcheck/netcheck.go      | 11 ++++++++++-
 net/netcheck/netcheck_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/net/netcheck/netcheck.go b/net/netcheck/netcheck.go
index 9398a1c2c..097b58704 100644
--- a/net/netcheck/netcheck.go
+++ b/net/netcheck/netcheck.go
@@ -362,7 +362,7 @@ func makeProbePlan(dm *tailcfg.DERPMap, ifState *interfaces.State, last *Report)
 			tries = 2
 		} else if hadBoth {
 			// For dual stack machines, make the 3rd & slower nodes alternate
-			// breetween
+			// beetween.
 			if ri%2 == 0 {
 				do4, do6 = true, false
 			} else {
@@ -373,6 +373,12 @@ func makeProbePlan(dm *tailcfg.DERPMap, ifState *interfaces.State, last *Report)
 			do6 = false
 		}
 
+		if reg.RegionID == last.PreferredDERP {
+			// But if we already had a DERP home, try extra hard to
+			// make sure it's there so we don't flip flop around.
+			tries = 4
+		}
+
 		for try := 0; try < tries; try++ {
 			if len(reg.Nodes) == 0 {
 				// Shouldn't be possible.
@@ -387,6 +393,9 @@ func makeProbePlan(dm *tailcfg.DERPMap, ifState *interfaces.State, last *Report)
 				prevLatency = defaultActiveRetransmitTime
 			}
 			delay := time.Duration(try) * prevLatency
+			if try > 1 {
+				delay += time.Duration(try) * 50 * time.Millisecond
+			}
 			if do4 {
 				p4 = append(p4, probe{delay: delay, node: n.Name, proto: probeIPv4})
 			}
diff --git a/net/netcheck/netcheck_test.go b/net/netcheck/netcheck_test.go
index d3c77b71e..8919da761 100644
--- a/net/netcheck/netcheck_test.go
+++ b/net/netcheck/netcheck_test.go
@@ -410,6 +410,35 @@ func TestMakeProbePlan(t *testing.T) {
 				"region-5-v6": []probe{p("5a", 6), p("5b", 6, 100*ms), p("5c", 6, 200*ms)},
 			},
 		},
+		{
+			name:    "try_harder_for_preferred_derp",
+			dm:      basicMap,
+			have6if: true,
+			last: &Report{
+				RegionLatency: map[int]time.Duration{
+					1: 10 * time.Millisecond,
+					2: 20 * time.Millisecond,
+					3: 30 * time.Millisecond,
+					4: 40 * time.Millisecond,
+				},
+				RegionV4Latency: map[int]time.Duration{
+					1: 10 * time.Millisecond,
+					2: 20 * time.Millisecond,
+				},
+				RegionV6Latency: map[int]time.Duration{
+					3: 30 * time.Millisecond,
+					4: 40 * time.Millisecond,
+				},
+				PreferredDERP: 1,
+			},
+			want: probePlan{
+				"region-1-v4": []probe{p("1a", 4), p("1a", 4, 12*ms), p("1a", 4, 124*ms), p("1a", 4, 186*ms)},
+				"region-1-v6": []probe{p("1a", 6), p("1a", 6, 12*ms), p("1a", 6, 124*ms), p("1a", 6, 186*ms)},
+				"region-2-v4": []probe{p("2a", 4), p("2b", 4, 24*ms)},
+				"region-2-v6": []probe{p("2a", 6), p("2b", 6, 24*ms)},
+				"region-3-v4": []probe{p("3a", 4)},
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {