Codereorg: Part 5, state

.gitignore

@@ -1,4 +1,5 @@
 ignored/
+tailscale/
 
 # Binaries for programs and plugins
 *.exe

.golangci.yaml

@@ -10,6 +10,8 @@ issues:
 linters:
   enable-all: true
   disable:
+    - depguard
+
     - exhaustivestruct
     - revive
     - lll

Dockerfile

@@ -1,5 +1,5 @@
 # Builder image
-FROM docker.io/golang:1.20-bullseye AS build
+FROM docker.io/golang:1.21-bookworm AS build
 ARG VERSION=dev
 ENV GOPATH /go
 WORKDIR /go/src/headscale
@@ -14,7 +14,7 @@ RUN strip /go/bin/headscale
 RUN test -e /go/bin/headscale
 
 # Production image
-FROM docker.io/debian:bullseye-slim
+FROM docker.io/debian:bookworm-slim
 
 RUN apt-get update \
     && apt-get install -y ca-certificates \

Dockerfile.debug

@@ -18,6 +18,10 @@ FROM docker.io/golang:1.20.0-bullseye
 COPY --from=build /go/bin/headscale /bin/headscale
 ENV TZ UTC
 
+RUN apt-get update \
+    && apt-get install --no-install-recommends --yes less jq \
+    && rm -rf /var/lib/apt/lists/* \
+    && apt-get clean
 RUN mkdir -p /var/run/headscale
 
 # Need to reset the entrypoint or everything will run as a busybox script

cmd/headscale/cli/root.go

@@ -51,7 +51,7 @@ func initConfig() {
 
 	cfg, err := types.GetHeadscaleConfig()
 	if err != nil {
-		log.Fatal().Caller().Err(err)
+		log.Fatal().Caller().Err(err).Msg("Failed to get headscale configuration")
 	}
 
 	machineOutput := HasMachineOutputFlag()

cmd/headscale/cli/utils.go

@@ -154,17 +154,17 @@ func SuccessOutput(result interface{}, override string, outputFormat string) {
 	case "json":
 		jsonBytes, err = json.MarshalIndent(result, "", "\t")
 		if err != nil {
-			log.Fatal().Err(err)
+			log.Fatal().Err(err).Msg("failed to unmarshal output")
 		}
 	case "json-line":
 		jsonBytes, err = json.Marshal(result)
 		if err != nil {
-			log.Fatal().Err(err)
+			log.Fatal().Err(err).Msg("failed to unmarshal output")
 		}
 	case "yaml":
 		jsonBytes, err = yaml.Marshal(result)
 		if err != nil {
-			log.Fatal().Err(err)
+			log.Fatal().Err(err).Msg("failed to unmarshal output")
 		}
 	default:
 		//nolint

flake.nix

@@ -33,7 +33,7 @@
 
           # When updating go.mod or go.sum, a new sha will need to be calculated,
           # update this if you have a mismatch after doing a change to thos files.
-          vendorSha256 = "sha256-9Hol8w8HB28AlulshMYYQwOgvGzR47qxzyPrB8G0XSQ=";
+          vendorSha256 = "sha256-dNE5wgR3oWXlYzPNXp0v/GGwY0/hvhOB5JWCb5EIbg8=";
 
           ldflags = ["-s" "-w" "-X github.com/juanfont/headscale/cmd/headscale/cli.Version=v${version}"];
         };

go.mod

@@ -61,7 +61,7 @@ require (
 	github.com/containerd/console v1.0.3 // indirect
 	github.com/containerd/continuity v0.3.0 // indirect
 	github.com/docker/cli v23.0.5+incompatible // indirect
-	github.com/docker/docker v23.0.5+incompatible // indirect
+	github.com/docker/docker v24.0.4+incompatible // indirect
 	github.com/docker/go-connections v0.4.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect

go.sum

@@ -109,8 +109,8 @@ github.com/deckarep/golang-set/v2 v2.3.0 h1:qs18EKUfHm2X9fA50Mr/M5hccg2tNnVqsiBI
 github.com/deckarep/golang-set/v2 v2.3.0/go.mod h1:VAky9rY/yGXJOLEDv3OMci+7wtDpOF4IN+y82NBOac4=
 github.com/docker/cli v23.0.5+incompatible h1:ufWmAOuD3Vmr7JP2G5K3cyuNC4YZWiAsuDEvFVVDafE=
 github.com/docker/cli v23.0.5+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8=
-github.com/docker/docker v23.0.5+incompatible h1:DaxtlTJjFSnLOXVNUBU1+6kXGz2lpDoEAH6QoxaSg8k=
-github.com/docker/docker v23.0.5+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/docker/docker v24.0.4+incompatible h1:s/LVDftw9hjblvqIeTiGYXBCD95nOEEl7qRsRrIOuQI=
+github.com/docker/docker v24.0.4+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/go-connections v0.4.0 h1:El9xVISelRB7BuFusrZozjnkIM5YnzCViNKohAFqRJQ=
 github.com/docker/go-connections v0.4.0/go.mod h1:Gbd7IOopHjR8Iph03tsViu4nIes5XhDvyHbTtUxmeec=
 github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=

hscontrol/app.go

@@ -8,9 +8,10 @@ import (
 	"io"
 	"net"
 	"net/http"
+	_ "net/http/pprof" //nolint
 	"os"
 	"os/signal"
-	"sort"
+	"runtime"
 	"strconv"
 	"strings"
 	"sync"
@@ -20,19 +21,19 @@ import (
 	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/gorilla/mux"
 	grpcMiddleware "github.com/grpc-ecosystem/go-grpc-middleware"
-	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
+	grpcRuntime "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 	"github.com/juanfont/headscale"
 	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 	"github.com/juanfont/headscale/hscontrol/db"
 	"github.com/juanfont/headscale/hscontrol/derp"
 	derpServer "github.com/juanfont/headscale/hscontrol/derp/server"
+	"github.com/juanfont/headscale/hscontrol/notifier"
 	"github.com/juanfont/headscale/hscontrol/policy"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/patrickmn/go-cache"
 	zerolog "github.com/philip-bui/grpc-zerolog"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
-	"github.com/puzpuzpuz/xsync/v2"
 	zl "github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"golang.org/x/crypto/acme"
@@ -84,7 +85,7 @@ type Headscale struct {
 
 	ACLPolicy *policy.ACLPolicy
 
-	lastStateChange *xsync.MapOf[string, time.Time]
+	nodeNotifier *notifier.Notifier
 
 	oidcProvider *oidc.Provider
 	oauth2Config *oauth2.Config
@@ -93,12 +94,13 @@ type Headscale struct {
 
 	shutdownChan       chan struct{}
 	pollNetMapStreamWG sync.WaitGroup
-
-	stateUpdateChan       chan struct{}
-	cancelStateUpdateChan chan struct{}
 }
 
 func NewHeadscale(cfg *types.Config) (*Headscale, error) {
+	if _, enableProfile := os.LookupEnv("HEADSCALE_PROFILING_ENABLED"); enableProfile {
+		runtime.SetBlockProfileRate(1)
+	}
+
 	privateKey, err := readOrCreatePrivateKey(cfg.PrivateKeyPath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read or create private key: %w", err)
@@ -158,19 +160,14 @@ func NewHeadscale(cfg *types.Config) (*Headscale, error) {
 		noisePrivateKey:    noisePrivateKey,
 		registrationCache:  registrationCache,
 		pollNetMapStreamWG: sync.WaitGroup{},
-		lastStateChange:    xsync.NewMapOf[time.Time](),
-
-		stateUpdateChan:       make(chan struct{}),
-		cancelStateUpdateChan: make(chan struct{}),
+		nodeNotifier:       notifier.NewNotifier(),
 	}
 
-	go app.watchStateChannel()
-
 	database, err := db.NewHeadscaleDatabase(
 		cfg.DBtype,
 		dbString,
 		app.dbDebug,
-		app.stateUpdateChan,
+		app.nodeNotifier,
 		cfg.IPPrefixes,
 		cfg.BaseDomain)
 	if err != nil {
@@ -203,7 +200,11 @@ func NewHeadscale(cfg *types.Config) (*Headscale, error) {
 
 	if cfg.DERP.ServerEnabled {
 		// TODO(kradalby): replace this key with a dedicated DERP key.
-		embeddedDERPServer, err := derpServer.NewDERPServer(cfg.ServerURL, key.NodePrivate(*privateKey), &cfg.DERP)
+		embeddedDERPServer, err := derpServer.NewDERPServer(
+			cfg.ServerURL,
+			key.NodePrivate(*privateKey),
+			&cfg.DERP,
+		)
 		if err != nil {
 			return nil, err
 		}
@@ -230,10 +231,14 @@ func (h *Headscale) expireEphemeralNodes(milliSeconds int64) {
 
 // expireExpiredMachines expires machines that have an explicit expiry set
 // after that expiry time has passed.
-func (h *Headscale) expireExpiredMachines(milliSeconds int64) {
-	ticker := time.NewTicker(time.Duration(milliSeconds) * time.Millisecond)
+func (h *Headscale) expireExpiredMachines(intervalMs int64) {
+	interval := time.Duration(intervalMs) * time.Millisecond
+	ticker := time.NewTicker(interval)
+
+	lastCheck := time.Unix(0, 0)
+
 	for range ticker.C {
-		h.db.ExpireExpiredMachines(h.getLastStateChange())
+		lastCheck = h.db.ExpireExpiredMachines(lastCheck)
 	}
 }
 
@@ -258,7 +263,10 @@ func (h *Headscale) scheduledDERPMapUpdateWorker(cancelChan <-chan struct{}) {
 				h.DERPMap.Regions[region.RegionID] = &region
 			}
 
-			h.setLastStateChangeToNow()
+			h.nodeNotifier.NotifyAll(types.StateUpdate{
+				Type:    types.StateDERPUpdated,
+				DERPMap: *h.DERPMap,
+			})
 		}
 	}
 }
@@ -433,8 +441,9 @@ func (h *Headscale) ensureUnixSocketIsAbsent() error {
 	return os.Remove(h.cfg.UnixSocket)
 }
 
-func (h *Headscale) createRouter(grpcMux *runtime.ServeMux) *mux.Router {
+func (h *Headscale) createRouter(grpcMux *grpcRuntime.ServeMux) *mux.Router {
 	router := mux.NewRouter()
+	router.PathPrefix("/debug/pprof/").Handler(http.DefaultServeMux)
 
 	router.HandleFunc(ts2021UpgradePath, h.NoiseUpgradeHandler).Methods(http.MethodPost)
 
@@ -541,7 +550,7 @@ func (h *Headscale) Serve() error {
 		return fmt.Errorf("failed change permission of gRPC socket: %w", err)
 	}
 
-	grpcGatewayMux := runtime.NewServeMux()
+	grpcGatewayMux := grpcRuntime.NewServeMux()
 
 	// Make the grpc-gateway connect to grpc over socket
 	grpcGatewayConn, err := grpc.Dial(
@@ -722,7 +731,9 @@ func (h *Headscale) Serve() error {
 						Str("path", aclPath).
 						Msg("ACL policy successfully reloaded, notifying nodes of change")
 
-					h.setLastStateChangeToNow()
+					h.nodeNotifier.NotifyAll(types.StateUpdate{
+						Type: types.StateFullUpdate,
+					})
 				}
 
 			default:
@@ -760,10 +771,6 @@ func (h *Headscale) Serve() error {
 				// Stop listening (and unlink the socket if unix type):
 				socketListener.Close()
 
-				<-h.cancelStateUpdateChan
-				close(h.stateUpdateChan)
-				close(h.cancelStateUpdateChan)
-
 				// Close db connections
 				err = h.db.Close()
 				if err != nil {
@@ -775,6 +782,8 @@ func (h *Headscale) Serve() error {
 
 				// And we're done:
 				cancel()
+
+				return
 			}
 		}
 	}
@@ -859,73 +868,6 @@ func (h *Headscale) getTLSSettings() (*tls.Config, error) {
 	}
 }
 
-// TODO(kradalby): baby steps, make this more robust.
-func (h *Headscale) watchStateChannel() {
-	for {
-		select {
-		case <-h.stateUpdateChan:
-			h.setLastStateChangeToNow()
-
-		case <-h.cancelStateUpdateChan:
-			return
-		}
-	}
-}
-
-func (h *Headscale) setLastStateChangeToNow() {
-	var err error
-
-	now := time.Now().UTC()
-
-	users, err := h.db.ListUsers()
-	if err != nil {
-		log.Error().
-			Caller().
-			Err(err).
-			Msg("failed to fetch all users, failing to update last changed state.")
-	}
-
-	for _, user := range users {
-		lastStateUpdate.WithLabelValues(user.Name, "headscale").Set(float64(now.Unix()))
-		if h.lastStateChange == nil {
-			h.lastStateChange = xsync.NewMapOf[time.Time]()
-		}
-		h.lastStateChange.Store(user.Name, now)
-	}
-}
-
-func (h *Headscale) getLastStateChange(users ...types.User) time.Time {
-	times := []time.Time{}
-
-	// getLastStateChange takes a list of users as a "filter", if no users
-	// are past, then use the entier list of users and look for the last update
-	if len(users) > 0 {
-		for _, user := range users {
-			if lastChange, ok := h.lastStateChange.Load(user.Name); ok {
-				times = append(times, lastChange)
-			}
-		}
-	} else {
-		h.lastStateChange.Range(func(key string, value time.Time) bool {
-			times = append(times, value)
-
-			return true
-		})
-	}
-
-	sort.Slice(times, func(i, j int) bool {
-		return times[i].After(times[j])
-	})
-
-	log.Trace().Msgf("Latest times %#v", times)
-
-	if len(times) == 0 {
-		return time.Now().UTC()
-	} else {
-		return times[0]
-	}
-}
-
 func notFoundHandler(
 	writer http.ResponseWriter,
 	req *http.Request,