atridad/goth.stack
Template
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Simplified the SSE system!

Browse Source
This commit is contained in:
Atridad Lahiji
2024-06-04 15:40:29 -06:00
parent b1d82ca609
commit 924906431f
10 changed files with 83 additions and 347 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
lib/logging.go
View File

@ -1,27 +1,29 @@
package lib
import "github.com/fatih/color"
import (
"github.com/fatih/color"
)
// Error logging
var red = color.New(color.FgRed)
var LogError = red.Add(color.Bold)
var red = color.New(color.FgRed, color.Bold)
var LogError = red
// Info logging
var cyan = color.New(color.FgCyan)
var LogInfo = cyan.Add(color.Bold)
var cyan = color.New(color.FgCyan, color.Bold)
var LogInfo = cyan
// Success logging
var green = color.New(color.FgGreen)
var LogSuccess = green.Add(color.Bold)
var green = color.New(color.FgGreen, color.Bold)
var LogSuccess = green
// Warning logging
var yellow = color.New(color.FgYellow)
var LogWarning = yellow.Add(color.Bold)
var yellow = color.New(color.FgYellow, color.Bold)
var LogWarning = yellow
// Debug logging
var magenta = color.New(color.FgMagenta)
var LogDebug = magenta.Add(color.Bold)
var magenta = color.New(color.FgMagenta, color.Bold)
var LogDebug = magenta
// Custom logging
var white = color.New(color.FgWhite)
var LogCustom = white.Add(color.Bold)
var white = color.New(color.FgWhite, color.Bold)
var LogCustom = white

86
lib/pubsub/adapters/localpubsub.go
View File

@ -1,86 +0,0 @@
package adapters
import (
"context"
"sync"
"atri.dad/lib"
"atri.dad/lib/pubsub"
)
type LocalPubSub struct {
subscribers map[string][]chan pubsub.Message
lock sync.RWMutex
}
type LocalPubSubMessage struct {
messages <-chan pubsub.Message
}
func (ps *LocalPubSub) SubscribeToChannel(channel string) (pubsub.PubSubMessage, error) {
ps.lock.Lock()
defer ps.lock.Unlock()
if ps.subscribers == nil {
ps.subscribers = make(map[string][]chan pubsub.Message)
}
ch := make(chan pubsub.Message, 100)
ps.subscribers[channel] = append(ps.subscribers[channel], ch)
lib.LogInfo.Printf("[PUBSUB/LOCAL] Subscribed to channel %s\n", channel)
return &LocalPubSubMessage{messages: ch}, nil
}
func (ps *LocalPubSub) PublishToChannel(channel string, message string) error {
subscribers, ok := ps.subscribers[channel]
if !ok {
lib.LogWarning.Printf("\n[PUBSUB/LOCAL] No subscribers for channel %s\n", channel)
return nil
}
ps.lock.Lock()
defer ps.lock.Unlock()
lib.LogInfo.Printf("\n[PUBSUB/LOCAL] Publishing message to channel %s: %s\n", channel, message)
for _, ch := range subscribers {
ch <- pubsub.Message{Payload: message}
}
return nil
}
func (ps *LocalPubSub) UnsubscribeFromChannel(channel string, ch <-chan pubsub.Message) {
ps.lock.Lock()
defer ps.lock.Unlock()
subscribers := ps.subscribers[channel]
for i, subscriber := range subscribers {
if subscriber == ch {
// Remove the subscriber from the slice
subscribers = append(subscribers[:i], subscribers[i+1:]...)
break
}
}
if len(subscribers) == 0 {
delete(ps.subscribers, channel)
} else {
ps.subscribers[channel] = subscribers
}
}
func (m *LocalPubSubMessage) ReceiveMessage(ctx context.Context) (*pubsub.Message, error) {
for {
select {
case <-ctx.Done():
// The client has disconnected. Stop trying to send messages.
return nil, ctx.Err()
case msg := <-m.messages:
// A message has been received. Send it to the client.
lib.LogInfo.Printf("\n[PUBSUB/LOCAL] Received message: %s\n", msg.Payload)
return &msg, nil
}
}
}

71
lib/pubsub/adapters/redispubsub.go
View File

