Pings and bigger buffer

Author: jazware

pkg/server/server.go

@@ -112,6 +112,8 @@ func (s *Server) HandleSubscribe(c echo.Context) error {
 	}
 	defer ws.Close()
 
+	writeWait := 10 * time.Second
+
 	log := slog.With("source", "server_handle_subscribe", "socket_addr", ws.RemoteAddr().String(), "real_ip", subIP)
 
 	sub, err := s.AddSubscriber(ws, subIP, subscriberOpts)
@@ -143,6 +145,8 @@ func (s *Server) HandleSubscribe(c echo.Context) error {
 					cancel()
 					return
 				}
+			case websocket.PongMessage:
+				log.Debug("received pong message from client")
 			case websocket.CloseMessage:
 				log.Info("received close message from client")
 				cancel()
@@ -241,6 +245,9 @@ func (s *Server) HandleSubscribe(c echo.Context) error {
 	}
 
 	// Read events from the outbox and send them to the subscriber
+	pingPeriod := 30 * time.Second
+	t := time.NewTicker(pingPeriod)
+	defer t.Stop()
 	for {
 		select {
 		case <-ctx.Done():
@@ -253,6 +260,8 @@ func (s *Server) HandleSubscribe(c echo.Context) error {
 				return fmt.Errorf("failed to wait for rate limiter: %w", err)
 			}
 
+			ws.SetWriteDeadline(time.Now().Add(writeWait))
+
 			// When compression is enabled, the msg is a zstd compressed message
 			if compress {
 				if err := sub.WriteMessage(websocket.BinaryMessage, *msg); err != nil {
@@ -267,6 +276,12 @@ func (s *Server) HandleSubscribe(c echo.Context) error {
 				log.Error("failed to write message to websocket", "error", err)
 				return nil
 			}
+		case <-t.C:
+			ws.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := sub.WriteMessage(websocket.PingMessage, nil); err != nil {
+				log.Error("failed to write ping to websocket", "error", err)
+				return nil
+			}
 		}
 	}
 }

pkg/server/subscriber.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log/slog"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -72,7 +73,7 @@ func emitToSubscriber(ctx context.Context, log *slog.Logger, sub *Subscriber, ti
 
 	if playback {
 		// Copy the event bytes so the playback iterator can reuse the buffer
-		evtBytes = append([]byte{}, evtBytes...)
+		evtBytes = slices.Clone(evtBytes)
 		select {
 		case <-ctx.Done():
 			log.Error("failed to send event to subscriber", "error", ctx.Err(), "subscriber", sub.id)
@@ -213,7 +214,7 @@ func (s *Server) AddSubscriber(ws *websocket.Conn, realIP string, opts *Subscrib
 	sub := Subscriber{
 		ws:                ws,
 		realIP:            realIP,
-		outbox:            make(chan *[]byte, 10_000),
+		outbox:            make(chan *[]byte, 50_000),
 		hello:             make(chan struct{}),
 		id:                s.nextSub,
 		wantedCollections: opts.WantedCollections,