services/ws/router.go

   1 package main
   2
   3 import (
   4         "log"
   5         "net/http"
   6         "sync"
   7
   8         "golang.org/x/net/websocket"
   9 )
  10
  11 type router struct {
  12         Config *Config
  13
  14         eventSource eventSource
  15         mux         *http.ServeMux
  16         setupOnce   sync.Once
  17 }
  18
  19 func (rtr *router) setup() {
  20         rtr.mux = http.NewServeMux()
  21         rtr.mux.Handle("/websocket", rtr.makeServer(&handlerV0{
  22                 QueueSize: rtr.Config.ClientEventQueue,
  23         }))
  24         rtr.mux.Handle("/arvados/v1/events.ws", rtr.makeServer(&handlerV1{
  25                 QueueSize: rtr.Config.ClientEventQueue,
  26         }))
  27 }
  28
  29 func (rtr *router) makeServer(handler handler) *websocket.Server {
  30         return &websocket.Server{
  31                 Handshake: func(c *websocket.Config, r *http.Request) error {
  32                         return nil
  33                 },
  34                 Handler: websocket.Handler(func(ws *websocket.Conn) {
  35                         log.Printf("socket request: %+v", ws.Request())
  36                         sink := rtr.eventSource.NewSink(nil)
  37                         handler.Handle(ws, sink.Channel())
  38                         sink.Stop()
  39                         ws.Close()
  40                         log.Printf("socket disconnect: %+v", ws.Request().RemoteAddr)
  41                 }),
  42         }
  43 }
  44
  45 func (rtr *router) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
  46         rtr.setupOnce.Do(rtr.setup)
  47         rtr.mux.ServeHTTP(resp, req)
  48 }