@ -1,71 +0,0 @@
package adapters
import (
"context"
"os"
"atri.dad/lib"
"atri.dad/lib/pubsub"
"github.com/joho/godotenv"
"github.com/redis/go-redis/v9"
)
var RedisClient *redis.Client
type RedisPubSubMessage struct {
pubsub *redis.PubSub
}
// RedisPubSub is a Redis implementation of the PubSub interface.
type RedisPubSub struct {
Client *redis.Client
}
func NewRedisClient() (*redis.Client, error) {
if RedisClient != nil {
return RedisClient, nil
}
godotenv.Load(".env")
redis_url := os.Getenv("REDIS_URL")
opts, err := redis.ParseURL(redis_url)
if err != nil {
return nil, err
}
lib.LogInfo.Printf("\n[PUBSUB/REDIS]Connecting to Redis at %s\n", opts.Addr)
RedisClient = redis.NewClient(opts)
return RedisClient, nil
}
func (m *RedisPubSubMessage) ReceiveMessage(ctx context.Context) (*pubsub.Message, error) {
msg, err := m.pubsub.ReceiveMessage(ctx)
if err != nil {
return nil, err
}
lib.LogInfo.Printf("\n[PUBSUB/REDIS] Received message: %s\n", msg.Payload)
return &pubsub.Message{Payload: msg.Payload}, nil
}
func (ps *RedisPubSub) SubscribeToChannel(channel string) (pubsub.PubSubMessage, error) {
pubsub := ps.Client.Subscribe(context.Background(), channel)
_, err := pubsub.Receive(context.Background())
if err != nil {
return nil, err
}
lib.LogInfo.Printf("\n[PUBSUB/REDIS] Subscribed to channel %s\n", channel)
return &RedisPubSubMessage{pubsub: pubsub}, nil
}
func (r *RedisPubSub) PublishToChannel(channel string, message string) error {
err := r.Client.Publish(context.Background(), channel, message).Err()
if err != nil {
return err
}
lib.LogInfo.Printf("\n[PUBSUB/REDIS] Publishing message to channel %s: %s\n", channel, message)
return nil
}

16
lib/pubsub/interface.go
View File

@ -1,16 +0,0 @@
package pubsub
import "context"
type Message struct {
Payload string
}
type PubSubMessage interface {
ReceiveMessage(ctx context.Context) (*Message, error)
}
type PubSub interface {
SubscribeToChannel(channel string) (PubSubMessage, error)
PublishToChannel(channel string, message string) error
}

98
lib/sse.go
View File

@ -1,14 +1,6 @@
package lib
import (
"context"
"fmt"
"net/http"
"sync"
"atri.dad/lib/pubsub"
"github.com/labstack/echo/v4"
)
import "sync"
type SSEServerType struct {
clients map[string]map[chan string]bool
@ -37,6 +29,8 @@ func (s *SSEServerType) AddClient(channel string, client chan string) {
s.clients[channel] = make(map[chan string]bool)
}
s.clients[channel][client] = true
LogInfo.Printf("\nClient connected to channel %s\n", channel)
}
func (s *SSEServerType) RemoveClient(channel string, client chan string) {
@ -47,6 +41,8 @@ func (s *SSEServerType) RemoveClient(channel string, client chan string) {
if len(s.clients[channel]) == 0 {
delete(s.clients, channel)
}
LogInfo.Printf("\nClient disconnected from channel %s\n", channel)
}
func (s *SSEServerType) ClientCount(channel string) int {
@ -56,87 +52,15 @@ func (s *SSEServerType) ClientCount(channel string) int {
return len(s.clients[channel])
}
func SendSSE(ctx context.Context, messageBroker pubsub.PubSub, channel string, message string) error {
LogInfo.Printf("Sending SSE message to channel %s", channel)
func SendSSE(channel string, message string) error {
SSEServer.mu.Lock()
defer SSEServer.mu.Unlock()
errCh := make(chan error, 1)
go func() {
select {
case <-ctx.Done():
errCh <- ctx.Err()
default:
err := messageBroker.PublishToChannel(channel, message)
errCh <- err
}
}()
err := <-errCh
if err != nil {
LogError.Printf("Error sending SSE message: %v", err)
return err
for client := range SSEServer.clients[channel] {
client <- message
}
LogSuccess.Printf("SSE message sent successfully")
LogDebug.Printf("\nMessage broadcast on channel %s: %s\n", channel, message)
return nil
}
func SetSSEHeaders(c echo.Context) {
c.Response().Header().Set(echo.HeaderContentType, "text/event-stream")
c.Response().Header().Set(echo.HeaderConnection, "keep-alive")
c.Response().Header().Set(echo.HeaderCacheControl, "no-cache")
c.Response().Header().Set("X-Accel-Buffering", "no")
}
func HandleIncomingMessages(c echo.Context, pubsub pubsub.PubSubMessage, client chan string) {
if c.Response().Writer == nil {
LogError.Printf("Cannot handle incoming messages: ResponseWriter is nil")
return
}
var mutex sync.Mutex
for {
// Receive messages using the context of the request, which is cancelled when the client disconnects
msg, err := pubsub.ReceiveMessage(c.Request().Context())
if err != nil {
if err == context.Canceled {
// Log when the client disconnects and stop the message forwarding
LogInfo.Printf("Client disconnected, stopping message forwarding")
return
}
// Log errors other than client disconnection
LogError.Printf("Failed to receive message: %v", err)
return
}
// Prepare the data string to be sent as an SSE
data := fmt.Sprintf("data: %s\n\n", msg.Payload)
// Locking before writing to the response writer to avoid concurrent write issues
mutex.Lock()
if c.Response().Writer != nil {
_, err := c.Response().Write([]byte(data))
if err != nil {
// Log failure to write and unlock before returning
LogError.Printf("Failed to write message: %v", err)
mutex.Unlock()
return
}
// Flush the response if possible
flusher, ok := c.Response().Writer.(http.Flusher)
if ok {
flusher.Flush()
} else {
LogError.Println("Failed to flush: ResponseWriter does not implement http.Flusher")
}
} else {
LogError.Println("Failed to write: ResponseWriter is nil")
}
// Ensure the mutex is unlocked after processing each message
mutex.Unlock()
}
}