Fix race condition in connection tests

kylecarbs · kylecarbs · commit bd7fb1314256 · 2022-10-15T23:18:43.000Z
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -283,11 +283,11 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 				for _, replica := range api.replicaManager.Regional() {
 					addresses = append(addresses, replica.RelayAddress)
 				}
-				api.derpMesh.SetAddresses(addresses)
+				api.derpMesh.SetAddresses(addresses, false)
 				_ = api.updateEntitlements(ctx)
 			})
 		} else {
-			api.derpMesh.SetAddresses([]string{})
+			api.derpMesh.SetAddresses([]string{}, false)
 			api.replicaManager.SetCallback(func() {
 				// If the amount of replicas change, so should our entitlements.
 				// This is to display a warning in the UI if the user is unlicensed.
diff --git a/enterprise/derpmesh/derpmesh.go b/enterprise/derpmesh/derpmesh.go
@@ -42,7 +42,10 @@ type Mesh struct {
 
 // SetAddresses performs a diff of the incoming addresses and adds
 // or removes DERP clients from the mesh.
-func (m *Mesh) SetAddresses(addresses []string) {
+//
+// Connect is only used for testing to ensure DERPs are meshed before
+// exchanging messages.
+func (m *Mesh) SetAddresses(addresses []string, connect bool) {
 	total := make(map[string]struct{}, 0)
 	for _, address := range addresses {
 		addressURL, err := url.Parse(address)
@@ -58,7 +61,7 @@ func (m *Mesh) SetAddresses(addresses []string) {
 		address = derpURL.String()
 
 		total[address] = struct{}{}
-		added, err := m.addAddress(address)
+		added, err := m.addAddress(address, connect)
 		if err != nil {
 			m.logger.Error(m.ctx, "failed to add address", slog.F("address", address), slog.Error(err))
 			continue
@@ -85,7 +88,7 @@ func (m *Mesh) SetAddresses(addresses []string) {
 // addAddress begins meshing with a new address. It returns false if the address is already being meshed with.
 // It's expected that this is a full HTTP address with a path.
 // e.g. http://127.0.0.1:8080/derp
-func (m *Mesh) addAddress(address string) (bool, error) {
+func (m *Mesh) addAddress(address string, connect bool) (bool, error) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 	_, isActive := m.active[address]
@@ -102,6 +105,9 @@ func (m *Mesh) addAddress(address string) (bool, error) {
 		var dialer net.Dialer
 		return dialer.DialContext(ctx, network, addr)
 	})
+	if connect {
+		_ = client.Connect(m.ctx)
+	}
 	ctx, cancelFunc := context.WithCancel(m.ctx)
 	closed := make(chan struct{})
 	closeFunc := func() {
diff --git a/enterprise/derpmesh/derpmesh_test.go b/enterprise/derpmesh/derpmesh_test.go
@@ -49,9 +49,9 @@ func TestDERPMesh(t *testing.T) {
 		defer firstServer.Close()
 		secondServer, secondServerURL := startDERP(t, tlsConfig)
 		firstMesh := derpmesh.New(slogtest.Make(t, nil).Named("first").Leveled(slog.LevelDebug), firstServer, tlsConfig)
-		firstMesh.SetAddresses([]string{secondServerURL})
+		firstMesh.SetAddresses([]string{secondServerURL}, true)
 		secondMesh := derpmesh.New(slogtest.Make(t, nil).Named("second").Leveled(slog.LevelDebug), secondServer, tlsConfig)
-		secondMesh.SetAddresses([]string{firstServerURL})
+		secondMesh.SetAddresses([]string{firstServerURL}, true)
 		defer firstMesh.Close()
 		defer secondMesh.Close()
 
@@ -78,9 +78,9 @@ func TestDERPMesh(t *testing.T) {
 		t.Parallel()
 		server, serverURL := startDERP(t, tlsConfig)
 		mesh := derpmesh.New(slogtest.Make(t, nil).Named("first").Leveled(slog.LevelDebug), server, tlsConfig)
-		mesh.SetAddresses([]string{"http://fake.com"})
+		mesh.SetAddresses([]string{"http://fake.com"}, false)
 		// This should trigger a removal...
-		mesh.SetAddresses([]string{})
+		mesh.SetAddresses([]string{}, false)
 		defer mesh.Close()
 
 		first := key.NewNode()
@@ -114,7 +114,7 @@ func TestDERPMesh(t *testing.T) {
 			meshes = append(meshes, mesh)
 		}
 		for _, mesh := range meshes {
-			mesh.SetAddresses(serverURLs)
+			mesh.SetAddresses(serverURLs, true)
 		}
 
 		first := key.NewNode()
