flesh out tests

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

hscontrol/poll.go

@@ -191,6 +191,7 @@ func (m *mapSession) serve() {
 //
 //nolint:gocyclo
 func (m *mapSession) serveLongPoll() {
+	start := time.Now()
 	m.beforeServeLongPoll()
 
 	// Clean up the session when the client disconnects
@@ -235,16 +236,6 @@ func (m *mapSession) serveLongPoll() {
 
 	m.pollFailoverRoutes("node connected", m.node)
 
-	// Upgrade the writer to a ResponseController
-	rc := http.NewResponseController(m.w)
-
-	// Longpolling will break if there is a write timeout,
-	// so it needs to be disabled.
-	rc.SetWriteDeadline(time.Time{})
-
-	ctx, cancel := context.WithCancel(context.WithValue(m.ctx, nodeNameContextKey, m.node.Hostname))
-	defer cancel()
-
 	m.keepAliveTicker = time.NewTicker(m.keepAlive)
 
 	m.h.nodeNotifier.AddNode(m.node.ID, m.ch)
@@ -258,12 +249,12 @@ func (m *mapSession) serveLongPoll() {
 		// consume channels with update, keep alives or "batch" blocking signals
 		select {
 		case <-m.cancelCh:
-			m.tracef("poll cancelled received")
+			m.tracef("poll cancelled received (%s)", time.Since(start).String())
 			mapResponseEnded.WithLabelValues("cancelled").Inc()
 			return
 
-		case <-ctx.Done():
-			m.tracef("poll context done")
+		case <-m.ctx.Done():
+			m.tracef("poll context done (%s): %s", time.Since(start).String(), m.ctx.Err().Error())
 			mapResponseEnded.WithLabelValues("done").Inc()
 			return
 
@@ -354,14 +345,7 @@ func (m *mapSession) serveLongPoll() {
 					m.errf(err, "could not write the map response(%s), for mapSession: %p", update.Type.String(), m)
 					return
 				}
-
-				err = rc.Flush()
-				if err != nil {
-					mapResponseSent.WithLabelValues("error", updateType).Inc()
-					m.errf(err, "flushing the map response to client, for mapSession: %p", m)
-					return
-				}
-
+				m.w.(http.Flusher).Flush()
 				log.Trace().Str("node", m.node.Hostname).TimeDiff("timeSpent", time.Now(), startWrite).Str("mkey", m.node.MachineKey.String()).Msg("finished writing mapresp to node")
 
 				if debugHighCardinalityMetrics {
@@ -375,22 +359,17 @@ func (m *mapSession) serveLongPoll() {
 		case <-m.keepAliveTicker.C:
 			data, err := m.mapper.KeepAliveResponse(m.req, m.node)
 			if err != nil {
-				m.errf(err, "Error generating the keep alive msg")
+				m.errf(err, "Error generating the keepalive msg")
 				mapResponseSent.WithLabelValues("error", "keepalive").Inc()
 				return
 			}
 			_, err = m.w.Write(data)
 			if err != nil {
-				m.errf(err, "Cannot write keep alive message")
-				mapResponseSent.WithLabelValues("error", "keepalive").Inc()
-				return
-			}
-			err = rc.Flush()
-			if err != nil {
-				m.errf(err, "flushing keep alive to client, for mapSession: %p", m)
+				m.errf(err, "Cannot write keepalive message")
 				mapResponseSent.WithLabelValues("error", "keepalive").Inc()
 				return
 			}
+			m.w.(http.Flusher).Flush()
 
 			if debugHighCardinalityMetrics {
 				mapResponseLastSentSeconds.WithLabelValues("keepalive", m.node.ID.String()).Set(float64(time.Now().Unix()))

integration/route_test.go

@@ -1492,7 +1492,7 @@ func TestHASubnetRouterFailoverWhenNodeDisconnects2129(t *testing.T) {
 	assert.True(t, nodeList[2].Online)
 
 	// Kill off one of the docker containers to simulate a disconnect
-	err = scenario.DisconnectContainers(subRouter1.Hostname())
+	err = scenario.DisconnectContainersFromScenario(subRouter1.Hostname())
 	assertNoErr(t, err)
 
 	time.Sleep(5 * time.Second)
@@ -1514,4 +1514,53 @@ func TestHASubnetRouterFailoverWhenNodeDisconnects2129(t *testing.T) {
 	assert.False(t, nodeListAfterDisconnect[0].Online)
 	assert.True(t, nodeListAfterDisconnect[1].Online)
 	assert.True(t, nodeListAfterDisconnect[2].Online)
+
+	var routesAfterDisconnect []*v1.Route
+	err = executeAndUnmarshal(
+		headscale,
+		[]string{
+			"headscale",
+			"routes",
+			"list",
+			"--output",
+			"json",
+		},
+		&routesAfterDisconnect,
+	)
+	assertNoErr(t, err)
+	assert.Len(t, routesAfterDisconnect, 2)
+
+	// Node 1 is primary
+	assert.Equal(t, true, routesAfterDisconnect[0].GetAdvertised())
+	assert.Equal(t, true, routesAfterDisconnect[0].GetEnabled())
+	assert.Equal(t, false, routesAfterDisconnect[0].GetIsPrimary(), "both subnet routers are up, expected r1 to be non-primary")
+
+	// Node 2 is not primary
+	assert.Equal(t, true, routesAfterDisconnect[1].GetAdvertised())
+	assert.Equal(t, true, routesAfterDisconnect[1].GetEnabled())
+	assert.Equal(t, true, routesAfterDisconnect[1].GetIsPrimary(), "both subnet routers are up, expected r2 to be primary")
+
+	// // Ensure the node can reconncet as expected
+	// err = scenario.ConnectContainersToScenario(subRouter1.Hostname())
+	// assertNoErr(t, err)
+
+	// time.Sleep(5 * time.Second)
+
+	// var nodeListAfterReconnect []v1.Node
+	// err = executeAndUnmarshal(
+	// 	headscale,
+	// 	[]string{
+	// 		"headscale",
+	// 		"nodes",
+	// 		"list",
+	// 		"--output",
+	// 		"json",
+	// 	},
+	// 	&nodeListAfterReconnect,
+	// )
+	// assert.Nil(t, err)
+	// assert.Len(t, nodeListAfterReconnect, 3)
+	// assert.True(t, nodeListAfterReconnect[0].Online)
+	// assert.True(t, nodeListAfterReconnect[1].Online)
+	// assert.True(t, nodeListAfterReconnect[2].Online)
 }

integration/scenario.go

@@ -650,8 +650,8 @@ func (s *Scenario) WaitForTailscaleLogout() error {
 	return nil
 }
 
-// DisconnectContainers disconnects a list of containers from the network.
-func (s *Scenario) DisconnectContainers(containers ...string) error {
+// DisconnectContainersFromScenario disconnects a list of containers from the network.
+func (s *Scenario) DisconnectContainersFromScenario(containers ...string) error {
 	for _, container := range containers {
 		if ctr, ok := s.pool.ContainerByName(container); ok {
 			err := ctr.DisconnectFromNetwork(s.network)
@@ -663,3 +663,17 @@ func (s *Scenario) DisconnectContainers(containers ...string) error {
 
 	return nil
 }
+
+// ConnectContainersToScenario disconnects a list of containers from the network.
+func (s *Scenario) ConnectContainersToScenario(containers ...string) error {
+	for _, container := range containers {
+		if ctr, ok := s.pool.ContainerByName(container); ok {
+			err := ctr.ConnectToNetwork(s.network)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}