hub.go 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329
  1. package sse
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "net/http"
  6. "time"
  7. "git.eugeniocarvalho.dev/eugeniucarvalho/apicodegen/api/errs"
  8. "github.com/gomodule/redigo/redis"
  9. context "github.com/kataras/iris/v12/context"
  10. )
  11. type Event struct {
  12. Timestamp int64 `json:"timestamp"`
  13. Payload string `json:"payload"`
  14. Kind string `json:"kind"`
  15. }
  16. type ChannelClient struct {
  17. // Context context.Context
  18. Flusher http.Flusher
  19. Response http.ResponseWriter
  20. Request *http.Request
  21. }
  22. type Channel struct {
  23. ID string
  24. Clients map[*ChannelClient]bool
  25. }
  26. func (channel *Channel) RemoveClient(client *ChannelClient) {
  27. delete(channel.Clients, client)
  28. }
  29. func (channel *Channel) AddClient(client *ChannelClient) {
  30. channel.Clients[client] = true
  31. }
  32. func (channel *Channel) Emit(event *Event) {
  33. var (
  34. err error
  35. )
  36. for client := range channel.Clients {
  37. if event.Kind != "" {
  38. if _, err = fmt.Fprintf(client.Response, "event: %s\n", event.Kind); err != nil {
  39. fmt.Println(fmt.Sprintf("%+v", err))
  40. continue
  41. }
  42. }
  43. if _, err = fmt.Fprintf(client.Response, "data: %s\n\n", event.Payload); err != nil {
  44. fmt.Println(fmt.Sprintf("%+v", err))
  45. continue
  46. }
  47. client.Flusher.Flush()
  48. }
  49. }
  50. type SSEHub struct {
  51. // New client connections
  52. newClients chan chan []byte
  53. // Closed client connections
  54. closingClients chan chan []byte
  55. consume chan bool
  56. // Client connections registry
  57. Channels map[string]*Channel
  58. ChannelCollection string
  59. RedisPool *redis.Pool
  60. }
  61. type SSEOptions struct {
  62. URI string
  63. Password string
  64. ChannelCollection string
  65. }
  66. func NewSSEHub(options *SSEOptions) *SSEHub {
  67. return &SSEHub{
  68. newClients: make(chan chan []byte),
  69. closingClients: make(chan chan []byte),
  70. consume: make(chan bool),
  71. Channels: make(map[string]*Channel),
  72. ChannelCollection: options.ChannelCollection,
  73. RedisPool: &redis.Pool{
  74. // Maximum number of idle connections in the pool.
  75. MaxIdle: 80,
  76. // max number of connections
  77. MaxActive: 12000,
  78. // Dial is an application supplied function for creating and
  79. // configuring a connection.
  80. Dial: func() (redis.Conn, error) {
  81. conn, err := redis.Dial("tcp", options.URI)
  82. if err != nil {
  83. panic(err.Error())
  84. }
  85. if options.Password != "" {
  86. if _, err := conn.Do("AUTH", options.Password); err != nil {
  87. conn.Close()
  88. return nil, err
  89. }
  90. }
  91. return conn, err
  92. },
  93. },
  94. }
  95. }
  96. func (hub *SSEHub) GetChannel(channelID string) *Channel {
  97. if _, exist := hub.Channels[channelID]; !exist {
  98. channel := &Channel{
  99. Clients: make(map[*ChannelClient]bool),
  100. ID: channelID,
  101. }
  102. hub.Channels[channelID] = channel
  103. }
  104. return hub.Channels[channelID]
  105. }
  106. func (hub *SSEHub) Dispatch(event *Event, channels ...string) {
  107. var (
  108. exists bool
  109. err error
  110. conn = hub.RedisPool.Get()
  111. )
  112. defer conn.Close()
  113. for _, channel := range channels {
  114. if exists, err = redis.Bool(conn.Do("HEXISTS", hub.ChannelCollection, channel)); err != nil || !exists {
  115. continue
  116. }
  117. eventBytes, _ := json.Marshal(event)
  118. conn.Do("RPUSH", channel, string(eventBytes))
  119. }
  120. hub.consume <- true
  121. }
  122. func (hub *SSEHub) UpgradeConnection(ctx context.Context, channelId string) (err *errs.Error) {
  123. var (
  124. conn = hub.RedisPool.Get()
  125. payload string
  126. count int
  127. ok bool
  128. consuming bool
  129. redisErr error
  130. flusher http.Flusher
  131. response = ctx.ResponseWriter().Naive()
  132. request = ctx.Request()
  133. )
  134. defer conn.Close()
  135. if flusher, ok = response.(http.Flusher); !ok {
  136. err = errs.HTTPVersionNotSupported().Details(&errs.Detail{
  137. Dominio: "",
  138. Reason: "Streaming unsupported",
  139. Location: "hook.beforePersist",
  140. LocationType: "",
  141. })
  142. return
  143. }
  144. if _, redisErr = redis.Int(conn.Do("HSET", hub.ChannelCollection, channelId, true)); redisErr != nil {
  145. err = errs.Internal().Details(&errs.Detail{
  146. Dominio: "",
  147. Reason: "Fail on register channel",
  148. Location: "hook.beforePersist",
  149. LocationType: "",
  150. })
  151. return
  152. }
  153. header := response.Header()
  154. header.Set("Access-Control-Allow-Origin", "*")
  155. header.Set("Content-Type", "text/event-stream")
  156. header.Set("Cache-Control", "no-cache")
  157. header.Set("Connection", "keep-alive")
  158. flusher.Flush()
  159. client := &ChannelClient{
  160. Flusher: flusher,
  161. Response: response,
  162. Request: request,
  163. }
  164. channel := hub.GetChannel(channelId)
  165. channel.AddClient(client)
  166. finalizeMain := make(chan bool)
  167. // finalizePing := make(chan bool)
  168. finalized := false
  169. // Listen to connection close and when the entire request handler chain exits(this handler here) and un-register messageChan.
  170. notify := response.(http.CloseNotifier).CloseNotify()
  171. go func() {
  172. <-notify
  173. fmt.Println("SSE request end")
  174. defer func() {
  175. // recover from panic caused by writing to a closed channel
  176. recover()
  177. return
  178. }()
  179. if finalized {
  180. return
  181. }
  182. finalized = true
  183. // ctx.Application().Logger().Infof("notification.channel.disconect(%s)", channelId)
  184. // Remove this client from the map of connected clients
  185. // when this handler exits.
  186. redis.Int(conn.Do("HDEL", hub.ChannelCollection, channelId))
  187. // closesd = true
  188. finalizeMain <- finalized
  189. // finalizePing <- finalized
  190. channel.RemoveClient(client)
  191. }()
  192. ticker5Second := time.NewTicker(5 * time.Second)
  193. defer ticker5Second.Stop()
  194. flusher.Flush()
  195. reset:
  196. for {
  197. select {
  198. case <-ticker5Second.C:
  199. go func () {
  200. if !consuming {
  201. hub.consume <- true
  202. }
  203. }()
  204. case <-hub.consume:
  205. ctx.Application().Logger().Warnf("init.notification.loop 5s")
  206. if count, redisErr = redis.Int(conn.Do("LLEN", channelId)); count == 0 || redisErr != nil {
  207. continue reset
  208. }
  209. consuming = true
  210. for ; count > 0; count-- {
  211. if payload, redisErr = redis.String(conn.Do("LPOP", channelId)); redisErr != nil {
  212. ctx.Application().Logger().Errorf("NotificationError:", redisErr)
  213. } else {
  214. event := &Event{}
  215. json.Unmarshal([]byte(payload), event)
  216. channel.Emit(event)
  217. break
  218. }
  219. }
  220. consuming = false
  221. case <-finalizeMain:
  222. for {
  223. if !consuming {
  224. break
  225. }
  226. time.Sleep(time.Second * 2)
  227. }
  228. // ctx.Application().Logger().Infof("notification.finalize.disconect(%s)", channelId)
  229. close(finalizeMain)
  230. return
  231. }
  232. }
  233. }
  234. // ctx.OnClose(func() {
  235. // defer func() {
  236. // // recover from panic caused by writing to a closed channel
  237. // recover()
  238. // return
  239. // }()
  240. // if finalized {
  241. // return
  242. // }
  243. // finalized = true
  244. // // ctx.Application().Logger().Infof("notification.channel.disconect(%s)", channelId)
  245. // // Remove this client from the map of connected clients
  246. // // when this handler exits.
  247. // redis.Int(conn.Do("HDEL", hub.ChannelCollection, channelId))
  248. // // closesd = true
  249. // finalizeMain <- finalized
  250. // finalizePing <- finalized
  251. // channel.RemoveClient(client)
  252. // })
  253. // Set the headers related to event streaming, you can omit the "application/json" if you send plain text.
  254. // If you develop a go client, you must have: "Accept" : "application/json, text/event-stream" header as well.
  255. // header := ctx.ResponseWriter().Header()
  256. // header.Set("Access-Control-Allow-Origin", "*")
  257. // header.Set("Content-Type", "text/event-stream")
  258. // header.Set("Cache-Control", "no-cache")
  259. // header.Set("Connection", "keep-alive")
  260. // ctx.ContentType("text/event-stream")
  261. // ctx.Header("Cache-Control", "no-cache")
  262. // ctx.Header("Connection", "keep-alive")
  263. // ctx.Request().Response.Header.Set("Access-Control-Allow-Origin", "*")
  264. // Block waiting for messages broadcast on this connection's messageChan.
  265. // check notification in redis
  266. // go func() {
  267. // ctx.Application().Logger().Warnf("init.notification.loop %v", closesd)