go-blocks/broker.go

package goblocks

import (
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/google/uuid"
)

type Message struct {
	Event string
	Data  any
}

type Client struct {
	ID       string
	Send     chan Message
	Close    chan struct{}
	IsActive bool
}

type SSEBroker struct {
	clients map[string]*Client
	mu      sync.RWMutex
	ticker  *time.Ticker
	done    chan struct{}
}

func newSSEBroker() *SSEBroker {
	s := &SSEBroker{
		clients: make(map[string]*Client),
		ticker:  time.NewTicker(time.Minute),
		done:    make(chan struct{}),
	}
	go s.run()
	return s
}

func (s *SSEBroker) run() {
	for {
		select {
		case <-s.ticker.C:
			s.Broadcast("ticker")
		case <-s.done:
			s.ticker.Stop()
			return
		}
	}
}

func (s *SSEBroker) registerClient(id string) *Client {
	s.mu.Lock()
	defer s.mu.Unlock()

	client := &Client{
		ID:       id,
		Send:     make(chan Message, 100),
		Close:    make(chan struct{}),
		IsActive: true,
	}
	s.clients[id] = client
	return client
}

func (s *SSEBroker) unregisterClient(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if client, exists := s.clients[id]; exists {
		close(client.Close)
		delete(s.clients, id)
	}
}

func (s *SSEBroker) Broadcast(event string, data ...any) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var dataValue any
	if len(data) > 0 {
		dataValue = data[0]
	}

	message := Message{
		Event: event,
		Data:  dataValue,
	}

	for _, client := range s.clients {
		if client.IsActive {
			select {
			case client.Send <- message:
			default:
				slog.Warn("client channel full, skipping message", "client_id", client.ID)
			}
		}
	}
}

func (s *SSEBroker) HandleSSE(w http.ResponseWriter, r *http.Request) {
	slog.Debug("SSE called")

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Transfer-Encoding", "chunked")

	clientID := uuid.New().String()
	client := s.registerClient(clientID)
	defer s.unregisterClient(clientID)

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
		return
	}

	slog.Info("SSE connection established", "client_id", clientID)

	fmt.Fprintf(w, "event: connection\ndata: %s\n\n", clientID)
	flusher.Flush()

	clientGone := r.Context().Done()

	for {
		select {
		case message := <-client.Send:
			var data string
			switch v := message.Data.(type) {
			case string:
				data = v
			case map[string]any, []any:
				jsonData, err := json.Marshal(v)
				if err != nil {
					slog.Error("error marshaling message", "error", err)
					continue
				}
				data = string(jsonData)
			default:
				jsonData, err := json.Marshal(v)
				if err != nil {
					slog.Error("error marshaling message", "error", err)
					continue
				}
				data = string(jsonData)
			}
			fmt.Fprintf(w, "event: %s\ndata: %s\n\n", message.Event, data)
		case <-client.Close:
			slog.Info("client closed", "client_id", clientID)
			return
		case <-clientGone:
			slog.Info("client gone", "client_id", clientID)
			return
		}

		flusher.Flush()
	}
}

func (s *SSEBroker) Shutdown() {
	close(s.done)
}