wip: add upstream heartbeats

Author: endigma

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client.go

@@ -114,6 +114,8 @@ func NewGraphQLSubscriptionClient(ctx context.Context, opts ...SubscriptionClien
 			UpgradeClient:   cfg.UpgradeClient,
 			StreamingClient: cfg.StreamingClient,
 			Logger:          cfg.Logger,
+			PingInterval:    cfg.PingInterval,
+			PingTimeout:     cfg.PingTimeout,
 		}),
 	}
 }

v2/pkg/engine/datasource/graphql_datasource/subscriptionclient/client/client.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"net/http"
+	"time"
 
 	"github.com/jensneuse/abstractlogger"
 
@@ -35,6 +36,8 @@ type Config struct {
 	UpgradeClient   *http.Client
 	StreamingClient *http.Client
 	Logger          abstractlogger.Logger
+	PingInterval    time.Duration
+	PingTimeout     time.Duration
 }
 
 // New creates a new subscription client with the provided config.
@@ -53,7 +56,7 @@ func New(ctx context.Context, cfg Config) *Client {
 		ctx: ctx,
 		log: cfg.Logger,
 
-		ws:  transport.NewWSTransport(ctx, cfg.UpgradeClient, cfg.Logger),
+		ws:  transport.NewWSTransport(ctx, cfg.UpgradeClient, cfg.Logger, cfg.PingInterval, cfg.PingTimeout),
 		sse: transport.NewSSETransport(ctx, cfg.StreamingClient, cfg.Logger),
 	}
 

v2/pkg/engine/datasource/graphql_datasource/subscriptionclient/transport/ws_conn.go

