package server import ( "encoding/json" "errors" "fmt" "github.com/gorilla/websocket" "log" "net/http" "regexp" "simplemq/lib/random" "simplemq/lib/types" "sync" ) func dbg(s string, p ...interface{}) { log.Printf(s, p...) } type MQServer struct { pool sync.Map subs sync.Map once sync.Map } func (s *MQServer) Get(id string) *MQConn { raw, ok := s.pool.Load(id) if !ok { return nil } return raw.(*MQConn) } func (s *MQServer) Put(id string, conn *websocket.Conn) *MQConn { mqconn := &MQConn{id: id, conn: conn, server: s} s.pool.Store(id, mqconn) //dbg("MQServer.Put:: Got new conn: %s", mqconn.id) return mqconn } func (s *MQServer) Del(id string) { s.pool.Delete(id) //dbg("MQServer.Del:: Removing conn: %s", id) } func (s *MQServer) Process(msg *types.Msg) error { if msg.Topic == "" { err := "Cant proceed with empty topic!" dbg("Src:%s Id:%s Err:%s", msg.Src, msg.Id, err) return errors.New(err) } switch msg.Cmd { case types.CMD_PUB: //dbg("MQServer.Process:: Got new msg: %s.%s=>%s", msg.Cmd, msg.Src, msg.Topic) s.Range(func(m *MQConn) bool { m.Process(msg) return true }) case types.CMD_RPC: //dbg("MQServer.Process:: Got new RPC: Src:%s Reply:%s Topic:%s", msg.Src, msg.ReplyTo, msg.Topic) s.Range(func(m *MQConn) bool { done, _ := m.ProcessFirst(msg) return !done }) case types.CMD_RPCREPLY: //dbg("MQServer.Process:: Got new RPC Reply: Src:%s Topic:%s", msg.Src, msg.Topic) s.Range(func(m *MQConn) bool { done, _ := m.ProcessFirst(msg) return !done }) default: } return nil } func (s *MQServer) Range(h func(m *MQConn) bool) { s.pool.Range(func(key, value interface{}) bool { mq := s.Get(key.(string)) if mq != nil { cont := h(mq) return cont } return true }) } func NewServer() *MQServer { ret := &MQServer{} ret.pool = sync.Map{} ret.subs = sync.Map{} ret.once = sync.Map{} return ret } type MQConn struct { id string server *MQServer subs sync.Map mtx sync.Mutex conn *websocket.Conn } func (m *MQConn) Recv() error { ret := &types.Msg{} //dbg("MQConn.Recv:: Getting new message: %s", m.id) _, mbs, err := m.conn.ReadMessage() if err != nil { dbg("MQConn.Recv:: Error retrieving: %s: %s", m.id, err.Error()) m.Close() return err } //dbg("MQConn.Recv:: Got new message: %s", m.id) err = json.Unmarshal(mbs, ret) if err != nil { return err } else { switch ret.Cmd { case types.CMD_SUB: //dbg("MQConn.Recv:: Got new sub: %s=>%s", m.id, ret.Topic) m.Sub(ret.Topic) case types.CMD_UNSUB: //dbg("MQConn.Recv:: Got new unsub: %s=>%s", m.id, ret.Topic) m.Unsub(ret.Topic) case types.CMD_CLOSE: //dbg("MQConn.Recv:: Got close: %s=>%s", m.id) m.Close() default: ret.Src = m.id m.server.Process(ret) } } return nil } func (m *MQConn) Loop() { for { err := m.Recv() if err != nil { return } } } func (m *MQConn) Sub(t string) error { rx, err := regexp.Compile(fmt.Sprintf(`^%s$`, t)) if err != nil { return err } m.subs.Store(t, rx) //dbg("MQConn.Sub: Subs: %s", strings.Join(m.Subs(), ",")) return nil } func (m *MQConn) Subs() []string { ret := make([]string, 0) m.subs.Range(func(key, value interface{}) bool { ret = append(ret, key.(string)) return true }) return ret } func (m *MQConn) Unsub(s string) { m.subs.Delete(s) //dbg("MQConn.Unsub: Subs: %s", strings.Join(m.Subs(), ",")) } func (m *MQConn) Send(msg *types.Msg) error { return m.server.Process(msg) } func (m *MQConn) Process(msg *types.Msg) error { var err error = nil m.subs.Range(func(key, value interface{}) bool { rx := value.(*regexp.Regexp) if rx.MatchString(msg.Topic) { //dbg("MQConn.Process: RX:%s matches topic: %s", rx.String(), msg.Topic) //msg.Topic = key.(string) err = m.Write(msg) } return true }) return err } func (m *MQConn) ProcessFirst(msg *types.Msg) (bool, error) { var err error = nil var ret bool = false m.subs.Range(func(key, value interface{}) bool { rx := value.(*regexp.Regexp) if rx.MatchString(msg.Topic) { //msg.Topic = key.(string) //dbg("MQConn.ProcessFirst: RX:%s matches topic: %s", rx.String(), msg.Topic) err = m.Write(msg) ret = true return false } else { //dbg("MQConn.ProcessFirst: RX:%s DOES NOT match topic: %s", rx.String(), msg.Topic) } return true }) return ret, err } func (m *MQConn) Close() { m.conn.Close() m.server.Del(m.id) } func (c *MQConn) Write(m *types.Msg) error { c.mtx.Lock() defer c.mtx.Unlock() return c.conn.WriteJSON(m) } func Serve() { var upgrader = websocket.Upgrader{} server := NewServer() http.HandleFunc("/ws", func(writer http.ResponseWriter, request *http.Request) { id := random.Str(32) con, err := upgrader.Upgrade(writer, request, nil) if err != nil { http.Error(writer, err.Error(), http.StatusInternalServerError) return } mqconn := server.Put(id, con) con.SetCloseHandler(func(code int, text string) error { mqconn.Close() return nil }) mqconn.Loop() }) http.HandleFunc("/stats", func(writer http.ResponseWriter, request *http.Request) { //server.Range(func(m *MQConn) bool { // m. // return true //}) }) }