Server-Sent Events Handler

package main

import (
	"bufio"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"time"
)

type Broker struct {
	mu      sync.Mutex
	clients map[chan string]struct{}
}

func NewBroker() *Broker {
	return &Broker{clients: make(map[chan string]struct{})}
}

func (b *Broker) Subscribe() chan string {
	events := make(chan string, 8)
	b.mu.Lock()
	b.clients[events] = struct{}{}
	b.mu.Unlock()
	return events
}

func (b *Broker) Unsubscribe(events chan string) {
	b.mu.Lock()
	delete(b.clients, events)
	b.mu.Unlock()
}

func (b *Broker) Broadcast(event string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for events := range b.clients {
		select {
		case events <- event:
		default:
		}
	}
}

func (b *Broker) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}
	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)
	flusher.Flush()

	events := b.Subscribe()
	defer b.Unsubscribe(events)
	for {
		select {
		case <-r.Context().Done():
			return
		case event := <-events:
			fmt.Fprintf(w, "data: %s\n\n", event)
			flusher.Flush()
		}
	}
}

func main() {
	broker := NewBroker()
	server := httptest.NewServer(broker)
	defer server.Close()

	go func() {
		time.Sleep(20 * time.Millisecond)
		broker.Broadcast("deploy started")
		time.Sleep(10 * time.Millisecond)
		broker.Broadcast("deploy finished")
	}()

	response, err := http.Get(server.URL)
	if err != nil {
		panic(err)
	}
	defer response.Body.Close()
	fmt.Println("content-type:", response.Header.Get("Content-Type"))

	scanner := bufio.NewScanner(response.Body)
	received := 0
	for scanner.Scan() && received < 2 {
		line := scanner.Text()
		if strings.HasPrefix(line, "data: ") {
			fmt.Println("event:", strings.TrimPrefix(line, "data: "))
			received++
		}
	}
}