From 001dd8789b763094e6bf73b02fa00af1fa82371b Mon Sep 17 00:00:00 2001 From: "Enot (ded) Skelly" Date: Fri, 29 May 2026 09:40:26 -0700 Subject: [PATCH] finish ws sub/unsub --- README.md | 13 ++++++++- internal/hub/hub.go | 64 ++++++++++++++++++++++++++++++------------ internal/ws/handler.go | 36 ++++++++++++++++++++---- 3 files changed, 89 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index 24dee97..37f6a7c 100644 --- a/README.md +++ b/README.md @@ -205,6 +205,17 @@ On connect the server sends a hello message: } ``` +- **Unsubscribe** + +```json +{ + "v": 1, + "type": "unsubscribe", + "id": "unsub-1", + "subscriptionId": "" +} +``` + - **Ping** (send every 30s; connection closes after 90s idle) ```json @@ -369,13 +380,13 @@ For paginated responses use the generic page wrapper: - [x] REST API: Stats - [x] Materialized view refresh (mv_hourly_iata_stats, mv_top_nodes_by_iata) - [x] Swagger/OpenAPI documentation via swaggo/swag +- [x] WebSocket subscription unsubscribe (scaffolded, not implemented) ### In progress / next - [ ] Path resolution (node short ID lookup) - [ ] Propagation time calculation - [ ] Routes and traces endpoints -- [ ] WebSocket subscription unsubscribe (scaffolded, not implemented) ### Future diff --git a/internal/hub/hub.go b/internal/hub/hub.go index 9a424c8..1faf849 100644 --- a/internal/hub/hub.go +++ b/internal/hub/hub.go @@ -52,19 +52,24 @@ type Scope struct { Events []EventType } +type LaggedNotification struct { + DroppedCount int +} + // Client represents a connected WebSocket consumer. type Client struct { - Send chan Event - scope []Scope // OR semantics: event matches if it matches any scope entry + Send chan Event + laggedCH chan LaggedNotification + subscriptions map[string]Scope // OR semantics: event matches if it matches any scope entry } // matches returns true if the event satisfies at least one of the client's // active subscriptions. func (c *Client) matches(e Event) bool { - if len(c.scope) == 0 { - return true // no subscriptions yet → receive nothing + if len(c.subscriptions) == 0 { + return false } - for _, s := range c.scope { + for _, s := range c.subscriptions { if scopeMatches(s, e) { return true } @@ -72,6 +77,12 @@ func (c *Client) matches(e Event) bool { return false } +// LaggedCH returns the channel on which lagged notifications are delivered. +// The WS write pump selects on this alongside Send. +func (c *Client) LaggedCH() <-chan LaggedNotification { + return c.laggedCH +} + func scopeMatches(s Scope, e Event) bool { if len(s.Events) > 0 && !containsEventType(s.Events, e.Type) { return false @@ -100,9 +111,9 @@ type Hub struct { } type subscribeMsg struct { - client *Client - scope Scope - hasScope bool // true when this is an AddScope call, false for NewClient registration + client *Client + scope Scope + subscriptionID string } type unsubscribeMsg struct { @@ -124,18 +135,27 @@ func New() *Hub { // The caller is responsible for calling Remove when the connection closes. func (h *Hub) NewClient() *Client { c := &Client{ - Send: make(chan Event, 256), + Send: make(chan Event, 256), + laggedCH: make(chan LaggedNotification, 8), + subscriptions: make(map[string]Scope), } // We don't add it to the map here; we send it through the channel so // Run() is the only goroutine that touches the client map. - h.subscribe <- subscribeMsg{client: c, hasScope: false} + h.subscribe <- subscribeMsg{client: c} return c } // AddScope appends a subscription scope to a client. Called by the WS handler // when it receives a "subscribe" message from the client. -func (h *Hub) AddScope(c *Client, s Scope) { - h.subscribe <- subscribeMsg{client: c, scope: s, hasScope: true} +func (h *Hub) AddScope(c *Client, id string, s Scope) { + h.subscribe <- subscribeMsg{client: c, scope: s, subscriptionID: id} +} + +// RemoveScope removes a single subscription by ID. Called by the WS handler +// when it receives an "unsubscribe" message from the client. Silently ignored +// if the ID is not found. +func (h *Hub) RemoveScope(c *Client, id string) { + h.unsubscribe <- unsubscribeMsg{client: c, subscriptionID: id} } // Remove deregisters a client and closes its Send channel. @@ -168,20 +188,26 @@ func (h *Hub) Run() { select { case msg := <-h.subscribe: - if !msg.hasScope { + if msg.subscriptionID == "" { // Registration with no scope yet (NewClient path). clients[msg.client] = struct{}{} } else { // AddScope path — client must already be registered. if _, ok := clients[msg.client]; ok { - msg.client.scope = append(msg.client.scope, msg.scope) + msg.client.subscriptions[msg.subscriptionID] = msg.scope } } + case msg := <-h.unsubscribe: + if _, ok := clients[msg.client]; ok { + delete(msg.client.subscriptions, msg.subscriptionID) + } + case c := <-h.remove: if _, ok := clients[c]; ok { delete(clients, c) close(c.Send) + close(c.laggedCH) } case evt := <-h.broadcast: @@ -192,14 +218,16 @@ func (h *Hub) Run() { select { case c.Send <- evt: default: - // Client send buffer full. The WS write pump is responsible - // for detecting its own lagged state and sending the lagged - // message. We just drain one slot so the broadcast loop - // doesn't block. + dropped := 1 select { case <-c.Send: default: } + select { + case c.laggedCH <- LaggedNotification{DroppedCount: dropped}: + default: + // laggedCh itself full; write pump will catch up on next drain + } log.Printf("hub: client send buffer full, dropped event type=%s", evt.Type) } } diff --git a/internal/ws/handler.go b/internal/ws/handler.go index 78a0291..c344825 100644 --- a/internal/ws/handler.go +++ b/internal/ws/handler.go @@ -74,7 +74,7 @@ func Handler(h *hub.Hub) http.HandlerFunc { select { case evt, ok := <-client.Send: if !ok { - return // hub closed the channel (client removed) + return } msg := map[string]any{ "v": 1, @@ -89,6 +89,24 @@ func Handler(h *hub.Hub) http.HandlerFunc { cancel() return } + + case lag, ok := <-client.LaggedCH(): + if !ok { + return + } + lagged := map[string]any{ + "v": 1, + "type": "lagged", + "droppedCount": lag.DroppedCount, + "since": time.Now().UnixMilli(), + } + lagBytes, _ := json.Marshal(lagged) + if err := conn.Write(ctx, websocket.MessageText, lagBytes); err != nil { + log.Printf("ws[%s]: failed to write lagged notice: %v", connID, err) + cancel() + return + } + case <-ctx.Done(): return } @@ -150,8 +168,8 @@ func handleClientMessage(ctx context.Context, client *hub.Client, h *hub.Hub, co Events: msg.Scope.Events, // RegionIATAs: TODO expand msg.Scope.RegionIDs → IATA slice via region_iatas DB lookup } - h.AddScope(client, scope) subID := uuid.NewString() + h.AddScope(client, subID, scope) reply, _ := json.Marshal(map[string]any{ "v": 1, "type": "subscribed", "id": msg.ID, "subscriptionId": subID, }) @@ -162,9 +180,17 @@ func handleClientMessage(ctx context.Context, client *hub.Client, h *hub.Hub, co } case "unsubscribe": - // TODO: remove the specific subscriptionId from client.scope. - // For now scope entries are append-only; implement removal when needed. - log.Printf("ws[%s]: unsubscribe %s (TODO)", connID, msg.SubscriptionID) + if msg.SubscriptionID == "" { + return + } + h.RemoveScope(client, msg.SubscriptionID) + reply, _ := json.Marshal(map[string]any{ + "v": 1, "type": "unsubscribed", "id": msg.ID, "subscriptionId": msg.SubscriptionID, + }) + log.Printf("ws[%s]: unsubscribed %s", connID, msg.SubscriptionID) + if err := conn.Write(ctx, websocket.MessageText, reply); err != nil { + log.Printf("ws[%s]: failed to send unsubscribed reply: %v", connID, err) + } case "ping": reply, _ := json.Marshal(map[string]any{"v": 1, "type": "pong", "id": msg.ID})