-
Notifications
You must be signed in to change notification settings - Fork 0
/
server.go
86 lines (74 loc) · 1.67 KB
/
server.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
package main
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"github.com/go-chi/chi"
"golang.org/x/sync/errgroup"
)
// eventsEndpoint returns an http.HandlerFunc that processes an http.Request
// to server sent event.
func eventsEndpoint(events <-chan Event) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
h := w.Header()
h.Set("Content-Type", "text/event-stream")
h.Set("Cache-Control", "no-cache")
h.Set("Connection", "keep-alive")
h.Set("X-Accel-Buffering", "no")
flush := func() {}
if f, ok := w.(http.Flusher); ok {
flush = f.Flush
}
defer func() {
fmt.Print("events: stream closed\n")
io.WriteString(w, "event:fin\n\n")
flush()
}()
for {
select {
case <-ctx.Done():
fmt.Print("events: stream cancelled\n")
return
case e, more := <-events:
if !more {
return
}
eventJSON, err := json.Marshal(e)
if err == nil {
io.WriteString(w, "data:"+string(eventJSON)+"\n\n")
} else {
io.WriteString(w, "error:"+err.Error()+"\n\n")
}
flush()
}
}
}
}
func startServer(ctx context.Context, addr, html string, events <-chan Event) {
r := chi.NewRouter()
r.Get("/", func(w http.ResponseWriter, r *http.Request) {
h := w.Header()
h.Set("Content-Type", "text/html")
fmt.Fprint(w, html)
})
r.Get("/events", eventsEndpoint(events))
srv := &http.Server{
Addr: addr,
// ReadTimeout: 75 * time.Second,
Handler: r,
}
var g errgroup.Group
g.Go(func() error {
return srv.ListenAndServe()
})
g.Go(func() error {
<-ctx.Done()
return srv.Shutdown(ctx)
})
if err := g.Wait(); err != nil {
fmt.Printf("shutdown server: %s", err)
}
}