@@ -27,8 +27,6 @@ type WSConnection struct {
 	protocol protocol.Protocol
 	log      abstractlogger.Logger
 
-	writeMu sync.Mutex
-
 	subsMu sync.RWMutex
 	subs   map[string]chan<- *common.Message
 
@@ -39,14 +37,19 @@ type WSConnection struct {
 	onEmpty func()
 
 	WriteTimeout time.Duration
+
+	// Ping/pong tracking for client-initiated heartbeats.
+	// Values stored as UnixNano timestamps.
+	lastPingSentAt atomic.Int64
+	lastPongAt     atomic.Int64
 }
 
 func NewWSConnection(ctx context.Context, conn *websocket.Conn, protocol protocol.Protocol, log abstractlogger.Logger, onEmpty func()) *WSConnection {
 	if log == nil {
 		log = abstractlogger.NoopLogger
 	}
 
-	return &WSConnection{
+	c := &WSConnection{
 		ctx:      ctx,
 		conn:     conn,
 		protocol: protocol,
@@ -57,6 +60,10 @@ func NewWSConnection(ctx context.Context, conn *websocket.Conn, protocol protoco
 
 		WriteTimeout: DefaultWriteTimeout,
 	}
+
+	c.lastPongAt.Store(time.Now().UnixNano())
+
+	return c
 }
 
 func (c *WSConnection) Subscribe(ctx context.Context, id string, req *common.Request) (<-chan *common.Message, func(), error) {
@@ -77,9 +84,7 @@ func (c *WSConnection) Subscribe(ctx context.Context, id string, req *common.Req
 	c.subs[id] = ch
 	c.subsMu.Unlock()
 
-	if err := c.withWriteLock(func() error {
-		return c.protocol.Subscribe(ctx, c.conn, id, req)
-	}); err != nil {
+	if err := c.protocol.Subscribe(ctx, c.conn, id, req); err != nil {
 		c.log.Error("wsConnection.Subscribe",
 			abstractlogger.String("id", id),
 			abstractlogger.Error(err),
@@ -110,9 +115,6 @@ func (c *WSConnection) removeSub(id string) {
 	}
 
 	if isEmpty {
-		if c.onEmpty != nil {
-			c.onEmpty()
-		}
 		c.Close()
 	}
 }
@@ -131,24 +133,11 @@ func (c *WSConnection) unsubscribe(id string) {
 	unsubscribeCtx, cancel := context.WithTimeout(context.Background(), c.WriteTimeout)
 	defer cancel()
 
-	_ = c.withWriteLock(func() error {
-		return c.protocol.Unsubscribe(unsubscribeCtx, c.conn, id)
-	})
+	_ = c.protocol.Unsubscribe(unsubscribeCtx, c.conn, id)
 
 	c.removeSub(id)
 }
 
-func (c *WSConnection) withWriteLock(f func() error) error {
-	c.writeMu.Lock()
-	defer c.writeMu.Unlock()
-
-	if c.closed.Load() {
-		return common.ErrConnectionClosed
-	}
-
-	return f()
-}
-
 func (c *WSConnection) ReadLoop() {
 	defer c.shutdown(errors.New("read loop exited"))
 
@@ -171,12 +160,10 @@ func (c *WSConnection) ReadLoop() {
 		case protocol.MessagePing:
 			c.log.Debug("wsConnection.ReadLoop", abstractlogger.String("message", "ping"))
 			pongCtx, cancel := context.WithTimeout(c.ctx, c.WriteTimeout)
-			_ = c.withWriteLock(func() error {
-				return c.protocol.Pong(pongCtx, c.conn)
-			})
+			_ = c.protocol.Pong(pongCtx, c.conn)
 			cancel()
 		case protocol.MessagePong:
-			// Do nothing, pongs can sometimes be used as unidirectional heartbeats
+			c.lastPongAt.Store(time.Now().UnixNano())
 			c.log.Debug("wsConnection.ReadLoop", abstractlogger.String("message", "pong"))
 		case protocol.MessageData, protocol.MessageError, protocol.MessageComplete:
 			c.dispatch(msg)
@@ -229,6 +216,10 @@ func (c *WSConnection) shutdown(err error) {
 		}
 		close(ch)
 	}
+
+	if c.onEmpty != nil {
+		c.onEmpty()
+	}
 }
 
 func (c *WSConnection) Close() error {
@@ -250,6 +241,30 @@ func (c *WSConnection) SubCount() int {
 	return len(c.subs)
 }
 
+// SendPing sends a protocol-level ping message and records the timestamp.
+func (c *WSConnection) SendPing(timeout time.Duration) error {
+	pingCtx, cancel := context.WithTimeout(c.ctx, timeout)
+	defer cancel()
+
+	err := c.protocol.Ping(pingCtx, c.conn)
+	if err != nil {
+		return err
+	}
+
+	c.lastPingSentAt.Store(time.Now().UnixNano())
+	return nil
+}
+
+// PongOverdue returns true if a pong has not been received since the last ping
+// and the ping timeout has elapsed.
+func (c *WSConnection) PongOverdue(timeout time.Duration) bool {
+	pingSent := c.lastPingSentAt.Load()
+	if pingSent == 0 {
+		return false
+	}
+	return c.lastPongAt.Load() < pingSent && time.Since(time.Unix(0, pingSent)) > timeout
+}
+
 func (c *WSConnection) IsClosed() bool {
 	return c.closed.Load()
 }

v2/pkg/engine/datasource/graphql_datasource/subscriptionclient/transport/ws_conn_test.go

@@ -297,6 +297,52 @@ func TestWSConnection_OnEmpty(t *testing.T) {
 			t.Error("onEmpty should be called after last subscription removed")
 		}
 	})
+
+	t.Run("calls callback on direct Close", func(t *testing.T) {
+		t.Parallel()
+
+		conn, _ := newTestConn(t)
+		proto := newMockProtocol()
+
+		emptyCalled := make(chan struct{}, 1)
+		wsc := transport.NewWSConnection(t.Context(), conn, proto, nil, func() {
+			emptyCalled <- struct{}{}
+		})
+
+		wsc.Close()
+
+		select {
+		case <-emptyCalled:
+			// success
+		case <-time.After(100 * time.Millisecond):
+			t.Error("onEmpty callback not called on Close")
+		}
+	})
+
+	t.Run("calls callback on read loop exit", func(t *testing.T) {
+		t.Parallel()
+
+		conn, _ := newTestConn(t)
+		proto := newMockProtocol()
+
+		emptyCalled := make(chan struct{}, 1)
+		ctx, cancel := context.WithCancel(context.Background())
+		wsc := transport.NewWSConnection(ctx, conn, proto, nil, func() {
+			emptyCalled <- struct{}{}
+		})
+
+		go wsc.ReadLoop()
+
+		// Cancel context to cause the read loop to exit
+		cancel()
+
+		select {
+		case <-emptyCalled:
+			// success
+		case <-time.After(time.Second):
+			t.Error("onEmpty callback not called on read loop exit")
+		}
+	})
 }
 
 func TestWSConnection_Close(t *testing.T) {

v2/pkg/engine/datasource/graphql_datasource/subscriptionclient/transport/ws_transport.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"sync"
+	"time"
 
 	"github.com/coder/websocket"
 	"github.com/jensneuse/abstractlogger"
@@ -36,6 +37,9 @@ type WSTransport struct {
 	upgradeClient *http.Client
 	log           abstractlogger.Logger
 
+	pingInterval time.Duration
+	pingTimeout  time.Duration
+
 	mu      sync.Mutex
 	dialing map[uint64]*dialResult
 	conns   map[uint64]*WSConnection
@@ -50,7 +54,11 @@ type dialResult struct {
 // NewWSTransport creates a new WSTransport with the provided http.Client
 // for WebSocket upgrade requests. The transport will automatically close
 // all connections when ctx is cancelled.
-func NewWSTransport(ctx context.Context, upgradeClient *http.Client, log abstractlogger.Logger) *WSTransport {
+//
+// If pingInterval > 0, a single goroutine sends protocol-level pings to all
+// connections at that cadence. If pingTimeout > 0, connections that fail to
+// respond with a pong within that window after a ping are shut down.
+func NewWSTransport(ctx context.Context, upgradeClient *http.Client, log abstractlogger.Logger, pingInterval, pingTimeout time.Duration) *WSTransport {
 	if log == nil {
 		log = abstractlogger.NoopLogger
 	}
@@ -59,12 +67,18 @@ func NewWSTransport(ctx context.Context, upgradeClient *http.Client, log abstrac
 		ctx:           ctx,
 		upgradeClient: upgradeClient,
 		log:           log,
+		pingInterval:  pingInterval,
+		pingTimeout:   pingTimeout,
 		conns:         make(map[uint64]*WSConnection),
 		dialing:       make(map[uint64]*dialResult),
 	}
 
 	context.AfterFunc(ctx, t.closeAll)
 
+	if pingInterval > 0 {
+		go t.pingLoop()
+	}
+
 	return t
 }
 
@@ -81,9 +95,6 @@ func (t *WSTransport) Subscribe(ctx context.Context, req *common.Request, opts c
 // closeAll closes all connections. Called automatically when context is cancelled.
 func (t *WSTransport) closeAll() {
 	t.mu.Lock()
-
-	// Copy because conn.Close -> shutdown -> onEmpty -> t.removeConn -> t.mu.Lock
-	// would cause a deadlock
 	conns := make([]*WSConnection, 0, len(t.conns))
 	for _, conn := range t.conns {
 		conns = append(conns, conn)
@@ -102,6 +113,46 @@ func (t *WSTransport) closeAll() {
 	}
 }
 
+// pingLoop sends periodic pings to all active connections and shuts down
+// any that have not responded with a pong in time.
+func (t *WSTransport) pingLoop() {
+	tick := time.Tick(t.pingInterval)
+	for {
+		select {
+		case <-t.ctx.Done():
+			return
+		case <-tick:
+			t.mu.Lock()
+			conns := make([]*WSConnection, 0, len(t.conns))
+			for _, conn := range t.conns {
+				conns = append(conns, conn)
+			}
+			t.mu.Unlock()
+
+			for _, conn := range conns {
+				if conn.IsClosed() {
+					continue
+				}
+
+				if t.pingTimeout > 0 && conn.PongOverdue(t.pingTimeout) {
+					t.log.Debug("wsTransport.pingLoop",
+						abstractlogger.String("action", "pong_timeout"),
+					)
+					conn.Close()
+					continue
+				}
+
+				if err := conn.SendPing(DefaultWriteTimeout); err != nil {
+					t.log.Debug("wsTransport.pingLoop",
+						abstractlogger.String("action", "ping_failed"),
+						abstractlogger.Error(err),
+					)
+				}
+			}
+		}
+	}
+}
+
 func (t *WSTransport) ConnCount() int {
 	t.mu.Lock()
 	defer t.mu.Unlock()