stuff

Signed-off-by: Kristoffer Dalby <[email protected]>
kradalby · Mar 14, 2024 · f4a53c4 · f4a53c4
1 parent 670e031
commit f4a53c4
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 28 deletions.
diff --git a/hscontrol/db/routes.go b/hscontrol/db/routes.go
@@ -138,14 +138,13 @@ func DisableRoute(tx *gorm.DB,
 	// https://github.com/juanfont/headscale/issues/804#issuecomment-1399314002
 	var update []types.NodeID
 	if !route.IsExitRoute() {
-		update, err = failoverRoute(tx, isConnected, route)
+		route.Enabled = false
+		err = tx.Save(route).Error
 		if err != nil {
 			return nil, err
 		}
 
-		route.Enabled = false
-		route.IsPrimary = false
-		err = tx.Save(route).Error
+		update, err = failoverRoute(tx, isConnected, route)
 		if err != nil {
 			return nil, err
 		}
@@ -159,6 +158,7 @@ func DisableRoute(tx *gorm.DB,
 			if routes[i].IsExitRoute() {
 				routes[i].Enabled = false
 				routes[i].IsPrimary = false
+
 				err = tx.Save(&routes[i]).Error
 				if err != nil {
 					return nil, err
@@ -167,15 +167,6 @@ func DisableRoute(tx *gorm.DB,
 		}
 	}
 
-	if routes == nil {
-		routes, err = GetNodeRoutes(tx, &node)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	node.Routes = routes
-
 	// If update is empty, it means that one was not created
 	// by failover (as a failover was not necessary), create
 	// one and return to the caller.
@@ -506,6 +497,7 @@ func failoverRoute(
 
 	// Find a new suitable route
 	for idx, route := range routes {
+		log.Trace().Uint64("node.id", route.NodeID).Interface("route", route).Msgf("ROUTE DEBUG CHECKING ROUTE")
 		if r.ID == route.ID {
 			continue
 		}
@@ -514,7 +506,9 @@ func failoverRoute(
 			continue
 		}
 
+		log.Trace().Uint64("node.id", route.NodeID).Interface("conmap", isConnected).Msgf("ROUTE DEBUG CHECKING IF ROUTE IS ONLINE")
 		if isConnected != nil && isConnected[route.Node.ID] {
+			log.Trace().Uint64("node.id", route.NodeID).Msgf("ROUTE DEBUG CHECKING ROUTE IS ONLINE")
 			newPrimary = &routes[idx]
 			break
 		}
@@ -526,6 +520,7 @@ func failoverRoute(
 	// the one currently marked as primary is the
 	// best we got.
 	if newPrimary == nil {
+		log.Trace().Uint64("node.id", r.NodeID).Msgf("ROUTE DEBUG DIDNT FIND NEW PRIMARY")
 		return nil, nil
 	}
 

diff --git a/hscontrol/grpcv1.go b/hscontrol/grpcv1.go
@@ -508,9 +508,8 @@ func (api headscaleV1APIServer) DisableRoute(
 	ctx context.Context,
 	request *v1.DisableRouteRequest,
 ) (*v1.DisableRouteResponse, error) {
-	isConnected := api.h.nodeNotifier.ConnectedMap()
 	update, err := db.Write(api.h.db.DB, func(tx *gorm.DB) ([]types.NodeID, error) {
-		return db.DisableRoute(tx, request.GetRouteId(), isConnected)
+		return db.DisableRoute(tx, request.GetRouteId(), api.h.nodeNotifier.ConnectedMap())
 	})
 	if err != nil {
 		return nil, err

diff --git a/hscontrol/noise.go b/hscontrol/noise.go
@@ -250,19 +250,33 @@ func (ns *noiseServer) NoisePollNetMapHandler(
 
 	// If a streaming mapSession exists for this node, close it
 	// and start a new one.
-	// if session.isStreaming() {
-	// 	ns.headscale.mapSessionMu.Lock()
-	// 	if oldSession, ok := ns.headscale.mapSessions[node.ID]; ok {
-	// 		log.Info().
-	// 			Caller().
-	// 			Int("node.id", int(node.ID)).
-	// 			Msg("Node has an open streaming session, replacing")
-	// 		oldSession.close()
-	// 	}
-
-	// 	ns.headscale.mapSessions[node.ID] = session
-	// 	ns.headscale.mapSessionMu.Unlock()
-	// }
+	if session.isStreaming() {
+		defer func() {
+			delete(ns.headscale.mapSessions, node.ID)
+		}()
+
+		log.Debug().
+			Caller().
+			Uint64("node.id", node.ID.Uint64()).
+			Int("cap_ver", int(mapRequest.Version)).
+			Msg("Aquiring lock to check stream")
+		ns.headscale.mapSessionMu.Lock()
+		if oldSession, ok := ns.headscale.mapSessions[node.ID]; ok {
+			log.Info().
+				Caller().
+				Int("node.id", int(node.ID)).
+				Msg("Node has an open streaming session, replacing")
+			oldSession.close()
+		}
+
+		ns.headscale.mapSessions[node.ID] = session
+		ns.headscale.mapSessionMu.Unlock()
+		log.Debug().
+			Caller().
+			Uint64("node.id", node.ID.Uint64()).
+			Int("cap_ver", int(mapRequest.Version)).
+			Msg("Releasing lock to check stream")
+	}
 
 	session.serve()
 }
diff --git a/hscontrol/poll.go b/hscontrol/poll.go
@@ -87,7 +87,9 @@ func (h *Headscale) newMapSession(
 }
 
 func (m *mapSession) close() {
+	m.tracef("ROUTE DEBUG: CLOSING SESSION VIA CHANNEL")
 	m.cancelCh <- struct{}{}
+	m.tracef("ROUTE DEBUG: CLOSING SESSION VIA CHANNEL SENT")
 }
 
 func (m *mapSession) isStreaming() bool {
@@ -333,6 +335,7 @@ func (m *mapSession) serve() {
 		// consume channels with update, keep alives or "batch" blocking signals
 		select {
 		case <-m.cancelCh:
+			m.tracef("ROUTE DEBUG: CLOSE RECEIVED RETURNING")
 			return
 		case <-ctx.Done():
 			return