Fix expired peer cleanup and disconnect handling

internal/signaling/handler.go

@@ -15,10 +15,11 @@ import (
 	"github.com/poki/netlib/internal/metrics"
 	"github.com/poki/netlib/internal/signaling/stores"
 	"github.com/poki/netlib/internal/util"
+	"github.com/rs/xid"
 	"go.uber.org/zap"
 )
 
-const LobbyCleanInterval = 30 * time.Minute
+const LobbyCleanInterval = 5 * time.Minute
 const LobbyCleanThreshold = 24 * time.Hour
 const peerPingDuration = 2 * time.Second
 const peerActiveUpdateInterval = 30 * time.Second
@@ -95,15 +96,17 @@ func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.Cre
 		}
 
 		peer := &Peer{
-			store: store,
-			conn:  conn,
+			connectionID: xid.New().String(),
+			store:        store,
+			conn:         conn,
 
 			retrievedIDCallback: manager.Reconnected,
 
 			Country: country,
 			Region:  region,
 		}
 		defer func() {
+			cancel()
 			logger.Debug("peer websocket closed", zap.String("peer", peer.ID), zap.String("game", peer.Game), zap.String("origin", r.Header.Get("Origin")))
 			conn.Close(websocket.StatusInternalError, "unexpected closure") // nolint:errcheck
 
@@ -122,21 +125,27 @@ func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.Cre
 			for {
 				select {
 				case <-ticker.C:
-					if err := peer.Send(ctx, PingPacket{Type: "ping"}); err != nil {
+					// Limit ping writes so a stalled connection cannot keep this goroutine alive indefinitely.
+					pingCtx, pingCancel := context.WithTimeout(ctx, 10*time.Second)
+					err := peer.Send(pingCtx, PingPacket{Type: "ping"})
+					pingCancel()
+					if err != nil {
 						if !util.ShouldIgnoreNetworkError(err) {
 							if strings.Contains(err.Error(), "write: broken pipe") || strings.Contains(err.Error(), "connection reset by peer") {
 								logger.Warn("failed to send ping packet", zap.String("peer", peer.ID), zap.Error(err))
 							} else {
 								logger.Error("failed to send ping packet", zap.String("peer", peer.ID), zap.Error(err))
 							}
 						}
+						cancel()
+						return
 					} else {
 						// If we can send a ping packet, and the peer has an ID, we update the peer as being active.
 						// If the peer doesn't have an ID yet, it's still in the process of connecting, so we don't update it.
 						if peer.ID != "" {
 							now := util.NowUTC(ctx)
 							if lastActiveUpdate.IsZero() || now.Sub(lastActiveUpdate) >= peerActiveUpdateInterval {
-								manager.MarkPeerAsActive(ctx, peer.ID)
+								manager.MarkPeerAsActive(ctx, peer.ID, peer.connectionID)
 								lastActiveUpdate = now
 							}
 						}

internal/signaling/peer.go

@@ -24,8 +24,9 @@ type Peer struct {
 	conn  *websocket.Conn
 
 	closedPacketReceived bool
+	connectionID         string
 
-	retrievedIDCallback func(context.Context, string, string, string) (bool, []string, error)
+	retrievedIDCallback func(context.Context, string, string, string, string) (bool, []string, error)
 
 	ID      string
 	Secret  string
@@ -197,7 +198,7 @@ func (p *Peer) HandleHelloPacket(ctx context.Context, packet HelloPacket) error
 	if packet.ID != "" && packet.Secret != "" {
 		logger.Debug("peer reconnecting", zap.String("game", packet.Game), zap.String("peer", packet.ID), zap.String("version", packet.Version))
 		var err error
-		hasReconnected, reconnectingLobbies, err = p.retrievedIDCallback(ctx, packet.ID, packet.Secret, packet.Game)
+		hasReconnected, reconnectingLobbies, err = p.retrievedIDCallback(ctx, packet.ID, packet.Secret, packet.Game, p.connectionID)
 		if err != nil {
 			return fmt.Errorf("unable to reconnect: %w", err)
 		}
@@ -223,7 +224,7 @@ func (p *Peer) HandleHelloPacket(ctx context.Context, packet HelloPacket) error
 		p.Secret = util.GenerateSecret(ctx)
 		logger.Debug("peer connecting", zap.String("game", p.Game), zap.String("peer", p.ID), zap.String("version", packet.Version))
 
-		if err := p.store.CreatePeer(ctx, p.ID, p.Secret, p.Game); err != nil {
+		if err := p.store.CreatePeer(ctx, p.ID, p.Secret, p.Game, p.connectionID); err != nil {
 			return fmt.Errorf("unable to create peer: %w", err)
 		}
 	}
@@ -288,7 +289,16 @@ func (p *Peer) HandleClosePacket(ctx context.Context, packet ClosePacket) error
 	logger := logging.GetLogger(ctx)
 	go metrics.Record(ctx, "client", "close", p.Game, p.ID, p.Lobby)
 
-	p.closedPacketReceived = true
+	if p.ID != "" {
+		disconnected, err := p.store.MarkPeerAsDisconnected(ctx, p.ID, p.connectionID)
+		if err != nil {
+			return fmt.Errorf("unable to mark peer as disconnected: %w", err)
+		}
+		if !disconnected {
+			p.closedPacketReceived = true
+			return nil
+		}
+	}
 
 	logger.Debug("client closed",
 		zap.String("game", p.Game),
@@ -322,6 +332,8 @@ func (p *Peer) HandleClosePacket(ctx context.Context, packet ClosePacket) error
 		p.Lobby = ""
 	}
 
+	// Keep deferred disconnect handling enabled if any cleanup above fails.
+	p.closedPacketReceived = true
 	return nil
 }
 

internal/signaling/stores/postgres.go

@@ -405,7 +405,7 @@ func (s *PostgresStore) ListLobbies(ctx context.Context, game string, country, r
 	return lobbies, nil
 }
 
-func (s *PostgresStore) CreatePeer(ctx context.Context, peerID, secret, gameID string) error {
+func (s *PostgresStore) CreatePeer(ctx context.Context, peerID, secret, gameID, connectionID string) error {
 	if len(peerID) > 20 {
 		logger := logging.GetLogger(ctx)
 		logger.Warn("peer id too long", zap.String("peerID", peerID))
@@ -414,9 +414,9 @@ func (s *PostgresStore) CreatePeer(ctx context.Context, peerID, secret, gameID s
 
 	now := util.NowUTC(ctx)
 	_, err := s.DB.Exec(ctx, `
-		INSERT INTO peers (peer, secret, game, last_seen, updated_at)
-		VALUES ($1, $2, $3, $4, $4)
-	`, peerID, secret, gameID, now)
+		INSERT INTO peers (peer, secret, game, last_seen, updated_at, connection_id)
+		VALUES ($1, $2, $3, $4, $4, $5)
+	`, peerID, secret, gameID, now, connectionID)
 	if err != nil {
 		return err
 	}
@@ -440,46 +440,49 @@ func (s *PostgresStore) UpdatePeerGeo(ctx context.Context, peerID string, countr
 	return err
 }
 
-func (s *PostgresStore) MarkPeerAsActive(ctx context.Context, peerID string) error {
+func (s *PostgresStore) MarkPeerAsActive(ctx context.Context, peerID, connectionID string) error {
+	// Only an authenticated reconnect may reactivate a disconnected peer.
+	// Late heartbeats must not revive it or postpone timeout cleanup.
 	now := util.NowUTC(ctx)
 
 	_, err := s.DB.Exec(ctx, `
 		UPDATE peers
 		SET
-			disconnected = FALSE,
 			last_seen = $1,
 			updated_at = $1
-		WHERE peer = $2
-	`, now, peerID)
+		WHERE peer = $2 AND connection_id = $3 AND disconnected = FALSE
+	`, now, peerID, connectionID)
 	return err
 }
 
-func (s *PostgresStore) MarkPeerAsDisconnected(ctx context.Context, peerID string) error {
+// MarkPeerAsDisconnected ignores handlers superseded by an authenticated reconnect.
+func (s *PostgresStore) MarkPeerAsDisconnected(ctx context.Context, peerID, connectionID string) (bool, error) {
 	now := util.NowUTC(ctx)
-
-	_, err := s.DB.Exec(ctx, `
+	result, err := s.DB.Exec(ctx, `
 		UPDATE peers
-		SET
-			disconnected = TRUE,
-			updated_at = $1
-		WHERE peer = $2
-	`, now, peerID)
-	return err
+		SET disconnected = TRUE, updated_at = $1
+		WHERE peer = $2 AND connection_id = $3
+	`, now, peerID, connectionID)
+	if err != nil {
+		return false, err
+	}
+	return result.RowsAffected() > 0, nil
 }
 
-func (s *PostgresStore) MarkPeerAsReconnected(ctx context.Context, peerID, secret, gameID string) (bool, []string, error) {
+func (s *PostgresStore) MarkPeerAsReconnected(ctx context.Context, peerID, secret, gameID, connectionID string) (bool, []string, error) {
 	now := util.NowUTC(ctx)
 
 	result, err := s.DB.Exec(ctx, `
 		UPDATE peers
 		SET
 			disconnected = FALSE,
+			connection_id = $5,
 			last_seen = $1,
 			updated_at = $1
 		WHERE peer = $2
 		AND secret = $3
 		AND game = $4
-	`, now, peerID, secret, gameID)
+	`, now, peerID, secret, gameID, connectionID)
 	if err != nil {
 		return false, nil, err
 	}
@@ -516,86 +519,93 @@ func (s *PostgresStore) MarkPeerAsReconnected(ctx context.Context, peerID, secre
 	return true, lobbies, nil
 }
 
-func (s *PostgresStore) ClaimNextTimedOutPeer(ctx context.Context, threshold time.Duration) (string, bool, map[string][]string, error) {
+func (s *PostgresStore) ClaimTimedOutPeers(ctx context.Context, threshold time.Duration, limit int) ([]TimedOutPeer, error) {
+	if limit <= 0 {
+		return nil, fmt.Errorf("cleanup limit must be positive")
+	}
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
 	now := util.NowUTC(ctx)
-
 	tx, err := s.DB.Begin(ctx)
 	if err != nil {
-		return "", false, nil, err
+		return nil, err
 	}
 	defer tx.Rollback(context.Background()) //nolint:errcheck
 
-	var peerID string
-	var disconnected bool
-	err = tx.QueryRow(ctx, `
-		WITH d AS (
-			SELECT peer, disconnected
-			FROM peers
+	rows, err := tx.Query(ctx, `
+		WITH expired AS (
+			SELECT peer FROM peers
 			WHERE last_seen < $1
-			LIMIT 1
+			ORDER BY last_seen, peer
+			LIMIT $2
+			FOR UPDATE SKIP LOCKED
 		)
-		DELETE FROM peers
-		USING d
-		WHERE peers.peer = d.peer
-		RETURNING d.peer, d.disconnected
-	`, now.Add(-threshold)).Scan(&peerID, &disconnected)
+		DELETE FROM peers p USING expired e
+		WHERE p.peer = e.peer
+		RETURNING p.peer
+	`, now.Add(-threshold), limit)
 	if err != nil {
-		if errors.Is(err, pgx.ErrNoRows) {
-			return "", false, nil, nil
+		return nil, err
+	}
+	var result []TimedOutPeer
+	var ids []string
+	indexes := make(map[string]int)
+	for rows.Next() {
+		var id string
+		if err := rows.Scan(&id); err != nil {
+			rows.Close()
+			return nil, err
 		}
-		return "", false, nil, err
+		indexes[id] = len(result)
+		ids = append(ids, id)
+		result = append(result, TimedOutPeer{ID: id, GameLobbies: make(map[string][]string)})
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	if len(result) == 0 {
+		return nil, nil
 	}
 
-	gameLobbies := make(map[string][]string)
-
-	rows, err := tx.Query(ctx, `
-		UPDATE lobbies
-		SET
-			peers = array_remove(peers, $1),
+	// Lock each affected lobby once, retaining its old membership for notifications.
+	// Array overlap uses lobbies_peers; ANY(peers) cannot use that GIN index.
+	rows, err = tx.Query(ctx, `
+		WITH affected AS MATERIALIZED (
+			SELECT game, code, peers FROM lobbies
+			WHERE peers && $1::varchar(20)[]
+			ORDER BY game, code
+			FOR UPDATE
+		)
+		UPDATE lobbies l
+		SET peers = ARRAY(SELECT p FROM unnest(l.peers) AS p WHERE p <> ALL($1::varchar(20)[])),
 			updated_at = $2
-		WHERE $1 = ANY(peers)
-		RETURNING game, code
-	`, peerID, now)
+		FROM affected a
+		WHERE l.game = a.game AND l.code = a.code
+		RETURNING a.game, a.code, a.peers
+	`, ids, now)
 	if err != nil {
-		return "", false, nil, err
+		return nil, err
 	}
-
 	for rows.Next() {
-		var game string
-		var lobby string
-
-		err = rows.Scan(&game, &lobby)
-		if err != nil {
-			return "", false, nil, err
+		var game, code string
+		var members []string
+		if err := rows.Scan(&game, &code, &members); err != nil {
+			rows.Close()
+			return nil, err
+		}
+		for _, id := range members {
+			if i, ok := indexes[id]; ok {
+				result[i].GameLobbies[game] = append(result[i].GameLobbies[game], code)
+			}
 		}
-
-		gameLobbies[game] = append(gameLobbies[game], lobby)
 	}
-
-	if err = rows.Err(); err != nil {
-		return "", false, nil, err
+	if err := rows.Err(); err != nil {
+		return nil, err
 	}
-
-	if err = tx.Commit(ctx); err != nil {
-		return "", false, nil, err
+	if err := tx.Commit(ctx); err != nil {
+		return nil, err
 	}
-
-	return peerID, disconnected, gameLobbies, nil
-}
-
-// ResetAllPeerLastSeen will reset all last_seen.
-// This is being called when the process restarts so it doesn't matter
-// how long the process was down.
-func (s *PostgresStore) ResetAllPeerLastSeen(ctx context.Context) error {
-	now := util.NowUTC(ctx)
-
-	_, err := s.DB.Exec(ctx, `
-		UPDATE peers
-		SET
-			last_seen = $1,
-			updated_at = $1
-	`, now)
-	return err
+	return result, nil
 }
 
 func (s *PostgresStore) CleanEmptyLobbies(ctx context.Context, olderThan time.Time) error {

internal/signaling/stores/shared.go

@@ -34,13 +34,12 @@ type Store interface {
 	Subscribe(ctx context.Context, callback SubscriptionCallback, game, lobby, peerID string)
 	Publish(ctx context.Context, topic string, data []byte) error
 
-	CreatePeer(ctx context.Context, peerID, secret, gameID string) error
+	CreatePeer(ctx context.Context, peerID, secret, gameID, connectionID string) error
 	UpdatePeerGeo(ctx context.Context, peerID string, country, region string) error
-	MarkPeerAsActive(ctx context.Context, peerID string) error
-	MarkPeerAsDisconnected(ctx context.Context, peerID string) error
-	MarkPeerAsReconnected(ctx context.Context, peerID, secret, gameID string) (bool, []string, error)
-	ClaimNextTimedOutPeer(ctx context.Context, threshold time.Duration) (string, bool, map[string][]string, error)
-	ResetAllPeerLastSeen(ctx context.Context) error
+	MarkPeerAsActive(ctx context.Context, peerID, connectionID string) error
+	MarkPeerAsDisconnected(ctx context.Context, peerID, connectionID string) (bool, error)
+	MarkPeerAsReconnected(ctx context.Context, peerID, secret, gameID, connectionID string) (bool, []string, error)
+	ClaimTimedOutPeers(ctx context.Context, threshold time.Duration, limit int) ([]TimedOutPeer, error)
 
 	CleanEmptyLobbies(ctx context.Context, olderThan time.Time) error
 
@@ -83,3 +82,9 @@ type ElectionResult struct {
 	Leader string
 	Term   int
 }
+
+// TimedOutPeer includes the lobby memberships removed in the same transaction.
+type TimedOutPeer struct {
+	ID          string
+	GameLobbies map[string][]string
+}