├── .gitignore ├── LICENSE ├── README.md ├── example ├── README.md └── ticker_pusher.go ├── notifier.go └── sessionmanager └── manager.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, built with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | 14 | # Dependency directories (remove the comment below to include it) 15 | # vendor/ 16 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 maoqide 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ws-notifier 2 | pushing messages to mutiple clients with **one** backend goroutine. 3 | 4 | Blog: [http://maoqide.live/post/golang/golang-websocket-message-pushing/](http://maoqide.live/post/golang/golang-websocket-message-pushing/) 5 | -------------------------------------------------------------------------------- /example/README.md: -------------------------------------------------------------------------------- 1 | continuous pushing specific message to a group using ticker. -------------------------------------------------------------------------------- /example/ticker_pusher.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "log" 7 | "net/http" 8 | "strings" 9 | "time" 10 | 11 | notifier "github.com/maoqide/ws-notifier" 12 | ) 13 | 14 | func main() { 15 | fmt.Println("hello") 16 | 17 | http.HandleFunc("/mon", handleNotifierMon) 18 | http.HandleFunc("/", handleWsFunc) 19 | http.ListenAndServe(":8080", nil) 20 | } 21 | 22 | func handleNotifierMon(w http.ResponseWriter, r *http.Request) { 23 | ret, _ := json.MarshalIndent(DebugInfo(), "", "\t") 24 | w.Write(ret) 25 | return 26 | } 27 | 28 | func handleWsFunc(w http.ResponseWriter, r *http.Request) { 29 | prefix := "ticker_" 30 | n := notifier.Default() 31 | 32 | group := strings.Trim(r.RequestURI, "/") 33 | // should be random generated 34 | sessionID := "123456" 35 | 36 | groupID := prefix + group 37 | n.Notify(groupID, tickerWorker, time.Hour*24) 38 | n.HandleRequestWithKeys(w, r, map[string]interface{}{"group": groupID, "id": groupID + "_" + sessionID}) 39 | return 40 | } 41 | 42 | func tickerWorker(groupID string, sigChan chan int8, n *notifier.Notifier) error { 43 | worker := fmt.Sprintf("ticker_worker_%s_%d", groupID, time.Now().Unix()) 44 | fmt.Printf("worker: %s\n", worker) 45 | 46 | defer func() { 47 | select { 48 | case sigChan <- 0: 49 | log.Printf("ticker worker: %s exit", worker) 50 | case <-time.After(time.Second * 3): 51 | log.Printf("ticker worker: %s exit after 3s delaying", worker) 52 | } 53 | }() 54 | ticker := time.NewTicker(time.Second * 2) 55 | count := 0 56 | for { 57 | fmt.Println(count) 58 | select { 59 | case signal := <-sigChan: 60 | log.Printf("receice stop signal %d for ticker worker: %s", signal, worker) 61 | return nil 62 | case <-ticker.C: 63 | err := n.GroupBroadcast([]byte(fmt.Sprintf("%s: %d", groupID, count)), groupID) 64 | if err != nil { 65 | log.Printf("err: %v", err) 66 | } 67 | } 68 | count++ 69 | } 70 | } 71 | 72 | // NotifierState describe notifier states 73 | type NotifierState struct { 74 | Sessions map[string][]string `json:"sessions"` 75 | Workers []string `json:"workers"` 76 | } 77 | 78 | // DebugInfo return debug info for websocket notifier 79 | func DebugInfo() *NotifierState { 80 | state := NotifierState{} 81 | n := notifier.Default() 82 | state.Sessions = n.SessionManager.ShowSessions() 83 | state.Workers = n.ShowWorkers() 84 | return &state 85 | } 86 | -------------------------------------------------------------------------------- /notifier.go: -------------------------------------------------------------------------------- 1 | package notifier 2 | 3 | import ( 4 | "errors" 5 | "net/http" 6 | "sync" 7 | "time" 8 | 9 | "github.com/gorilla/websocket" 10 | "github.com/maoqide/melody" 11 | 12 | "github.com/maoqide/ws-notifier/sessionmanager" 13 | ) 14 | 15 | var notifier = New() 16 | 17 | // NotifyMessage is common struct for notifier 18 | type NotifyMessage struct { 19 | Type string `json:"type"` 20 | Code int32 `json:"code"` 21 | Message string `json:"message"` 22 | Data interface{} `json:"data"` 23 | } 24 | 25 | type workerFunc func(string, chan int8, *Notifier) error 26 | 27 | // Notifier wrapped websocket operation for notifier 28 | type Notifier struct { 29 | SessionManager *sessionmanager.SessionManager 30 | Melody *melody.Melody 31 | workers map[string]chan int8 32 | lock *sync.RWMutex 33 | } 34 | 35 | // Default return default initialized notifier, recommended. 36 | func Default() *Notifier { 37 | return notifier 38 | } 39 | 40 | // New creates a Notifier 41 | func New() *Notifier { 42 | m := melody.New() 43 | upgrader := websocket.Upgrader{} 44 | upgrader.HandshakeTimeout = time.Second * 2 45 | upgrader.CheckOrigin = func(r *http.Request) bool { 46 | return true 47 | } 48 | m.Upgrader = &upgrader 49 | 50 | sm := sessionmanager.New() 51 | m.HandleConnect(func(s *melody.Session) { 52 | g, exists := s.Get("group") 53 | if !exists { 54 | return 55 | } 56 | sm.Join(g.(string), s) 57 | if v, ok := s.Get("message"); ok { 58 | s.Write([]byte(v.(string))) 59 | s.Del("message") 60 | } 61 | }) 62 | 63 | m.HandleMessage(func(s *melody.Session, msg []byte) { 64 | }) 65 | m.HandlePong(func(s *melody.Session) { 66 | }) 67 | m.HandleDisconnect(func(s *melody.Session) { 68 | g, exists := s.Get("group") 69 | if !exists { 70 | s.Close() 71 | } 72 | sm.Release(g.(string), s) 73 | }) 74 | // m.Config.PongWait = 600 * time.Second 75 | // m.Config.PingPeriod = 601 * time.Second 76 | return &Notifier{ 77 | Melody: m, 78 | SessionManager: sm, 79 | workers: make(map[string]chan int8), 80 | lock: new(sync.RWMutex), 81 | } 82 | } 83 | 84 | // Notify start notify worker process 85 | func (n *Notifier) Notify(groupID string, f workerFunc, timeout time.Duration) error { 86 | n.lock.Lock() 87 | defer n.lock.Unlock() 88 | if _, ok := n.workers[groupID]; ok { 89 | return nil 90 | } 91 | n.workers[groupID] = make(chan int8) 92 | go f(groupID, n.workers[groupID], n) 93 | go func() { 94 | timer := time.NewTimer(timeout) 95 | for { 96 | select { 97 | case <-n.workers[groupID]: 98 | n.ReleaseWorker(groupID) 99 | // close all sessions of the group if worker exited, reconnection is needed from frontend. 100 | n.CloseGroupWithMsg(groupID, []byte{}) 101 | return 102 | case <-timer.C: 103 | n.workers[groupID] <- 1 104 | // kill worker goroutine when all session closed. 105 | case <-time.Tick(time.Second): 106 | if n.GroupLen(groupID) == 0 { 107 | n.workers[groupID] <- 2 108 | } 109 | } 110 | } 111 | }() 112 | return nil 113 | } 114 | 115 | // ReleaseWorker release worker for a group, usually called from a workerFunc when goroutine exited 116 | func (n *Notifier) ReleaseWorker(groupID string) { 117 | n.lock.Lock() 118 | defer n.lock.Unlock() 119 | if _, ok := n.workers[groupID]; !ok { 120 | return 121 | } 122 | close(n.workers[groupID]) 123 | delete(n.workers, groupID) 124 | return 125 | } 126 | 127 | // GroupBroadcast broadcast message to a group 128 | func (n *Notifier) GroupBroadcast(msg []byte, groupID string) error { 129 | if n.GroupLen(groupID) == 0 { 130 | return errors.New("no active session") 131 | } 132 | 133 | return n.Melody.BroadcastFilter(msg, func(s *melody.Session) bool { 134 | group, exists := s.Get("group") 135 | return exists && (group.(string) == groupID) 136 | }) 137 | } 138 | 139 | // Broadcast broadcast message to all 140 | func (n *Notifier) Broadcast(msg []byte) error { 141 | return n.Melody.Broadcast(msg) 142 | } 143 | 144 | // Close close all websocket connections 145 | func (n *Notifier) Close() error { 146 | n.SessionManager = nil 147 | for _, c := range n.workers { 148 | close(c) 149 | } 150 | return n.Melody.Close() 151 | } 152 | 153 | // CloseWithMsg close all websocket connections with messages. 154 | // Use the FormatCloseMessage function to format a proper close message payload. 155 | func (n *Notifier) CloseWithMsg(msg []byte) error { 156 | n.SessionManager = nil 157 | for _, c := range n.workers { 158 | close(c) 159 | } 160 | return n.Melody.CloseWithMsg(msg) 161 | } 162 | 163 | // CloseGroupWithMsg close all websocket connections of a group with messages. 164 | // Use the FormatCloseMessage function to format a proper close message payload. 165 | func (n *Notifier) CloseGroupWithMsg(groupID string, msg []byte) error { 166 | sessions := n.SessionManager.GetSessions(groupID) 167 | for _, s := range sessions { 168 | s.CloseWithMsg(msg) 169 | } 170 | return nil 171 | } 172 | 173 | // IsClosed return status of websocket 174 | func (n *Notifier) IsClosed() bool { 175 | return n.Melody.IsClosed() 176 | } 177 | 178 | // Len return the number of connected sessions. 179 | func (n *Notifier) Len() int { 180 | return n.Melody.Len() 181 | } 182 | 183 | // GroupLen return the number of connected sessions of a group. 184 | func (n *Notifier) GroupLen(groupID string) int { 185 | return len(n.SessionManager.GetSessions(groupID)) 186 | } 187 | 188 | // HandleRequest upgrades http requests to websocket connections and dispatches them to be handled by the melody instance. 189 | func (n *Notifier) HandleRequest(w http.ResponseWriter, r *http.Request) error { 190 | return n.Melody.HandleRequest(w, r) 191 | } 192 | 193 | // HandleRequestWithKeys does the same as HandleRequest but populates session.Keys with keys. 194 | func (n *Notifier) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]interface{}) error { 195 | return n.Melody.HandleRequestWithKeys(w, r, keys) 196 | } 197 | 198 | // FormatCloseMessage formats closeCode and text as a WebSocket close message. 199 | func FormatCloseMessage(closeCode int, text string) []byte { 200 | return websocket.FormatCloseMessage(closeCode, text) 201 | } 202 | 203 | // ShowWorkers shows all workers 204 | func (n *Notifier) ShowWorkers() []string { 205 | n.lock.RLock() 206 | defer n.lock.RUnlock() 207 | res := make([]string, 0) 208 | for w := range n.workers { 209 | res = append(res, w) 210 | } 211 | return res 212 | } 213 | -------------------------------------------------------------------------------- /sessionmanager/manager.go: -------------------------------------------------------------------------------- 1 | package sessionmanager 2 | 3 | import ( 4 | "sync" 5 | 6 | "github.com/maoqide/melody" 7 | ) 8 | 9 | // SessionManager manage sessions 10 | type SessionManager struct { 11 | sessionGroup map[string]map[*melody.Session]bool 12 | rwLock *sync.RWMutex 13 | } 14 | 15 | // Join join a session into a group 16 | func (sm *SessionManager) Join(groupID string, session *melody.Session) { 17 | sm.rwLock.Lock() 18 | defer sm.rwLock.Unlock() 19 | if _, ok := sm.sessionGroup[groupID]; !ok { 20 | sm.sessionGroup[groupID] = make(map[*melody.Session]bool) 21 | } 22 | sm.sessionGroup[groupID][session] = true 23 | return 24 | } 25 | 26 | // Release release a session from a group 27 | func (sm *SessionManager) Release(groupID string, session *melody.Session) { 28 | sm.rwLock.Lock() 29 | defer sm.rwLock.Unlock() 30 | if !session.IsClosed() { 31 | session.Close() 32 | } 33 | delete(sm.sessionGroup[groupID], session) 34 | if len(sm.sessionGroup[groupID]) == 0 { 35 | delete(sm.sessionGroup, groupID) 36 | } 37 | return 38 | } 39 | 40 | // GetSessions get sessions of a group 41 | func (sm *SessionManager) GetSessions(groupID string) []*melody.Session { 42 | sessions := make([]*melody.Session, 0) 43 | sm.rwLock.RLock() 44 | defer sm.rwLock.RUnlock() 45 | for s := range sm.sessionGroup[groupID] { 46 | sessions = append(sessions, s) 47 | } 48 | return sessions 49 | } 50 | 51 | // New create a new SessionManager 52 | func New() *SessionManager { 53 | return &SessionManager{ 54 | sessionGroup: make(map[string]map[*melody.Session]bool), 55 | rwLock: new(sync.RWMutex), 56 | } 57 | } 58 | 59 | // ShowSessions shows all sessions 60 | func (sm *SessionManager) ShowSessions() map[string][]string { 61 | res := make(map[string][]string) 62 | sm.rwLock.RLock() 63 | defer sm.rwLock.RUnlock() 64 | for group, sessions := range sm.sessionGroup { 65 | res[group] = make([]string, 0) 66 | for s := range sessions { 67 | id, exists := s.Get("id") 68 | if !exists { 69 | id = "unknown" 70 | } 71 | res[group] = append(res[group], id.(string)) 72 | } 73 | } 74 | return res 75 | } 76 | --------------------------------------------------------------------------------