Skip to content

Instantly share code, notes, and snippets.

@nf
Created July 14, 2014 01:23
Show Gist options
  • Save nf/7c03729770315c05570f to your computer and use it in GitHub Desktop.
Save nf/7c03729770315c05570f to your computer and use it in GitHub Desktop.
package main
import (
"flag"
"log"
"net/http"
"strings"
"sync"
"time"
"github.com/gorilla/websocket"
)
var httpAddr = flag.String("http", ":8080", "HTTP listen address")
func main() {
flag.Parse()
http.Handle("/", NewServer())
log.Fatal(http.ListenAndServe(*httpAddr, nil))
}
type Server struct {
mu sync.Mutex
m map[string][]chan string
}
func NewServer() *Server {
return &Server{m: make(map[string][]chan string)}
}
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
const (
broadcast = "/broadcast/"
listen = "/listen/"
)
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
p := r.URL.Path
switch {
default:
http.NotFound(w, r)
case strings.HasPrefix(p, broadcast):
p = strings.TrimPrefix(p, broadcast)
s.broadcast(p, r.FormValue("message"))
case strings.HasPrefix(p, listen):
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Print(err)
return
}
p = strings.TrimPrefix(r.URL.Path, listen)
c := s.listen(p)
defer s.hangup(p, c)
for m := range c {
err := conn.WriteMessage(websocket.TextMessage, []byte(m))
if err != nil {
log.Print(err)
return
}
}
}
}
func (s *Server) listen(p string) <-chan string {
c := make(chan string)
s.mu.Lock()
s.m[p] = append(s.m[p], c)
s.mu.Unlock()
return c
}
func (s *Server) hangup(p string, c <-chan string) {
// Remove channel from listener map.
s.mu.Lock()
ls := s.m[p]
for i := range ls {
if ls[i] == c {
ls = append(ls[:i], ls[i+1:]...)
break
}
}
s.m[p] = ls
s.mu.Unlock()
// Drain channel for a minute, to unblock any in-flight senders.
go func() {
timeout := time.After(1 * time.Minute)
for {
select {
case <-c:
case <-timeout:
return
}
}
}()
}
func (s *Server) broadcast(p, m string) {
s.mu.Lock()
ls := append([]chan string{}, s.m[p]...) // copy
s.mu.Unlock()
for _, c := range ls {
c <- m
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment