You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

103 lines
2.0 KiB
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
// }
// }
// })
// }