mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-12 02:06:37 +08:00
[优化] Phase2
This commit is contained in:
ryan
@@ -3,7 +3,6 @@ package service
|
||||
import (
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"sync"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -22,11 +21,6 @@ type AgentWSInboundMessage struct {
|
||||
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"`
|
||||
@@ -35,133 +29,29 @@ type AgentWSBroadcastResult struct {
|
||||
FailedNodes []string `json:"failed_nodes"`
|
||||
}
|
||||
|
||||
type AgentWSClient struct {
|
||||
nodeID string
|
||||
send chan AgentWSOutboundMessage
|
||||
done chan struct{}
|
||||
once sync.Once
|
||||
var DefaultAgentWSHub = NewWSHub("agent")
|
||||
|
||||
func RegisterAgentWSClient(nodeID string) *WSClient {
|
||||
return DefaultAgentWSHub.Register(nodeID)
|
||||
}
|
||||
|
||||
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 UnregisterAgentWSClient(client *WSClient) {
|
||||
DefaultAgentWSHub.Unregister(client)
|
||||
}
|
||||
|
||||
func DisconnectAgentWSClient(nodeID string) {
|
||||
defaultAgentWSHub.mu.Lock()
|
||||
client := defaultAgentWSHub.clients[nodeID]
|
||||
if client != nil {
|
||||
delete(defaultAgentWSHub.clients, nodeID)
|
||||
}
|
||||
count := len(defaultAgentWSHub.clients)
|
||||
defaultAgentWSHub.mu.Unlock()
|
||||
|
||||
if client != nil {
|
||||
client.Close()
|
||||
slog.Debug("agent ws connection forcefully disconnected", "node_id", nodeID, "client_count", count)
|
||||
}
|
||||
DefaultAgentWSHub.Disconnect(nodeID)
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
return DefaultAgentWSHub.IsConnected(nodeID)
|
||||
}
|
||||
|
||||
func SendAgentWSSettings(nodeID string, settings *AgentSettings) bool {
|
||||
if settings == nil {
|
||||
return false
|
||||
}
|
||||
return sendAgentWSMessage(nodeID, AgentWSOutboundMessage{
|
||||
return DefaultAgentWSHub.SendMessage(nodeID, WSMessage{
|
||||
Type: AgentWSMessageTypeSettings,
|
||||
Payload: settings,
|
||||
})
|
||||
@@ -171,7 +61,7 @@ func SendAgentWSActiveConfig(nodeID string, activeConfig *ActiveConfigMeta) bool
|
||||
if activeConfig == nil {
|
||||
return false
|
||||
}
|
||||
return sendAgentWSMessage(nodeID, AgentWSOutboundMessage{
|
||||
return DefaultAgentWSHub.SendMessage(nodeID, WSMessage{
|
||||
Type: AgentWSMessageTypeActiveConfig,
|
||||
Payload: activeConfig,
|
||||
})
|
||||
@@ -181,60 +71,37 @@ func SendAgentWSForceSyncConfig(nodeID string, activeConfig *ActiveConfigMeta) b
|
||||
if activeConfig == nil {
|
||||
return false
|
||||
}
|
||||
return sendAgentWSMessage(nodeID, AgentWSOutboundMessage{
|
||||
return DefaultAgentWSHub.SendMessage(nodeID, WSMessage{
|
||||
Type: AgentWSMessageTypeForceSyncConfig,
|
||||
Payload: activeConfig,
|
||||
})
|
||||
}
|
||||
|
||||
func SendAgentWSPong(nodeID string) bool {
|
||||
return sendAgentWSMessage(nodeID, AgentWSOutboundMessage{
|
||||
return DefaultAgentWSHub.SendMessage(nodeID, WSMessage{
|
||||
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
|
||||
return AgentWSBroadcastResult{}
|
||||
}
|
||||
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{
|
||||
res := DefaultAgentWSHub.Broadcast(WSMessage{
|
||||
Type: AgentWSMessageTypeActiveConfig,
|
||||
Payload: activeConfig,
|
||||
})
|
||||
|
||||
result := AgentWSBroadcastResult{
|
||||
Version: activeConfig.Version,
|
||||
Checksum: activeConfig.Checksum,
|
||||
ClientCount: res.ClientCount,
|
||||
SuccessCount: res.SuccessCount,
|
||||
FailedNodes: res.FailedIDs,
|
||||
}
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user