github.com/m-lab/locate@v0.17.6/cmd/heartbeat/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	compute "cloud.google.com/go/compute/apiv1"
	md "cloud.google.com/go/compute/metadata"
	"github.com/gorilla/websocket"
	"github.com/m-lab/go/flagx"
	"github.com/m-lab/go/memoryless"
	"github.com/m-lab/go/prometheusx"
	"github.com/m-lab/go/rtx"
	v2 "github.com/m-lab/locate/api/v2"
	"github.com/m-lab/locate/cmd/heartbeat/health"
	"github.com/m-lab/locate/cmd/heartbeat/metadata"
	"github.com/m-lab/locate/cmd/heartbeat/registration"
	"github.com/m-lab/locate/connection"
	"github.com/m-lab/locate/metrics"
	"github.com/m-lab/locate/static"
)

var (
	heartbeatURL        string
	hostname			flagx.StringFile
	experiment          string
	pod                 string
	node                string
	namespace           string
	kubernetesAuth      = "/var/run/secrets/kubernetes.io/serviceaccount/"
	kubernetesURL       = flagx.URL{}
	registrationURL     = flagx.URL{}
	services            = flagx.KeyValueArray{}
	heartbeatPeriod     = static.HeartbeatPeriod
	mainCtx, mainCancel = context.WithCancel(context.Background())
	lbPath              = "/metadata/loadbalanced"
)

// Checker generates a health score for the heartbeat instance (0, 1).
type Checker interface {
	GetHealth(ctx context.Context) float64 // Health score.
}

func init() {
	flag.StringVar(&heartbeatURL, "heartbeat-url", "ws://localhost:8080/v2/platform/heartbeat",
		"URL for locate service")
	flag.Var(&hostname, "hostname", "The service hostname (may be read from @/path/file)")
	flag.StringVar(&experiment, "experiment", "", "Experiment name")
	flag.StringVar(&pod, "pod", "", "Kubernetes pod name")
	flag.StringVar(&node, "node", "", "Kubernetes node name")
	flag.StringVar(&namespace, "namespace", "", "Kubernetes namespace")
	flag.Var(&kubernetesURL, "kubernetes-url", "URL for Kubernetes API")
	flag.Var(&registrationURL, "registration-url", "URL for site registration")
	flag.Var(&services, "services", "Maps experiment target names to their set of services")
}

func main() {
	flag.Parse()
	rtx.Must(flagx.ArgsFromEnvWithLog(flag.CommandLine, false), "failed to read args from env")

	// Start metrics server.
	prom := prometheusx.MustServeMetrics()
	defer prom.Close()

	// Load registration data.
	ldrConfig := memoryless.Config{
		Min:      static.RegistrationLoadMin,
		Expected: static.RegistrationLoadExpected,
		Max:      static.RegistrationLoadMax,
	}
	svcs := services.Get()
	ldr, err := registration.NewLoader(mainCtx, registrationURL.URL, hostname.Value, experiment, svcs, ldrConfig)
	rtx.Must(err, "could not initialize registration loader")
	r, err := ldr.GetRegistration(mainCtx)
	rtx.Must(err, "could not load registration data")
	hbm := v2.HeartbeatMessage{Registration: r}

	// Establish a connection.
	conn := connection.NewConn()
	err = conn.Dial(heartbeatURL, http.Header{}, hbm)
	rtx.Must(err, "failed to establish a websocket connection with %s", heartbeatURL)

	probe := health.NewPortProbe(svcs)
	ec := health.NewEndpointClient(static.HealthEndpointTimeout)
	var hc Checker

	// TODO(kinkade): cause a fatal error if lberr is not nil. Not fatally
	// exiting on lberr is just a workaround to get this rolled out while we
	// wait for every physical machine on the platform to actually have that
	// file, which won't be the case until the rolling reboot in production
	// completes in 4 or 5 days, as of this comment 2024-08-06.
	lbbytes, lberr := os.ReadFile(lbPath)

	// If the "loadbalanced" file exists, then make sure that the content of the
	// file is "true". If the file doesn't exist, then, for now, just consider
	// the machine as not loadbalanced.
	if lberr == nil && string(lbbytes) == "true" {
		gcpmd, err := metadata.NewGCPMetadata(md.NewClient(http.DefaultClient), hostname.Value)
		rtx.Must(err, "failed to get VM metadata")
		gceClient, err := compute.NewRegionBackendServicesRESTClient(mainCtx)
		rtx.Must(err, "failed to create GCE client")
		hc = health.NewGCPChecker(gceClient, gcpmd)
	} else if kubernetesURL.URL == nil {
		hc = health.NewChecker(probe, ec)
	} else {
		k8s := health.MustNewKubernetesClient(kubernetesURL.URL, pod, node, namespace, kubernetesAuth)
		hc = health.NewCheckerK8S(probe, k8s, ec)
	}

	write(conn, hc, ldr)
}

// write starts a write loop to send health messages every
// HeartbeatPeriod.
func write(ws *connection.Conn, hc Checker, ldr *registration.Loader) {
	defer ws.Close()
	hbTicker := *time.NewTicker(heartbeatPeriod)
	defer hbTicker.Stop()

	// Register the channel to receive SIGTERM events.
	sigterm := make(chan os.Signal, 1)
	defer close(sigterm)
	signal.Notify(sigterm, syscall.SIGTERM)

	defer ldr.Ticker.Stop()

	for {
		select {
		case <-mainCtx.Done():
			log.Println("context cancelled")
			sendExitMessage(ws)
			return
		case <-sigterm:
			log.Println("received SIGTERM")
			sendExitMessage(ws)
			mainCancel()
			return
		case <-ldr.Ticker.C:
			reg, err := ldr.GetRegistration(mainCtx)
			if err != nil {
				log.Printf("could not load registration data, err: %v", err)
			}
			if reg != nil {
				sendMessage(ws, v2.HeartbeatMessage{Registration: reg}, "registration")
				log.Printf("updated registration to %v", reg)
			}
		case <-hbTicker.C:
			t := time.Now()
			score := getHealth(hc)
			healthMsg := v2.Health{Score: score}
			hbm := v2.HeartbeatMessage{Health: &healthMsg}
			sendMessage(ws, hbm, "health")

			// Record duration metric.
			fmtScore := fmt.Sprintf("%.1f", score)
			metrics.HealthTransmissionDuration.WithLabelValues(fmtScore).Observe(time.Since(t).Seconds())
		}
	}
}

func getHealth(hc Checker) float64 {
	ctx, cancel := context.WithTimeout(mainCtx, heartbeatPeriod)
	defer cancel()
	return hc.GetHealth(ctx)
}

func sendMessage(ws *connection.Conn, hbm v2.HeartbeatMessage, msgType string) {
	// If a new registration message was found, update the websocket's dial message.
	// The message is sent whenever the connection is restarted (i.e., once per hour in App Engine).
	if msgType == "registration" {
		ws.DialMessage = hbm
	}

	err := ws.WriteMessage(websocket.TextMessage, hbm)
	if err != nil {
		log.Printf("failed to write %s message, err: %v", msgType, err)
	}
}

func sendExitMessage(ws *connection.Conn) {
	// Notify the receiver that the health score should now be 0.
	hbm := v2.HeartbeatMessage{
		Health: &v2.Health{
			Score: 0,
		},
	}
	sendMessage(ws, hbm, "final health")
}