notify/server.go

// Package notify is a package which provide common tcp/udp/unix socket service
package notify

import (
	"net"
	"strings"
	"time"

	"b612.me/starainrt"
)

var builder *starainrt.StarQueue

func init() {
	builder = starainrt.NewQueue()
}

// StarNotifyS 为Server端
type StarNotifyS struct {
	// Queue 是用来处理收发信息的简单消息队列
	Queue *starainrt.StarQueue
	// FuncLists 记录了被通知项所记录的函数
	FuncLists      map[string]func(SMsg) string
	defaultFunc    func(SMsg) string
	serverStopSign chan int
	notifychan     chan int
	connPool       map[string]net.Conn
}

// SMsg 指明当前服务端被通知的关键字
type SMsg struct {
	Conn  net.Conn
	Key   string
	Value string
}

// Reply 用于向client端回复数据
func (nmsg *SMsg) Reply(msg string) error {
	_, err := nmsg.Conn.Write(builder.BuildMessage(nmsg.Key + "||" + msg))
	return err
}

// Send 用于向client端发送key-value数据
func (nmsg *SMsg) Send(key, value string) error {
	_, err := nmsg.Conn.Write(builder.BuildMessage(key + "||" + value))
	return err
}

func (star *StarNotifyS) starinits() {
	star.serverStopSign, star.notifychan = make(chan int, 1), make(chan int, 5)
	star.Queue = starainrt.NewQueue()
	star.FuncLists = make(map[string]func(SMsg) string)
	star.connPool = make(map[string]net.Conn)
}

// NewNotifyS 开启一个新的Server端通知
func NewNotifyS(netype, value string) (*StarNotifyS, error) {
	var star StarNotifyS
	star.starinits()
	listener, err := net.Listen(netype, value)
	if err != nil {
		return nil, err
	}
	go star.notify()
	go func() {
		for {
			go func() {
				<-star.serverStopSign
				star.notifychan <- 1
				star.notifychan <- 1
				for k, v := range star.connPool {
					v.Close()
					delete(star.connPool, k)
				}
				listener.Close()
				return
			}()

			conn, err := listener.Accept()
			if err != nil {
				select {
				case <-star.notifychan:
					listener.Close()
					return
				default:
					continue
				}
			}
			go func(conn net.Conn) {
				for {
					buf := make([]byte, 8192)
					n, err := conn.Read(buf)
					if n != 0 {
						star.Queue.ParseMessage(buf[0:n], conn)
					}
					if err != nil {
						conn.Close()
						delete(star.connPool, conn.RemoteAddr().String())
						break
					}
				}
			}(conn)
			star.connPool[conn.RemoteAddr().String()] = conn
		}
	}()
	return &star, nil
}

// SetNotify 用于设置通知关键词的调用函数
func (star *StarNotifyS) SetNotify(name string, data func(SMsg) string) {
	star.FuncLists[name] = data
}

// SetDefaultNotify 用于设置默认关键词的调用函数
func (star *StarNotifyS) SetDefaultNotify(name string, data func(SMsg) string) {
	star.defaultFunc = data
}

func (star *StarNotifyS) notify() {
	for {
		select {
		case <-star.notifychan:
			return
		default:
		}
		data, err := star.Queue.RestoreOne()
		if err != nil {
			time.Sleep(time.Millisecond * 20)
			continue
		}
		key, value := analyseData(data.Msg)
		if msg, ok := star.FuncLists[key]; ok {
			sdata := msg(SMsg{data.Conn.(net.Conn), key, value})
			if sdata == "" {
				continue
			}
			sdata = key + "||" + sdata
			data.Conn.(net.Conn).Write(star.Queue.BuildMessage(sdata))
		} else {
			if star.defaultFunc != nil {
				sdata := star.defaultFunc(SMsg{data.Conn.(net.Conn), key, value})
				if sdata == "" {
					continue
				}
				sdata = key + "||" + sdata
				data.Conn.(net.Conn).Write(star.Queue.BuildMessage(sdata))
			}
		}
	}
}

func analyseData(msg string) (key, value string) {
	slice := strings.SplitN(msg, "||", 2)
	if len(slice) == 1 {
		return msg, ""
	}
	return slice[0], slice[1]
}

// ServerStop 用于终止Server端运行
func (star *StarNotifyS) ServerStop() {
	star.serverStopSign <- 0
}