apicodegen/api/sse/hub.go

package sse

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"git.eugeniocarvalho.dev/eugeniucarvalho/apicodegen/api/errs"
	"github.com/gomodule/redigo/redis"
	context "github.com/kataras/iris/v12/context"
)

type Event struct {
	Timestamp int64  `json:"timestamp"`
	Payload   string `json:"payload"`
	Kind      string `json:"kind"`
}

type ChannelClient struct {
	// Context context.Context
	Flusher  http.Flusher
	Response http.ResponseWriter
	Request  *http.Request
}

type Channel struct {
	ID      string
	Clients map[*ChannelClient]bool
}

func (channel *Channel) RemoveClient(client *ChannelClient) {
	delete(channel.Clients, client)
}

func (channel *Channel) AddClient(client *ChannelClient) {
	channel.Clients[client] = true
}

func (channel *Channel) Emit(event *Event) {

	var (
		err error
	)
	for client := range channel.Clients {

		if event.Kind != "" {
			if _, err = fmt.Fprintf(client.Response, "event: %s\n", event.Kind); err != nil {
				fmt.Println(fmt.Sprintf("%+v", err))
				continue
			}
		}

		if _, err = fmt.Fprintf(client.Response, "data: %s\n\n", event.Payload); err != nil {
			fmt.Println(fmt.Sprintf("%+v", err))
			continue
		}

		client.Flusher.Flush()
	}
}

type SSEHub struct {
	// New client connections
	newClients chan chan []byte

	// Closed client connections
	closingClients chan chan []byte

	consume chan bool
	// Client connections registry
	Channels map[string]*Channel

	ChannelCollection string

	RedisPool *redis.Pool
}

type SSEOptions struct {
	URI               string
	Password          string
	ChannelCollection string
}

func NewSSEHub(options *SSEOptions) *SSEHub {

	return &SSEHub{
		newClients:        make(chan chan []byte),
		closingClients:    make(chan chan []byte),
		consume:    make(chan bool),
		Channels:          make(map[string]*Channel),
		ChannelCollection: options.ChannelCollection,
		RedisPool: &redis.Pool{
			// Maximum number of idle connections in the pool.
			MaxIdle: 80,
			// max number of connections
			MaxActive: 12000,
			// Dial is an application supplied function for creating and
			// configuring a connection.
			Dial: func() (redis.Conn, error) {
				conn, err := redis.Dial("tcp", options.URI)
				if err != nil {
					panic(err.Error())
				}

				if options.Password != "" {
					if _, err := conn.Do("AUTH", options.Password); err != nil {
						conn.Close()
						return nil, err
					}
				}
				return conn, err
			},
		},
	}
}

func (hub *SSEHub) GetChannel(channelID string) *Channel {

	if _, exist := hub.Channels[channelID]; !exist {
		channel := &Channel{
			Clients: make(map[*ChannelClient]bool),
			ID:      channelID,
		}
		hub.Channels[channelID] = channel
	}

	return hub.Channels[channelID]
}

func (hub *SSEHub) Dispatch(event *Event, channels ...string) {
	var (
		exists bool
		err    error
		conn   = hub.RedisPool.Get()
	)
	defer conn.Close()
	
	eventBytes, _ := json.Marshal(event)

	for _, channel := range channels {

		if exists, err = redis.Bool(conn.Do("HEXISTS", hub.ChannelCollection, channel)); err != nil || !exists {
			continue
		}

		conn.Do("RPUSH", channel, string(eventBytes))
	}
	go func(){ hub.consume <- true }()
}

func (hub *SSEHub) Close(channels ...string) {
	var (
		exists bool
		count int
		err    error
		conn   = hub.RedisPool.Get()
	)

	for _, channel := range channels {
		if exists, err = redis.Bool(conn.Do("HEXISTS", hub.ChannelCollection, channel)); err != nil || !exists {
			continue
		}
		go func(channelID string) {
			for {
				if count, err = redis.Int(conn.Do("LLEN", channelID)); count > 0 {
					continue
				}
				conn.Do("HDEL", hub.ChannelCollection, channelID)
			}
		}(channel)
	}
}

