diff --git a/main.go b/main.go index b9da54a..d8a5a0d 100644 --- a/main.go +++ b/main.go @@ -6,6 +6,7 @@ import ( "log" "net/http" "os" + "sync" "github.com/gorilla/websocket" ) @@ -18,8 +19,19 @@ var upgrader = websocket.Upgrader{ }, } +type client struct { + nudges chan [2]int + quit chan bool +} + +type clientsStore struct { + set map[*client]struct{} + sync.RWMutex +} + func main() { - var nudges = make(chan [2]int) + var clients clientsStore + clients.set = make(map[*client]struct{}) go func() { mux := http.NewServeMux() @@ -45,7 +57,17 @@ func main() { w.WriteHeader(http.StatusBadRequest) return } - nudges <- nudge + (func() { + clients.RLock() + defer clients.RUnlock() + for c, _ := range clients.set { + select { + case c.nudges <- nudge: + default: + log.Println("Client too slow -> dropping nudge") + } + } + })() w.WriteHeader(http.StatusNoContent) case "GET": w.Header().Add("Content-Type", "text/html") @@ -82,22 +104,29 @@ func main() { log.Println(err) return } - var quit = make(chan bool) + defer con.Close() + var c client + c.nudges = make(chan [2]int) + (func() { + clients.Lock() + defer clients.Unlock() + clients.set[&c] = struct{}{} + })() + defer (func() { + clients.Lock() + defer clients.Unlock() + delete(clients.set, &c) + log.Println("Closing client connection") + })() con.SetCloseHandler(func(code int, text string) error { - quit <- true + close(c.nudges) return nil }) - for { - select { - case <-quit: + for n := range c.nudges { + err := con.WriteJSON(n) + if err != nil { + log.Println(err) return - case n := <-nudges: - log.Println(n) - err := con.WriteJSON(n) - if err != nil { - log.Println(err) - return - } } } })