Skip to content

Commit 386f2ed

Browse files
erikdubbelboererikdubbelboer
authored
Fix expired peer cleanup and disconnect handling (#409)
Tic Tac Toe Ultimate is generating peers quicker than our cleanup process can currently handle. The game is very popular. Clean up expired peers in batches of 250 and update each affected lobby once per batch. Add a GIN index for lobby membership lookups and elect replacement leaders once per affected lobby. Replace the startup timestamp reset with a reconnect grace period so restarts don't revive stale peers or block reconnections. Mark peers disconnected on explicit close and prevent late heartbeats from refreshing them. Cancel connections after failed pings and bound heartbeat database updates with a timeout.
1 parent 14e6dc0 commit 386f2ed

8 files changed

Lines changed: 179 additions & 128 deletions

File tree

‎internal/signaling/handler.go‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,11 @@ import (
1515
"github.com/poki/netlib/internal/metrics"
1616
"github.com/poki/netlib/internal/signaling/stores"
1717
"github.com/poki/netlib/internal/util"
18+
"github.com/rs/xid"
1819
"go.uber.org/zap"
1920
)
2021

21-
const LobbyCleanInterval = 30 * time.Minute
22+
const LobbyCleanInterval = 5 * time.Minute
2223
const LobbyCleanThreshold = 24 * time.Hour
2324
const peerPingDuration = 2 * time.Second
2425
const peerActiveUpdateInterval = 30 * time.Second
@@ -95,15 +96,17 @@ func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.Cre
9596
}
9697

