package main import ( "fmt" "log" "net/http" "os" "sync" ) func main() { host := os.Getenv("HOST") if host == "" { host = "0.0.0.0" } port := os.Getenv("PORT") if port == "" { port = "8080" } address := host + ":" + port chat := NewChat() router := http.NewServeMux() router.HandleFunc("/", handleIndex) router.HandleFunc("/chat", handleChat) router.HandleFunc("/chat/events", newHandleEvents(chat)) router.HandleFunc("/chat/msg", newHandleMsg(chat)) log.Println("Starting server on", address) http.ListenAndServe(address, router) } func handleIndex(w http.ResponseWriter, r *http.Request) { fmt.Fprintln(w, ` Chat
`) } func handleChat(w http.ResponseWriter, r *http.Request) { fmt.Fprintln(w, `

Chat

`) } func newHandleEvents(chat *Chat) http.HandlerFunc { return 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.WriteHeader(http.StatusOK) msgs, unsubscribe := chat.Register() defer unsubscribe() for msg := range msgs { fmt.Fprintf(w, "data:
%v
\n\n", msg) if f, ok := w.(http.Flusher); ok { f.Flush() } } } } const msgTemplate = `
` func newHandleMsg(chat *Chat) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { msg := r.FormValue("msg") if msg != "" { chat.Broadcast(msg) } fmt.Fprintln(w, msgTemplate) } } type Chat struct { mu sync.Mutex clients map[int]chan<- string counter int } func NewChat() *Chat { return &Chat{ clients: make(map[int]chan<- string), } } func (c *Chat) Register() (<-chan string, func()) { c.mu.Lock() defer c.mu.Unlock() ch := make(chan string, 1) key := c.counter c.counter += 1 c.clients[key] = ch return ch, func() { c.mu.Lock() defer c.mu.Unlock() delete(c.clients, key) } } func (c *Chat) Broadcast(msg string) { c.mu.Lock() defer c.mu.Unlock() var toDelete []int for key, ch := range c.clients { select { case ch <- msg: default: close(ch) toDelete = append(toDelete, key) } } for _, key := range toDelete { delete(c.clients, key) } }