up

stats/stats.go

@@ -0,0 +1,165 @@
+package stats
+
+import (
+	"bytes"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	mqtt "github.com/mochi-mqtt/server/v2"
+	"github.com/mochi-mqtt/server/v2/packets"
+)
+
+// ---------------------------------------------------------------------------
+// 数据结构
+// ---------------------------------------------------------------------------
+
+// ClientInfo 当前在线客户端信息
+type ClientInfo struct {
+	ID          string    `json:"id"`
+	RemoteAddr  string    `json:"remote_addr"`
+	Username    string    `json:"username"`
+	ConnectedAt time.Time `json:"connected_at"`
+	SubsCount   int       `json:"subs_count"`
+}
+
+// Stats 当前统计快照
+type Stats struct {
+	Connections    int64              `json:"connections"`     // 当前连接数
+	MessagesTotal int64              `json:"messages_total"`  // 累计消息数（所有主题）
+	MessagesMsh   int64              `json:"messages_msh"`    // msh/# 消息数
+	Uptime        int64              `json:"uptime"`          // 服务运行时长（秒）
+	Clients       []ClientInfo       `json:"clients"`         // 在线客户端列表
+	Topics        map[string]int64   `json:"topics"`         // 各主题消息数
+}
+
+// ---------------------------------------------------------------------------
+// 全局统计（atomic + mutex 无锁热点路径）
+// ---------------------------------------------------------------------------
+
+var (
+	connections    atomic.Int64
+	messagesTotal  atomic.Int64
+	messagesMsh    atomic.Int64
+	startTime      = time.Now()
+	clientsMu      sync.RWMutex
+	clients        = make(map[string]ClientInfo)        // clientID → info
+	subs           = make(map[string][]string)          // clientID → []filter
+	topicsMu       sync.RWMutex
+	topics         = make(map[string]int64)             // topic → count
+)
+
+// GetStats 返回当前统计快照（只读副本）
+func GetStats() Stats {
+	clientsMu.RLock()
+	clientList := make([]ClientInfo, 0, len(clients))
+	for id, info := range clients {
+		info.SubsCount = len(subs[id])
+		clientList = append(clientList, info)
+	}
+	clientsMu.RUnlock()
+
+	topicsMu.RLock()
+	topicsCopy := make(map[string]int64, len(topics))
+	for k, v := range topics {
+		topicsCopy[k] = v
+	}
+	topicsMu.RUnlock()
+
+	return Stats{
+		Connections:    connections.Load(),
+		MessagesTotal:  messagesTotal.Load(),
+		MessagesMsh:    messagesMsh.Load(),
+		Uptime:         int64(time.Since(startTime).Seconds()),
+		Clients:        clientList,
+		Topics:         topicsCopy,
+	}
+}
+
+// ---------------------------------------------------------------------------
+// Hook 实现
+// ---------------------------------------------------------------------------
+
+// Hook 收集 MQTT 运行统计
+type Hook struct {
+	mqtt.HookBase
+}
+
+func (h *Hook) ID() string { return "meshgo-stats" }
+
+func (h *Hook) Provides(b byte) bool {
+	return bytes.Contains([]byte{
+		mqtt.OnSessionEstablished,
+		mqtt.OnDisconnect,
+		mqtt.OnPublish,
+		mqtt.OnSubscribe,
+		mqtt.OnUnsubscribe,
+	}, []byte{b})
+}
+
+// OnSessionEstablished 客户端连接成功
+func (h *Hook) OnSessionEstablished(cl *mqtt.Client, pk packets.Packet) {
+	username := string(pk.Connect.Username)
+	if username == "" {
+		username = "(anonymous)"
+	}
+	clientsMu.Lock()
+	clients[cl.ID] = ClientInfo{
+		ID:          cl.ID,
+		RemoteAddr:  cl.Net.Remote,
+		Username:    username,
+		ConnectedAt: time.Now(),
+	}
+	subs[cl.ID] = []string{}
+	clientsMu.Unlock()
+	connections.Add(1)
+}
+
+// OnDisconnect 客户端断开
+func (h *Hook) OnDisconnect(cl *mqtt.Client, err error, expire bool) {
+	clientsMu.Lock()
+	delete(clients, cl.ID)
+	delete(subs, cl.ID)
+	clientsMu.Unlock()
+	connections.Add(-1)
+}
+
+// OnPublish 收到发布消息
+func (h *Hook) OnPublish(cl *mqtt.Client, pk packets.Packet) (packets.Packet, error) {
+	messagesTotal.Add(1)
+	if len(pk.TopicName) >= 4 && pk.TopicName[:4] == "msh/" {
+		messagesMsh.Add(1)
+	}
+	topicsMu.Lock()
+	topics[pk.TopicName]++
+	topicsMu.Unlock()
+	return pk, nil
+}
+
+// OnSubscribe 客户端订阅
+func (h *Hook) OnSubscribe(cl *mqtt.Client, pk packets.Packet) packets.Packet {
+	clientsMu.Lock()
+	for _, f := range pk.Filters {
+		subs[cl.ID] = append(subs[cl.ID], f.Filter)
+	}
+	clientsMu.Unlock()
+	return pk
+}
+
+// OnUnsubscribe 客户端取消订阅
+func (h *Hook) OnUnsubscribe(cl *mqtt.Client, pk packets.Packet) packets.Packet {
+	clientsMu.Lock()
+	for _, f := range pk.Filters {
+		filters := subs[cl.ID]
+		for i, ff := range filters {
+			if ff == f.Filter {
+				subs[cl.ID] = append(filters[:i], filters[i+1:]...)
+				break
+			}
+		}
+	}
+	clientsMu.Unlock()
+	return pk
+}
+
+var _ mqtt.Hook = (*Hook)(nil)