update node-agent Heartbeat

Signed-off-by: gaoyuan <[email protected]>
kosmos-io · Nov 28, 2024 · 087b53c · 087b53c
1 parent c87aed4
commit 087b53c
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 11 deletions.
diff --git a/cmd/kubenest/node-agent/app/root.go b/cmd/kubenest/node-agent/app/root.go
@@ -1,7 +1,6 @@
 package app
 
 import (
-	"os"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -30,11 +29,12 @@ var RootCmd = &cobra.Command{
 
 func initConfig() {
 	// Tell Viper to automatically look for a .env file
-	viper.SetConfigFile("agent.env")
-	currentDir, _ := os.Getwd()
-	viper.AddConfigPath(currentDir)
-	viper.AddConfigPath("/srv/node-agent/agent.env")
-	viper.SetConfigType("toml")
+	//viper.SetConfigFile("agent.env")
+	viper.SetConfigFile("/srv/node-agent/agent.env")
+	//currentDir, _ := os.Getwd()
+	//viper.AddConfigPath(currentDir)
+	//viper.AddConfigPath("/srv/node-agent/agent.env")
+	//viper.SetConfigType("toml")
 	// If a agent.env file is found, read it in.
 	if err := viper.ReadInConfig(); err != nil {
 		log.Warnf("Load config file error, %s", err)

diff --git a/cmd/kubenest/node-agent/app/serve/serve.go b/cmd/kubenest/node-agent/app/serve/serve.go
@@ -2,6 +2,7 @@ package serve
 
 import (
 	"bufio"
+	"context"
 	"crypto/sha256"
 	"crypto/tls"
 	"encoding/base64"
@@ -20,10 +21,13 @@ import (
 
 	"github.com/creack/pty"
 	"github.com/gorilla/websocket"
+	"github.com/kosmos.io/kosmos/cmd/kubenest/node-agent/app/logger"
+	"github.com/kosmos.io/kosmos/pkg/generated/clientset/versioned"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
-
-	"github.com/kosmos.io/kosmos/cmd/kubenest/node-agent/app/logger"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/tools/clientcmd"
 )
 
 var (
@@ -36,6 +40,7 @@ var (
 	certFile string // SSL certificate file
 	keyFile  string // SSL key file
 	addr     string // server listen address
+	nodeName string // server nodename
 	log      = logger.GetLogger()
 )
 
@@ -52,9 +57,20 @@ func init() {
 	ServeCmd.PersistentFlags().StringVarP(&addr, "addr", "a", ":5678", "websocket service address")
 	ServeCmd.PersistentFlags().StringVarP(&certFile, "cert", "c", "cert.pem", "SSL certificate file")
 	ServeCmd.PersistentFlags().StringVarP(&keyFile, "key", "k", "key.pem", "SSL key file")
+	ServeCmd.PersistentFlags().StringVarP(&nodeName, "nodename", "n", "", "set nodename")
 }
 
 func serveCmdRun(_ *cobra.Command, _ []string) error {
+
+	//start heartbeatCheck Goroutine
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	if len(nodeName) == 0 {
+		nodeName = viper.GetString("NODE_NAME")
+	}
+	go heartbeatCheck(ctx, nodeName)
+
 	user := viper.GetString("WEB_USER")
 	password := viper.GetString("WEB_PASS")
 	port := viper.GetString("WEB_PORT")
@@ -65,11 +81,57 @@ func serveCmdRun(_ *cobra.Command, _ []string) error {
 	if port != "" {
 		addr = ":" + port
 	}
-	return Start(addr, certFile, keyFile, user, password)
+
+	return Start(ctx, addr, certFile, keyFile, user, password)
+	// err := Start(ctx, addr, certFile, keyFile, user, password)
+	// if err != nil {
+	// 	log.Errorf("Start server failed: %v", err)
+	// }
+	// return err
+}
+
+func heartbeatCheck(ctx context.Context, nodeName string) {
+	kubeconfigPath := "/srv/node-agent/kubeconfig"
+	config, err := clientcmd.BuildConfigFromFlags("", kubeconfigPath)
+	if err != nil {
+		log.Errorf("Failed to load kubeconfig: %v", err)
+	}
+
+	kosmosClient, err := versioned.NewForConfig(config)
+	if err != nil {
+		log.Errorf("Failed to get config: %v", err)
+	}
+
+	ticker := time.NewTicker(10 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			log.Infof("Heartbeat for node %s stopped", nodeName)
+			return
+		case <-ticker.C:
+			node, err := kosmosClient.KosmosV1alpha1().GlobalNodes().Get(ctx, nodeName, metav1.GetOptions{})
+			if err != nil {
+				log.Errorf("Failed to get node: %v", err)
+			}
+			heartbeatTime := metav1.Now()
+			node.Status.Conditions = []corev1.NodeCondition{
+				{
+					LastHeartbeatTime: heartbeatTime,
+				},
+			}
+			if _, err := kosmosClient.KosmosV1alpha1().GlobalNodes().UpdateStatus(ctx, node, metav1.UpdateOptions{}); err != nil {
+				log.Errorf("update node %s status for globalnode failed, %v", node.Name, err)
+			} else {
+				log.Infof("GlobalnodeHeartbeat: successfully updated global node %s, Status.Conditions: %+v", node.Name, node.Status.Conditions)
+			}
+		}
+	}
 }
 
 // start server
-func Start(addr, certFile, keyFile, user, password string) error {
+func Start(ctx context.Context, addr, certFile, keyFile, user, password string) error {
 	passwordHash := sha256.Sum256([]byte(password))
 
 	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
@@ -146,6 +208,12 @@ func Start(addr, certFile, keyFile, user, password string) error {
 		TLSConfig:         tlsConfig,
 		ReadHeaderTimeout: 10 * time.Second,
 	}
+
+	go func() {
+		time.Sleep(30 * time.Second)
+		_ = server.Close()
+	}()
+
 	err := server.ListenAndServeTLS("", "")
 	if err != nil {
 		log.Errorf("failed to start server %v", err)

diff --git a/deploy/virtual-cluster-operator.yml b/deploy/virtual-cluster-operator.yml
@@ -230,6 +230,16 @@ spec:
                 secretKeyRef:
                   name: node-agent-secret
                   key: password
+            - name: KUBECONFIG
+              valueFrom:
+                secretKeyRef:
+                  name: virtual-cluster-operator
+                  key: kubeconfig
+            - name: NODE_NAME
+              valueFrom:
+                fieldRef:
+                  apiVersion: v1
+                  fieldPath: spec.nodeName
           command: ["/bin/bash"]
           args:
             - "/app/init.sh"

diff --git a/hack/node-agent/agent.env b/hack/node-agent/agent.env
@@ -1,3 +1,4 @@
 WEB_USER={{WEB_USER}}
 WEB_PASS={{WEB_PASS}}
 WEB_PORT={{WEB_PORT}}
+NODE_NAME={{NODE_NAME}}
diff --git a/hack/node-agent/init.sh b/hack/node-agent/init.sh
@@ -1,9 +1,26 @@
 #!/usr/bin/env bash
 
+if [ -z "$KUBECONFIG" ]; then
+  echo "Error: KUBECONFIG_CONTENT is not set."
+  exit 1
+fi
+
+echo "$KUBECONFIG" > /app/kubeconfig
+if [ $? -ne 0 ]; then
+  echo "Error: Failed to write kubeconfig to $KUBECONFIG."
+  exit 1
+fi
+
+echo "KUBECONFIG has been written to $KUBECONFIG."
+
+
 WEB_USER="$WEB_USER" sed -i 's/^WEB_USER=.*/WEB_USER="'"$WEB_USER"'"/' /app/agent.env
 WEB_PASS="$WEB_PASS" sed -i 's/^WEB_PASS=.*/WEB_PASS="'"$WEB_PASS"'"/' /app/agent.env
 WEB_PORT="$WEB_PORT" sed -i 's/^WEB_PORT=.*/WEB_PORT="'"$WEB_PORT"'"/' /app/agent.env
+NODE_NAME="$NODE_NAME" sed -i 's/^NODE_NAME=.*/NODE_NAME="'"$NODE_NAME"'"/' /app/agent.env
+
 sha256sum /app/node-agent > /app/node-agent.sum
 sha256sum /host-path/node-agent >> /app/node-agent.sum
 rsync -avz /app/ /host-path/
-cp /app/node-agent.service /host-systemd/node-agent.service
+cp /app/node-agent.service /host-systemd/node-agent.service 
+