func (hub *SSEHub) UpgradeConnection(ctx context.Context, channelId string) (err *errs.Error) {

	var (
		conn     = hub.RedisPool.Get()
		payload  string
		count    int
		ok       bool
		consuming bool
		redisErr error
		flusher  http.Flusher
		response = ctx.ResponseWriter().Naive()
		request  = ctx.Request()
	)

	defer conn.Close()

	if flusher, ok = response.(http.Flusher); !ok {
		err = errs.HTTPVersionNotSupported().Details(&errs.Detail{
			Dominio:      "",
			Reason:       "Streaming unsupported",
			Location:     "hook.beforePersist",
			LocationType: "",
		})
		return
	}

	if _, redisErr = redis.Int(conn.Do("HSET", hub.ChannelCollection, channelId, true)); redisErr != nil {
		err = errs.Internal().Details(&errs.Detail{
			Dominio:      "",
			Reason:       "Fail on register channel",
			Location:     "hook.beforePersist",
			LocationType: "",
		})
		return
	}

	header := response.Header()
	header.Set("Access-Control-Allow-Origin", "*")
	header.Set("Content-Type", "text/event-stream")
	header.Set("Cache-Control", "no-cache")
	header.Set("Connection", "keep-alive")

	flusher.Flush()

	client := &ChannelClient{
		Flusher:  flusher,
		Response: response,
		Request:  request,
	}

	channel := hub.GetChannel(channelId)
	channel.AddClient(client)

	finalizeMain := make(chan bool)
	// finalizePing := make(chan bool)
	finalized := false

	// Listen to connection close and when the entire request handler chain exits(this handler here) and un-register messageChan.

	notify := response.(http.CloseNotifier).CloseNotify()

	go func() {
		<-notify
		fmt.Println("SSE request end")
		defer func() {
			// recover from panic caused by writing to a closed channel
			recover()
			return
		}()

		if finalized {
			return
		}

		finalized = true
		// ctx.Application().Logger().Infof("notification.channel.disconect(%s)", channelId)
		// Remove this client from the map of connected clients
		// when this handler exits.
		redis.Int(conn.Do("HDEL", hub.ChannelCollection, channelId))
		// closesd = true
		finalizeMain <- finalized
		// finalizePing <- finalized
		channel.RemoveClient(client)
	}()

	ticker5Second := time.NewTicker(5 * time.Second)
	defer ticker5Second.Stop()

	flusher.Flush()
	
reset:
	for {
		select {
		case <-ticker5Second.C:
			go func ()  {
				if !consuming {
					hub.consume <- true
				}	
			}()
		case <-hub.consume:
			// ctx.Application().Logger().Warnf("init.notification.loop 5s")
			
			if count, redisErr = redis.Int(conn.Do("LLEN", channelId)); count == 0 || redisErr != nil {
				continue reset
			}
			
			consuming = true

			for ; count > 0; count-- {

				if payload, redisErr = redis.String(conn.Do("LPOP", channelId)); redisErr != nil {
					// ctx.Application().Logger().Errorf("NotificationError:", redisErr)
				} else {
					event := &Event{}
					json.Unmarshal([]byte(payload), event)
					channel.Emit(event)
					break
				}
			}

			consuming = false

		case <-finalizeMain:
			
			for {
				if !consuming {
					break
				}
				time.Sleep(time.Second * 2)
			}
			// ctx.Application().Logger().Infof("notification.finalize.disconect(%s)", channelId)
			close(finalizeMain)
			return
		}
	}
}

// ctx.OnClose(func() {
// 	defer func() {
// 		// recover from panic caused by writing to a closed channel
// 		recover()
// 		return
// 	}()

// 	if finalized {
// 		return
// 	}

// 	finalized = true
// 	// ctx.Application().Logger().Infof("notification.channel.disconect(%s)", channelId)
// 	// Remove this client from the map of connected clients
// 	// when this handler exits.
// 	redis.Int(conn.Do("HDEL", hub.ChannelCollection, channelId))
// 	// closesd = true
// 	finalizeMain <- finalized
// 	finalizePing <- finalized
// 	channel.RemoveClient(client)
// })

// Set the headers related to event streaming, you can omit the "application/json" if you send plain text.
// If you develop a go client, you must have: "Accept" : "application/json, text/event-stream" header as well.
// header := ctx.ResponseWriter().Header()
// header.Set("Access-Control-Allow-Origin", "*")
// header.Set("Content-Type", "text/event-stream")
// header.Set("Cache-Control", "no-cache")
// header.Set("Connection", "keep-alive")
// ctx.ContentType("text/event-stream")
// ctx.Header("Cache-Control", "no-cache")
// ctx.Header("Connection", "keep-alive")

// ctx.Request().Response.Header.Set("Access-Control-Allow-Origin", "*")

// Block waiting for messages broadcast on this connection's messageChan.
// check notification in redis
// go func() {

// ctx.Application().Logger().Warnf("init.notification.loop %v", closesd)