meshgo/stats/stats.go

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)