hmq/broker/retain.go

package broker

import (
	"sync"
)

type RetainList struct {
	sync.RWMutex
	root *rlevel
}
type rlevel struct {
	nodes map[string]*rnode
}
type rnode struct {
	next *rlevel
	msg  []byte
}
type RetainResult struct {
	msg [][]byte
}

func newRNode() *rnode {
	return &rnode{msg: make([]byte, 0, 4)}
}

func newRLevel() *rlevel {
	return &rlevel{nodes: make(map[string]*rnode)}
}

func NewRetainList() *RetainList {
	return &RetainList{root: newRLevel()}
}

func (r *RetainList) Insert(topic, buf []byte) error {

	tokens, err := PublishTopicCheckAndSpilt(topic)
	if err != nil {
		return err
	}
	// log.Info("insert tokens:", tokens)
	r.Lock()

	l := r.root
	var n *rnode
	for _, t := range tokens {
		n = l.nodes[t]
		if n == nil {
			n = newRNode()
			l.nodes[t] = n
		}
		if n.next == nil {
			n.next = newRLevel()
		}
		l = n.next
	}
	n.msg = buf
	r.Unlock()
	return nil
}

func (r *RetainList) Match(topic []byte) [][]byte {

	tokens, err := SubscribeTopicCheckAndSpilt(topic)
	if err != nil {
		return nil
	}
	results := &RetainResult{}

	r.Lock()
	l := r.root
	matchRLevel(l, tokens, results)
	r.Unlock()
	// log.Info("results: ", results)
	return results.msg

}
func matchRLevel(l *rlevel, toks []string, results *RetainResult) {
	var n *rnode
	for i, t := range toks {
		if l == nil {
			return
		}
		// log.Info("l info :", l.nodes)
		if t == "#" {
			for _, n := range l.nodes {
				n.GetAll(results)
			}
		}
		if t == "+" {
			for _, n := range l.nodes {
				if len(t[i+1:]) == 0 {
					results.msg = append(results.msg, n.msg)
				} else {
					matchRLevel(n.next, toks[i+1:], results)
				}
			}
		}

		n = l.nodes[t]
		if n != nil {
			l = n.next
		} else {
			l = nil
		}
	}
	if n != nil {
		results.msg = append(results.msg, n.msg)
	}
}

func (r *rnode) GetAll(results *RetainResult) {
	// log.Info("node 's message: ", string(r.msg))
	if r.msg != nil && string(r.msg) != "" {
		results.msg = append(results.msg, r.msg)
	}
	l := r.next
	for _, n := range l.nodes {
		n.GetAll(results)
	}
}