websocket.go 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. package handlers
  2. import (
  3. "log"
  4. "net/http"
  5. "sync"
  6. "github.com/gorilla/websocket"
  7. )
  8. var upgrader = websocket.Upgrader{
  9. CheckOrigin: func(r *http.Request) bool { return true },
  10. }
  11. type VisitorDot struct {
  12. UserID string `json:"user_id,omitempty"`
  13. Lat float64 `json:"lat"`
  14. Lng float64 `json:"lng"`
  15. }
  16. type WSHub struct {
  17. mu sync.RWMutex
  18. clients map[*websocket.Conn]VisitorDot
  19. }
  20. func NewWSHub() *WSHub {
  21. return &WSHub{clients: make(map[*websocket.Conn]VisitorDot)}
  22. }
  23. func (h *WSHub) HandleWS(w http.ResponseWriter, r *http.Request) {
  24. conn, err := upgrader.Upgrade(w, r, nil)
  25. if err != nil {
  26. log.Printf("ws upgrade: %v", err)
  27. return
  28. }
  29. dot := VisitorDot{}
  30. h.mu.Lock()
  31. h.clients[conn] = dot
  32. h.mu.Unlock()
  33. defer func() {
  34. h.mu.Lock()
  35. delete(h.clients, conn)
  36. h.mu.Unlock()
  37. conn.Close()
  38. }()
  39. for {
  40. var msg VisitorDot
  41. if err := conn.ReadJSON(&msg); err != nil {
  42. break
  43. }
  44. h.mu.Lock()
  45. h.clients[conn] = msg
  46. visitors := make([]VisitorDot, 0, len(h.clients))
  47. for _, v := range h.clients {
  48. visitors = append(visitors, v)
  49. }
  50. h.mu.Unlock()
  51. if err := conn.WriteJSON(map[string]interface{}{"visitors": visitors}); err != nil {
  52. break
  53. }
  54. }
  55. }