posti-dm/server/httputil/serversentevents.go

package httputil

import (
	"fmt"
	"log"
	"net/http"
)

type SSEHandler struct {
	clients map[chan string]bool

	Connected    func(chan string)
	Disconnected func(chan string)
}

func (sse *SSEHandler) init() {
	if sse.clients == nil {
		sse.clients = make(map[chan string]bool)
	}
}

func (sse *SSEHandler) Broadcast(message string) {
	sse.init()

	for client := range sse.clients {
		client <- message
	}
}

func (sse *SSEHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	sse.init()

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

	client := make(chan string)
	sse.clients[client] = true

	log.Printf(`New connection`)

	if sse.Connected != nil {
		go sse.Connected(client)
	}

	defer func() {
		log.Printf(`Connection closed`)
		close(client)
		delete(sse.clients, client)
		if sse.Disconnected != nil {
			go sse.Disconnected(client)
		}
	}()

	flusher, _ := w.(http.Flusher)
	for {
		select {
		case message, ok := <-client:
			if !ok {
				return
			}

			fmt.Fprintf(w, "data: %s\n\n", message)
			flusher.Flush()
		case <-r.Context().Done():
			return
		}
	}
}

// func HandleSSE(handler func(broadcast, single chan<- string)) http.Handler {
// 	broadcast := make(chan string)

// 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// 		w.Header().Set("Content-Type", "text/event-stream")
// 		w.Header().Set("Cache-Control", "no-cache")
// 		w.Header().Set("Connection", "keep-alive")
// 		w.Header().Set("Access-Control-Allow-Origin", "*")

// 		client := make(chan string)
// 		defer func() {
// 			close(client)
// 		}()

// 		go handler(broadcast, client)

// 		flusher, _ := w.(http.Flusher)
// 		for {
// 			select {
// 			case message := <-broadcast:
// 				fmt.Fprintf(w, "data: %s\n\n", message)
// 				flusher.Flush()
// 			case message := <-client:
// 				fmt.Fprintf(w, "data: %s\n\n", message)
// 				flusher.Flush()
// 			case <-r.Context().Done():
// 				return
// 			}
// 		}
// 	})
// }