9798
peer := &Peer{
98-
store: store,
99-
conn: conn,
99+
connectionID: xid.New().String(),
100+
store: store,
101+
conn: conn,
100102

101103
retrievedIDCallback: manager.Reconnected,
102104

103105
Country: country,
104106
Region: region,
105107
}
106108
defer func() {
109+
cancel()
107110
logger.Debug("peer websocket closed", zap.String("peer", peer.ID), zap.String("game", peer.Game), zap.String("origin", r.Header.Get("Origin")))
108111
conn.Close(websocket.StatusInternalError, "unexpected closure") // nolint:errcheck
109112

@@ -122,21 +125,27 @@ func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.Cre
122125
for {
123126
select {
124127
case <-ticker.C:
125-
if err := peer.Send(ctx, PingPacket{Type: "ping"}); err != nil {
128+
// Limit ping writes so a stalled connection cannot keep this goroutine alive indefinitely.
129+
pingCtx, pingCancel := context.WithTimeout(ctx, 10*time.Second)
130+
err := peer.Send(pingCtx, PingPacket{Type: "ping"})
131+
pingCancel()
132+
if err != nil {
126133
if !util.ShouldIgnoreNetworkError(err) {
127134
if strings.Contains(err.Error(), "write: broken pipe") || strings.Contains(err.Error(), "connection reset by peer") {
128135
logger.Warn("failed to send ping packet", zap.String("peer", peer.ID), zap.Error(err))
129136
} else {
130137
logger.Error("failed to send ping packet", zap.String("peer", peer.ID), zap.Error(err))
131138
}
132139
}
140+
cancel()
141+
return
133142
} else {
134143
// If we can send a ping packet, and the peer has an ID, we update the peer as being active.
135144
// If the peer doesn't have an ID yet, it's still in the process of connecting, so we don't update it.
136145
if peer.ID != "" {
137146
now := util.NowUTC(ctx)
138147
if lastActiveUpdate.IsZero() || now.Sub(lastActiveUpdate) >= peerActiveUpdateInterval {
139-
manager.MarkPeerAsActive(ctx, peer.ID)
148+
manager.MarkPeerAsActive(ctx, peer.ID, peer.connectionID)
140149
lastActiveUpdate = now
141150
}
142151
}

‎internal/signaling/peer.go‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,9 @@ type Peer struct {
2424
conn *websocket.Conn
2525

2626
closedPacketReceived bool
27+
connectionID string
2728

28-
retrievedIDCallback func(context.Context, string, string, string) (bool, []string, error)
29+
retrievedIDCallback func(context.Context, string, string, string, string) (bool, []string, error)
2930

3031
ID string
3132
Secret string
@@ -197,7 +198,7 @@ func (p *Peer) HandleHelloPacket(ctx context.Context, packet HelloPacket) error
197198
if packet.ID != "" && packet.Secret != "" {
198199
logger.Debug("peer reconnecting", zap.String("game", packet.Game), zap.String("peer", packet.ID), zap.String("version", packet.Version))
199200
var err error
200-
hasReconnected, reconnectingLobbies, err = p.retrievedIDCallback(ctx, packet.ID, packet.Secret, packet.Game)
201+
hasReconnected, reconnectingLobbies, err = p.retrievedIDCallback(ctx, packet.ID, packet.Secret, packet.Game, p.connectionID)
201202
if err != nil {
202203
return fmt.Errorf("unable to reconnect: %w", err)
203204
}
@@ -223,7 +224,7 @@ func (p *Peer) HandleHelloPacket(ctx context.Context, packet HelloPacket) error
223224
p.Secret = util.GenerateSecret(ctx)
224225
logger.Debug("peer connecting", zap.String("game", p.Game), zap.String("peer", p.ID), zap.String("version", packet.Version))
225226

226-
if err := p.store.CreatePeer(ctx, p.ID, p.Secret, p.Game); err != nil {
227+
if err := p.store.CreatePeer(ctx, p.ID, p.Secret, p.Game, p.connectionID); err != nil {
227228
return fmt.Errorf("unable to create peer: %w", err)
228229
}
229230
}
@@ -288,7 +289,16 @@ func (p *Peer) HandleClosePacket(ctx context.Context, packet ClosePacket) error
288289
logger := logging.GetLogger(ctx)
289290
go metrics.Record(ctx, "client", "close", p.Game, p.ID, p.Lobby)
290291

291-
p.closedPacketReceived = true
292+
if p.ID != "" {
293+
disconnected, err := p.store.MarkPeerAsDisconnected(ctx, p.ID, p.connectionID)
294+
if err != nil {
295+
return fmt.Errorf("unable to mark peer as disconnected: %w", err)
296+
}
297+
if !disconnected {
298+
p.closedPacketReceived = true
299+
return nil
300+
}
301+
}
292302

293303
logger.Debug("client closed",
294304
zap.String("game", p.Game),
@@ -322,6 +332,8 @@ func (p *Peer) HandleClosePacket(ctx context.Context, packet ClosePacket) error
322332
p.Lobby = ""
323333
}
324334

335+
// Keep deferred disconnect handling enabled if any cleanup above fails.
336+
p.closedPacketReceived = true
325337
return nil
326338
}
327339

‎internal/signaling/stores/postgres.go‎

Lines changed: 89 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -405,7 +405,7 @@ func (s *PostgresStore) ListLobbies(ctx context.Context, game string, country, r
405405
return lobbies, nil
406406
}
407407

408-
func (s *PostgresStore) CreatePeer(ctx context.Context, peerID, secret, gameID string) error {
408+
func (s *PostgresStore) CreatePeer(ctx context.Context, peerID, secret, gameID, connectionID string) error {
409409
if len(peerID) > 20 {
410410
logger := logging.GetLogger(ctx)
411411
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
414414

415415
now := util.NowUTC(ctx)
416416
_, err := s.DB.Exec(ctx, `
417-
INSERT INTO peers (peer, secret, game, last_seen, updated_at)
418-
VALUES ($1, $2, $3, $4, $4)
419-
`, peerID, secret, gameID, now)
417+
INSERT INTO peers (peer, secret, game, last_seen, updated_at, connection_id)
418+
VALUES ($1, $2, $3, $4, $4, $5)
419+
`, peerID, secret, gameID, now, connectionID)
420420
if err != nil {
421421
return err
422422
}
@@ -440,46 +440,49 @@ func (s *PostgresStore) UpdatePeerGeo(ctx context.Context, peerID string, countr
440440
return err
441441
}
442442

443-
func (s *PostgresStore) MarkPeerAsActive(ctx context.Context, peerID string) error {
443+
func (s *PostgresStore) MarkPeerAsActive(ctx context.Context, peerID, connectionID string) error {
444+
// Only an authenticated reconnect may reactivate a disconnected peer.
445+
// Late heartbeats must not revive it or postpone timeout cleanup.
444446
now := util.NowUTC(ctx)
445447

446448
_, err := s.DB.Exec(ctx, `
447449
UPDATE peers
448450
SET
449-
disconnected = FALSE,
450451
last_seen = $1,
451452
updated_at = $1
452-
WHERE peer = $2
453-
`, now, peerID)
453+
WHERE peer = $2 AND connection_id = $3 AND disconnected = FALSE
454+
`, now, peerID, connectionID)
454455
return err
455456
}
456457

457-
func (s *PostgresStore) MarkPeerAsDisconnected(ctx context.Context, peerID string) error {
458+
// MarkPeerAsDisconnected ignores handlers superseded by an authenticated reconnect.
459+
func (s *PostgresStore) MarkPeerAsDisconnected(ctx context.Context, peerID, connectionID string) (bool, error) {
458460
now := util.NowUTC(ctx)
459-
460-
_, err := s.DB.Exec(ctx, `
461+
result, err := s.DB.Exec(ctx, `
461462
UPDATE peers
462-
SET
463-
disconnected = TRUE,
464-
updated_at = $1
465-
WHERE peer = $2
466-
`, now, peerID)
467-
return err
463+
SET disconnected = TRUE, updated_at = $1
464+
WHERE peer = $2 AND connection_id = $3
465+
`, now, peerID, connectionID)
466+
if err != nil {
467+
return false, err
468+
}
469+
return result.RowsAffected() > 0, nil
468470
}
469471

470-
func (s *PostgresStore) MarkPeerAsReconnected(ctx context.Context, peerID, secret, gameID string) (bool, []string, error) {
472+
func (s *PostgresStore) MarkPeerAsReconnected(ctx context.Context, peerID, secret, gameID, connectionID string) (bool, []string, error) {
471473
now := util.NowUTC(ctx)
472474

473475
result, err := s.DB.Exec(ctx, `
474476
UPDATE peers
475477
SET
476478
disconnected = FALSE,
479+
connection_id = $5,
477480
last_seen = $1,
478481
updated_at = $1
479482
WHERE peer = $2
480483
AND secret = $3
481484
AND game = $4
482-
`, now, peerID, secret, gameID)
485+
`, now, peerID, secret, gameID, connectionID)
483486
if err != nil {
484487
return false, nil, err
485488
}
@@ -516,86 +519,93 @@ func (s *PostgresStore) MarkPeerAsReconnected(ctx context.Context, peerID, secre
516519
return true, lobbies, nil
517520
}
518521

519-
func (s *PostgresStore) ClaimNextTimedOutPeer(ctx context.Context, threshold time.Duration) (string, bool, map[string][]string, error) {
522+
func (s *PostgresStore) ClaimTimedOutPeers(ctx context.Context, threshold time.Duration, limit int) ([]TimedOutPeer, error) {
523+
if limit <= 0 {
524+
return nil, fmt.Errorf("cleanup limit must be positive")
525+
}
526+
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
527+
defer cancel()
520528
now := util.NowUTC(ctx)
521-
522529
tx, err := s.DB.Begin(ctx)
523530
if err != nil {
524-
return "", false, nil, err
531+
return nil, err
525532
}
526533
defer tx.Rollback(context.Background()) //nolint:errcheck
527534

528-
var peerID string
529-
var disconnected bool
530-
err = tx.QueryRow(ctx, `
531-
WITH d AS (
532-
SELECT peer, disconnected
533-
FROM peers
535+
rows, err := tx.Query(ctx, `
536+
WITH expired AS (
537+
SELECT peer FROM peers
534538
WHERE last_seen < $1
535-
LIMIT 1
539+
ORDER BY last_seen, peer
540+
LIMIT $2
541+
FOR UPDATE SKIP LOCKED
536542
)
537-
DELETE FROM peers
538-
USING d
539-
WHERE peers.peer = d.peer
540-
RETURNING d.peer, d.disconnected
541-
`, now.Add(-threshold)).Scan(&peerID, &disconnected)
543+
DELETE FROM peers p USING expired e
544+
WHERE p.peer = e.peer
545+
RETURNING p.peer
546+
`, now.Add(-threshold), limit)
542547
if err != nil {
543-
if errors.Is(err, pgx.ErrNoRows) {
544-
return "", false, nil, nil
548+
return nil, err
549+
}
550+
var result []TimedOutPeer
551+
var ids []string
552+
indexes := make(map[string]int)
553+
for rows.Next() {
554+
var id string
555+
if err := rows.Scan(&id); err != nil {
556+
rows.Close()
557+
return nil, err
545558
}
546-
return "", false, nil, err
559+
indexes[id] = len(result)
560+
ids = append(ids, id)
561+
result = append(result, TimedOutPeer{ID: id, GameLobbies: make(map[string][]string)})
562+
}
563+
if err := rows.Err(); err != nil {
564+
return nil, err
565+
}
566+
if len(result) == 0 {
567+
return nil, nil
547568
}
548569

549-
gameLobbies := make(map[string][]string)
550-
551-
rows, err := tx.Query(ctx, `
552-
UPDATE lobbies
553-
SET
554-
peers = array_remove(peers, $1),
570+
// Lock each affected lobby once, retaining its old membership for notifications.
571+
// Array overlap uses lobbies_peers; ANY(peers) cannot use that GIN index.
572+
rows, err = tx.Query(ctx, `
573+
WITH affected AS MATERIALIZED (
574+
SELECT game, code, peers FROM lobbies
575+
WHERE peers && $1::varchar(20)[]
576+
ORDER BY game, code
577+
FOR UPDATE
578+
)
579+
UPDATE lobbies l
580+
SET peers = ARRAY(SELECT p FROM unnest(l.peers) AS p WHERE p <> ALL($1::varchar(20)[])),
555581
updated_at = $2
556-
WHERE $1 = ANY(peers)
557-
RETURNING game, code
558-
`, peerID, now)
582+
FROM affected a
583+
WHERE l.game = a.game AND l.code = a.code
584+
RETURNING a.game, a.code, a.peers
585+
`, ids, now)
559586
if err != nil {
560-
return "", false, nil, err
587+
return nil, err
561588
}
562-
563589
for rows.Next() {
564-
var game string
565-
var lobby string
566-
567-
err = rows.Scan(&game, &lobby)
568-
if err != nil {
569-
return "", false, nil, err
590+
var game, code string
591+
var members []string
592+
if err := rows.Scan(&game, &code, &members); err != nil {
593+
rows.Close()
594+
return nil, err
595+
}
596+
for _, id := range members {
597+
if i, ok := indexes[id]; ok {
598+
result[i].GameLobbies[game] = append(result[i].GameLobbies[game], code)
599+
}
570600
}
571-
572-
gameLobbies[game] = append(gameLobbies[game], lobby)
573601
}
574-
575-
if err = rows.Err(); err != nil {
576-
return "", false, nil, err
602+
if err := rows.Err(); err != nil {
603+
return nil, err
577604
}
578-
579-
if err = tx.Commit(ctx); err != nil {
580-
return "", false, nil, err
605+
if err := tx.Commit(ctx); err != nil {
606+
return nil, err
581607
}
582-
583-
return peerID, disconnected, gameLobbies, nil
584-
}
585-
586-
// ResetAllPeerLastSeen will reset all last_seen.
587-
// This is being called when the process restarts so it doesn't matter
588-
// how long the process was down.
589-
func (s *PostgresStore) ResetAllPeerLastSeen(ctx context.Context) error {
590-
now := util.NowUTC(ctx)
591-
592-
_, err := s.DB.Exec(ctx, `
593-
UPDATE peers
594-
SET
595-
last_seen = $1,
596-
updated_at = $1
597-
`, now)
598-
return err
608+
return result, nil
599609
}
600610

601611
func (s *PostgresStore) CleanEmptyLobbies(ctx context.Context, olderThan time.Time) error {

‎internal/signaling/stores/shared.go‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,13 +34,12 @@ type Store interface {
3434
Subscribe(ctx context.Context, callback SubscriptionCallback, game, lobby, peerID string)
3535
Publish(ctx context.Context, topic string, data []byte) error
3636

37-
CreatePeer(ctx context.Context, peerID, secret, gameID string) error
37+
CreatePeer(ctx context.Context, peerID, secret, gameID, connectionID string) error
3838
UpdatePeerGeo(ctx context.Context, peerID string, country, region string) error
39-
MarkPeerAsActive(ctx context.Context, peerID string) error
40-
MarkPeerAsDisconnected(ctx context.Context, peerID string) error
41-
MarkPeerAsReconnected(ctx context.Context, peerID, secret, gameID string) (bool, []string, error)
42-
ClaimNextTimedOutPeer(ctx context.Context, threshold time.Duration) (string, bool, map[string][]string, error)
43-
ResetAllPeerLastSeen(ctx context.Context) error
39+
MarkPeerAsActive(ctx context.Context, peerID, connectionID string) error
40+
MarkPeerAsDisconnected(ctx context.Context, peerID, connectionID string) (bool, error)
41+
MarkPeerAsReconnected(ctx context.Context, peerID, secret, gameID, connectionID string) (bool, []string, error)
42+
ClaimTimedOutPeers(ctx context.Context, threshold time.Duration, limit int) ([]TimedOutPeer, error)
4443

4544
CleanEmptyLobbies(ctx context.Context, olderThan time.Time) error
4645

@@ -83,3 +82,9 @@ type ElectionResult struct {
8382
Leader string
8483
Term int
8584
}
85+
86+
// TimedOutPeer includes the lobby memberships removed in the same transaction.
87+
type TimedOutPeer struct {
88+
ID string
89+
GameLobbies map[string][]string
90+
}

0 commit comments

Comments
 (0)