110 lines
2.1 KiB
Go
110 lines
2.1 KiB
Go
package topics
|
|
|
|
import (
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type Message struct {
|
|
ID string `json:"id"`
|
|
Topic string `json:"topic"`
|
|
Payload string `json:"payload"`
|
|
QOS byte `json:"qos"`
|
|
Retained bool `json:"retained"`
|
|
Timestamp time.Time `json:"timestamp"`
|
|
Size int `json:"size"`
|
|
}
|
|
|
|
type Tree struct {
|
|
mu sync.RWMutex
|
|
root *node
|
|
}
|
|
|
|
type node struct {
|
|
name string
|
|
fullName string
|
|
children map[string]*node
|
|
messageCount int
|
|
lastMessage *Message
|
|
}
|
|
|
|
type NodeSnapshot struct {
|
|
Name string `json:"name"`
|
|
FullName string `json:"fullName"`
|
|
MessageCount int `json:"messageCount"`
|
|
LastMessage *Message `json:"lastMessage,omitempty"`
|
|
Children []NodeSnapshot `json:"children"`
|
|
}
|
|
|
|
func NewTree() *Tree {
|
|
return &Tree{
|
|
root: &node{
|
|
name: "root",
|
|
fullName: "",
|
|
children: make(map[string]*node),
|
|
},
|
|
}
|
|
}
|
|
|
|
func (t *Tree) AddMessage(msg Message) {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
|
|
parts := strings.Split(msg.Topic, "/")
|
|
current := t.root
|
|
for i, part := range parts {
|
|
child, ok := current.children[part]
|
|
if !ok {
|
|
full := strings.Join(parts[:i+1], "/")
|
|
child = &node{
|
|
name: part,
|
|
fullName: full,
|
|
children: make(map[string]*node),
|
|
}
|
|
current.children[part] = child
|
|
}
|
|
current = child
|
|
current.messageCount++
|
|
if i == len(parts)-1 {
|
|
copyMsg := msg
|
|
current.lastMessage = ©Msg
|
|
}
|
|
}
|
|
}
|
|
|
|
func (t *Tree) Snapshot() NodeSnapshot {
|
|
t.mu.RLock()
|
|
defer t.mu.RUnlock()
|
|
|
|
return snapshotNode(t.root)
|
|
}
|
|
|
|
func snapshotNode(n *node) NodeSnapshot {
|
|
children := make([]NodeSnapshot, 0, len(n.children))
|
|
keys := make([]string, 0, len(n.children))
|
|
for key := range n.children {
|
|
keys = append(keys, key)
|
|
}
|
|
sort.Strings(keys)
|
|
for _, key := range keys {
|
|
child := n.children[key]
|
|
children = append(children, snapshotNode(child))
|
|
}
|
|
|
|
var last *Message
|
|
if n.lastMessage != nil {
|
|
copyMsg := *n.lastMessage
|
|
last = ©Msg
|
|
}
|
|
|
|
return NodeSnapshot{
|
|
Name: n.name,
|
|
FullName: n.fullName,
|
|
MessageCount: n.messageCount,
|
|
LastMessage: last,
|
|
Children: children,
|
|
}
|
|
}
|