hub.go 6.6 KB

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