[新增] 实现 Agent WebSocket 连接升级功能，支持状态上报和配置广播

2026-10-07 08:06:37 +08:00 · 2026-05-29 09:52:34 +08:00
parent 4257b6fd5a
commit 79ed8fd6ab
29 changed files with 1128 additions and 95 deletions
@@ -89,13 +89,14 @@ type AgentConfigResponse struct {
 }

 type AgentSettings struct {
-	HeartbeatInterval   int    `json:"heartbeat_interval"`
-	AutoUpdate          bool   `json:"auto_update"`
-	UpdateRepo          string `json:"update_repo"`
-	UpdateNow           bool   `json:"update_now"`
-	UpdateChannel       string `json:"update_channel"`
-	UpdateTag           string `json:"update_tag"`
-	RestartOpenrestyNow bool   `json:"restart_openresty_now"`
+	HeartbeatInterval       int    `json:"heartbeat_interval"`
+	WebsocketUpgradeEnabled bool   `json:"websocket_upgrade_enabled"`
+	AutoUpdate              bool   `json:"auto_update"`
+	UpdateRepo              string `json:"update_repo"`
+	UpdateNow               bool   `json:"update_now"`
+	UpdateChannel           string `json:"update_channel"`
+	UpdateTag               string `json:"update_tag"`
+	RestartOpenrestyNow     bool   `json:"restart_openresty_now"`
 }

 type ActiveConfigMeta struct {
@@ -130,7 +131,7 @@ type NodeView struct {
 	OpenrestyMessage          string     `json:"openresty_message"`
 	Status                    string     `json:"status"`
 	CurrentVersion            string     `json:"current_version"`
-	LastSeenAt                time.Time  `json:"last_seen_at"`
+	LastSeenAt                any        `json:"last_seen_at"`
 	LastError                 string     `json:"last_error"`
 	LatestApplyResult         string     `json:"latest_apply_result"`
 	LatestApplyMessage        string     `json:"latest_apply_message"`
@@ -177,20 +178,32 @@ func HeartbeatNode(node *model.Node, payload AgentNodePayload) (*HeartbeatRespon
 		return nil, err
 	}
 	return &HeartbeatResponse{
-		Node: node,
-		AgentSettings: &AgentSettings{
-			HeartbeatInterval:   common.AgentHeartbeatInterval,
-			AutoUpdate:          node.AutoUpdateEnabled,
-			UpdateRepo:          common.AgentUpdateRepo,
-			UpdateNow:           updateNow,
-			UpdateChannel:       updateChannel.String(),
-			UpdateTag:           updateTag,
-			RestartOpenrestyNow: restartOpenrestyNow,
-		},
-		ActiveConfig: activeConfig,
+		Node:          node,
+		AgentSettings: buildAgentSettings(node, updateNow, updateChannel.String(), updateTag, restartOpenrestyNow),
+		ActiveConfig:  activeConfig,
 	}, nil
 }

+func buildAgentSettings(node *model.Node, updateNow bool, updateChannel string, updateTag string, restartOpenrestyNow bool) *AgentSettings {
+	autoUpdate := false
+	if node != nil {
+		autoUpdate = node.AutoUpdateEnabled
+	}
+	if strings.TrimSpace(updateChannel) == "" {
+		updateChannel = ReleaseChannelStable.String()
+	}
+	return &AgentSettings{
+		HeartbeatInterval:       common.AgentHeartbeatInterval,
+		WebsocketUpgradeEnabled: common.AgentWebsocketUpgradeEnabled,
+		AutoUpdate:              autoUpdate,
+		UpdateRepo:              common.AgentUpdateRepo,
+		UpdateNow:               updateNow,
+		UpdateChannel:           updateChannel,
+		UpdateTag:               strings.TrimSpace(updateTag),
+		RestartOpenrestyNow:     restartOpenrestyNow,
+	}
+}
+
 func GetActiveConfigMetaForAgent() (*ActiveConfigMeta, error) {
 	version, err := model.GetActiveConfigVersion()
 	if err != nil {
@@ -438,6 +451,9 @@ func computeNodeStatus(node *model.Node) string {
 	if node == nil {
 		return NodeStatusOffline
 	}
+	if IsAgentWSConnected(node.NodeID) {
+		return NodeStatusOnline
+	}
 	if node.LastSeenAt.IsZero() {
 		return NodeStatusPending
 	}
@@ -0,0 +1,226 @@
+package service
+
+import (
+	"encoding/json"
+	"log/slog"
+	"sync"
+)
+
+const (
+	AgentWSMessageTypeStatus       = "status"
+	AgentWSMessageTypeSettings     = "settings"
+	AgentWSMessageTypeActiveConfig = "active_config"
+	AgentWSMessageTypePing         = "ping"
+	AgentWSMessageTypePong         = "pong"
+
+	AgentWSConnectedLastSeenValue = "__OPENFLARE_WS_CONNECTED__"
+)
+
+type AgentWSInboundMessage struct {
+	Type    string          `json:"type"`
+	Payload json.RawMessage `json:"payload,omitempty"`
+}
+
+type AgentWSOutboundMessage struct {
+	Type    string `json:"type"`
+	Payload any    `json:"payload,omitempty"`
+}
+
+type AgentWSBroadcastResult struct {
+	Version      string   `json:"version"`
+	Checksum     string   `json:"checksum"`
+	ClientCount  int      `json:"client_count"`
+	SuccessCount int      `json:"success_count"`
+	FailedNodes  []string `json:"failed_nodes"`
+}
+
+type AgentWSClient struct {
+	nodeID string
+	send   chan AgentWSOutboundMessage
+	done   chan struct{}
+	once   sync.Once
+}
+
+func (client *AgentWSClient) NodeID() string {
+	if client == nil {
+		return ""
+	}
+	return client.nodeID
+}
+
+func (client *AgentWSClient) Messages() <-chan AgentWSOutboundMessage {
+	if client == nil {
+		return nil
+	}
+	return client.send
+}
+
+func (client *AgentWSClient) Done() <-chan struct{} {
+	if client == nil {
+		return nil
+	}
+	return client.done
+}
+
+func (client *AgentWSClient) Send(message AgentWSOutboundMessage) bool {
+	if client == nil {
+		return false
+	}
+	select {
+	case <-client.done:
+		return false
+	case client.send <- message:
+		return true
+	default:
+		return false
+	}
+}
+
+func (client *AgentWSClient) Close() {
+	if client == nil {
+		return
+	}
+	client.once.Do(func() {
+		close(client.done)
+	})
+}
+
+type agentWSHub struct {
+	mu      sync.RWMutex
+	clients map[string]*AgentWSClient
+}
+
+var defaultAgentWSHub = &agentWSHub{
+	clients: make(map[string]*AgentWSClient),
+}
+
+func RegisterAgentWSClient(nodeID string) *AgentWSClient {
+	client := &AgentWSClient{
+		nodeID: nodeID,
+		send:   make(chan AgentWSOutboundMessage, 16),
+		done:   make(chan struct{}),
+	}
+	defaultAgentWSHub.mu.Lock()
+	if existing := defaultAgentWSHub.clients[nodeID]; existing != nil {
+		slog.Debug("agent ws replacing existing connection", "node_id", nodeID)
+		existing.Close()
+	}
+	defaultAgentWSHub.clients[nodeID] = client
+	count := len(defaultAgentWSHub.clients)
+	defaultAgentWSHub.mu.Unlock()
+	slog.Debug("agent ws connection registered", "node_id", nodeID, "client_count", count)
+	return client
+}
+
+func UnregisterAgentWSClient(client *AgentWSClient) {
+	if client == nil {
+		return
+	}
+	defaultAgentWSHub.mu.Lock()
+	if current := defaultAgentWSHub.clients[client.nodeID]; current == client {
+		delete(defaultAgentWSHub.clients, client.nodeID)
+	}
+	count := len(defaultAgentWSHub.clients)
+	defaultAgentWSHub.mu.Unlock()
+	client.Close()
+	slog.Debug("agent ws connection unregistered", "node_id", client.nodeID, "client_count", count)
+}
+
+func IsAgentWSConnected(nodeID string) bool {
+	defaultAgentWSHub.mu.RLock()
+	client := defaultAgentWSHub.clients[nodeID]
+	defaultAgentWSHub.mu.RUnlock()
+	if client == nil {
+		return false
+	}
+	select {
+	case <-client.done:
+		return false
+	default:
+		return true
+	}
+}
+
+func AgentWSClientCount() int {
+	defaultAgentWSHub.mu.RLock()
+	defer defaultAgentWSHub.mu.RUnlock()
+	return len(defaultAgentWSHub.clients)
+}
+
+func SendAgentWSSettings(nodeID string, settings *AgentSettings) bool {
+	if settings == nil {
+		return false
+	}
+	return sendAgentWSMessage(nodeID, AgentWSOutboundMessage{
+		Type:    AgentWSMessageTypeSettings,
+		Payload: settings,
+	})
+}
+
+func SendAgentWSActiveConfig(nodeID string, activeConfig *ActiveConfigMeta) bool {
+	if activeConfig == nil {
+		return false
+	}
+	return sendAgentWSMessage(nodeID, AgentWSOutboundMessage{
+		Type:    AgentWSMessageTypeActiveConfig,
+		Payload: activeConfig,
+	})
+}
+
+func SendAgentWSPong(nodeID string) bool {
+	return sendAgentWSMessage(nodeID, AgentWSOutboundMessage{
+		Type: AgentWSMessageTypePong,
+	})
+}
+
+func sendAgentWSMessage(nodeID string, message AgentWSOutboundMessage) bool {
+	defaultAgentWSHub.mu.RLock()
+	client := defaultAgentWSHub.clients[nodeID]
+	defaultAgentWSHub.mu.RUnlock()
+	if client == nil {
+		return false
+	}
+	ok := client.Send(message)
+	if !ok {
+		slog.Debug("agent ws send queued message failed", "node_id", nodeID, "type", message.Type)
+	}
+	return ok
+}
+
+func BroadcastAgentWSActiveConfig(activeConfig *ActiveConfigMeta) AgentWSBroadcastResult {
+	result := AgentWSBroadcastResult{}
+	if activeConfig == nil {
+		slog.Debug("agent ws broadcast skipped because active config is nil")
+		return result
+	}
+	result.Version = activeConfig.Version
+	result.Checksum = activeConfig.Checksum
+
+	defaultAgentWSHub.mu.RLock()
+	clients := make([]*AgentWSClient, 0, len(defaultAgentWSHub.clients))
+	for _, client := range defaultAgentWSHub.clients {
+		clients = append(clients, client)
+	}
+	defaultAgentWSHub.mu.RUnlock()
+
+	result.ClientCount = len(clients)
+	message := AgentWSOutboundMessage{
+		Type:    AgentWSMessageTypeActiveConfig,
+		Payload: activeConfig,
+	}
+	for _, client := range clients {
+		if client.Send(message) {
+			result.SuccessCount++
+			continue
+		}
+		result.FailedNodes = append(result.FailedNodes, client.NodeID())
+	}
+	slog.Debug("agent ws broadcast active config",
+		"version", result.Version,
+		"checksum", result.Checksum,
+		"client_count", result.ClientCount,
+		"success_count", result.SuccessCount,
+		"failed_nodes", result.FailedNodes,
+	)
+	return result
+}
@@ -385,6 +385,10 @@ func PublishConfigVersion(createdBy string, force bool) (*ReleaseResult, error)
 		}
 		return nil, err
 	}
+	BroadcastAgentWSActiveConfig(&ActiveConfigMeta{
+		Version:  record.Version,
+		Checksum: record.Checksum,
+	})
 	return &ReleaseResult{
 		Version: record,
 		Routes:  bundle.Routes,
@@ -409,6 +413,10 @@ func ActivateConfigVersion(id uint) (*model.ConfigVersion, error) {
 		return nil, err
 	}
 	version.IsActive = true
+	BroadcastAgentWSActiveConfig(&ActiveConfigMeta{
+		Version:  version.Version,
+		Checksum: version.Checksum,
+	})
 	return version, nil
 }

@@ -48,23 +48,23 @@ type DashboardTrends struct {
 }

 type DashboardNodeHealth struct {
-	ID                  uint      `json:"id"`
-	NodeID              string    `json:"node_id"`
-	Name                string    `json:"name"`
-	GeoName             string    `json:"geo_name"`
-	GeoLatitude         *float64  `json:"geo_latitude"`
-	GeoLongitude        *float64  `json:"geo_longitude"`
-	Status              string    `json:"status"`
-	OpenrestyStatus     string    `json:"openresty_status"`
-	CurrentVersion      string    `json:"current_version"`
-	LastSeenAt          time.Time `json:"last_seen_at"`
-	ActiveEventCount    int       `json:"active_event_count"`
-	CPUUsagePercent     float64   `json:"cpu_usage_percent"`
-	MemoryUsagePercent  float64   `json:"memory_usage_percent"`
-	StorageUsagePercent float64   `json:"storage_usage_percent"`
-	RequestCount        int64     `json:"request_count"`
-	ErrorCount          int64     `json:"error_count"`
-	UniqueVisitorCount  int64     `json:"unique_visitor_count"`
+	ID                  uint     `json:"id"`
+	NodeID              string   `json:"node_id"`
+	Name                string   `json:"name"`
+	GeoName             string   `json:"geo_name"`
+	GeoLatitude         *float64 `json:"geo_latitude"`
+	GeoLongitude        *float64 `json:"geo_longitude"`
+	Status              string   `json:"status"`
+	OpenrestyStatus     string   `json:"openresty_status"`
+	CurrentVersion      string   `json:"current_version"`
+	LastSeenAt          any      `json:"last_seen_at"`
+	ActiveEventCount    int      `json:"active_event_count"`
+	CPUUsagePercent     float64  `json:"cpu_usage_percent"`
+	MemoryUsagePercent  float64  `json:"memory_usage_percent"`
+	StorageUsagePercent float64  `json:"storage_usage_percent"`
+	RequestCount        int64    `json:"request_count"`
+	ErrorCount          int64    `json:"error_count"`
+	UniqueVisitorCount  int64    `json:"unique_visitor_count"`
 }

 func GetDashboardOverview() (*DashboardOverviewView, error) {
@@ -139,7 +139,7 @@ func GetDashboardOverview() (*DashboardOverviewView, error) {
 			Status:           computedStatus,
 			OpenrestyStatus:  node.OpenrestyStatus,
 			CurrentVersion:   node.CurrentVersion,
-			LastSeenAt:       node.LastSeenAt,
+			LastSeenAt:       nodeViewLastSeenAt(node),
 			ActiveEventCount: len(nodeActiveEvents),
 		}

@@ -275,7 +275,7 @@ func buildNodeView(node *model.Node) *NodeView {
 		OpenrestyMessage:          strings.TrimSpace(node.OpenrestyMessage),
 		Status:                    status,
 		CurrentVersion:            node.CurrentVersion,
-		LastSeenAt:                node.LastSeenAt,
+		LastSeenAt:                nodeViewLastSeenAt(node),
 		LastError:                 node.LastError,
 		CreatedAt:                 node.CreatedAt,
 		UpdatedAt:                 node.UpdatedAt,
@@ -288,6 +288,16 @@ func buildNodeView(node *model.Node) *NodeView {
 	return view
 }

+func nodeViewLastSeenAt(node *model.Node) any {
+	if node != nil && IsAgentWSConnected(node.NodeID) {
+		return AgentWSConnectedLastSeenValue
+	}
+	if node == nil {
+		return time.Time{}
+	}
+	return node.LastSeenAt
+}
+
 func normalizeNodeInput(input NodeInput) (string, string, string, *float64, *float64, bool, error) {
 	name := strings.TrimSpace(input.Name)
 	ip := strings.TrimSpace(input.IP)
@@ -615,6 +615,43 @@ func TestListNodeViewsDoesNotPersistComputedStatus(t *testing.T) {
 	}
 }

+func TestAgentWSConnectionMarksNodeViewOnline(t *testing.T) {
+	node := &model.Node{
+		ID:         99,
+		NodeID:     "node-ws-view",
+		Name:       "edge-ws",
+		IP:         "10.0.0.9",
+		Status:     NodeStatusOffline,
+		LastSeenAt: time.Now().Add(-common.NodeOfflineThreshold - time.Minute),
+	}
+	client := RegisterAgentWSClient(node.NodeID)
+	defer UnregisterAgentWSClient(client)
+
+	view := buildNodeView(node)
+	if view.Status != NodeStatusOnline {
+		t.Fatalf("expected websocket-connected node to be online, got %s", view.Status)
+	}
+	if view.LastSeenAt != AgentWSConnectedLastSeenValue {
+		t.Fatalf("expected websocket special last_seen_at, got %#v", view.LastSeenAt)
+	}
+
+	result := BroadcastAgentWSActiveConfig(&ActiveConfigMeta{
+		Version:  "20260529-001",
+		Checksum: "checksum-ws",
+	})
+	if result.ClientCount != 1 || result.SuccessCount != 1 || len(result.FailedNodes) != 0 {
+		t.Fatalf("unexpected broadcast result: %+v", result)
+	}
+	select {
+	case message := <-client.Messages():
+		if message.Type != AgentWSMessageTypeActiveConfig {
+			t.Fatalf("unexpected websocket message type: %s", message.Type)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("expected websocket broadcast message")
+	}
+}
+
 func TestHeartbeatNodePersistsObservabilityPayload(t *testing.T) {
 	setupServiceTestDB(t)
 	withFakeAccessLogGeoProvider(t, &geoip.GeoInfo{