Warp/server/internal/syncapi/slidingconn.go

package syncapi

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"maps"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Sliding sync connections: what the server remembers about what each
// client has already been sent.
//
// A classic sync token is a position in every stream and nothing else; the
// client is sent everything after it. Sliding sync cannot work that way,
// because what a client is sent depends on what it asked for — which rooms
// are in the window of which list, with which pieces of state — and that
// changes from request to request. So the server remembers, per connection,
// which rooms it has sent and as of which position, and `pos` names one of
// those memories.
//
// The memory is in process and is lost on restart. That is what the
// protocol expects: an unknown `pos` is answered with M_UNKNOWN_POS, and
// the client starts a new connection — one initial request, after which it
// is in step again. Persisting it would buy a slightly cheaper restart for
// a write on every request.

// ErrUnknownPos is a `pos` this server does not remember: expired, from
// before a restart, or never handed out.
var ErrUnknownPos = errors.New("unknown sliding sync position")

// Limits of the connection store.
const (
	// connectionTTL is how long an idle connection is remembered. A client
	// that comes back later starts again.
	connectionTTL = 30 * time.Minute

	// maxConnectionsPerDevice caps the connections one device may hold
	// open, each conn_id being one. The oldest goes first.
	maxConnectionsPerDevice = 10

	// keptGenerations is how many of a connection's latest positions stay
	// usable. A client that never saw a response retries with the pos it
	// sent, so the one before the latest has to still work.
	keptGenerations = 3
)

// sentRoom is what a connection has been sent of one room.
type sentRoom struct {
	// pos is the event stream position the room was last sent up to.
	pos int64

	// config is the room's required_state and timeline_limit as of that
	// send, rendered as a string: when it changes, the state the client
	// holds is no longer what it asked for.
	config string

	// membership is the user's membership as of that send.
	membership string

	// typing is the typing notification last sent, as its JSON.
	typing string
}

// connSnapshot is a connection as of one `pos`.
type connSnapshot struct {
	// token is the stream positions the response at this pos read up to.
	token Token

	// startPDU is the event stream position the connection began at: rooms
	// the user left before it are not news to it.
	startPDU int64

	rooms map[string]sentRoom

	// lists is each list's room count as last sent.
	lists map[string]int

	// toDevice is where the to_device extension has got to.
	toDevice int64
}

// clone copies a snapshot so that a new generation can be built from it
// without touching the one a retry might still need.
func (c *connSnapshot) clone() *connSnapshot {
	return &connSnapshot{
		token:    c.token,
		startPDU: c.startPDU,
		rooms:    maps.Clone(c.rooms),
		lists:    maps.Clone(c.lists),
		toDevice: c.toDevice,
	}
}

// slidingConn is one connection: a client device's conn_id.
type slidingConn struct {
	// nonce makes the connection's positions its own: a pos from another
	// connection, or from before a restart, cannot name one of these
	// generations by accident.
	nonce string

	generations map[int64]*connSnapshot
	latest      int64
	lastUsed    time.Time
}

type connKey struct {
	user, device, connID string
}

// connStore holds every connection.
type connStore struct {
	mu    sync.Mutex
	conns map[connKey]*slidingConn
	now   func() time.Time
}

func newConnStore(now func() time.Time) *connStore {
	return &connStore{conns: map[connKey]*slidingConn{}, now: now}
}

// load returns the snapshot a request's pos names, or a fresh one for a
// request with none. A request without a pos starts the connection over,
// which is what a client does after M_UNKNOWN_POS and on first launch.
func (s *connStore) load(key connKey, pos string) (*connSnapshot, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.sweep()

	if pos == "" {
		delete(s.conns, key)
		return nil, nil
	}
	conn, ok := s.conns[key]
	if !ok {
		return nil, fmt.Errorf("%w: the connection has expired", ErrUnknownPos)
	}
	nonce, raw, ok := strings.Cut(pos, ".")
	gen, err := strconv.ParseInt(raw, 10, 64)
	if !ok || err != nil || nonce != conn.nonce {
		return nil, fmt.Errorf("%w: %q is not a position of this connection", ErrUnknownPos, pos)
	}
	snap, ok := conn.generations[gen]
	if !ok {
		return nil, fmt.Errorf("%w: %q is too old", ErrUnknownPos, pos)
	}
	conn.lastUsed = s.now()
	return snap, nil
}

// save stores a new generation of a connection and returns its pos.
func (s *connStore) save(key connKey, snap *connSnapshot) string {
	s.mu.Lock()
	defer s.mu.Unlock()

	conn, ok := s.conns[key]
	if !ok {
		conn = &slidingConn{nonce: randomNonce(), generations: map[int64]*connSnapshot{}, lastUsed: s.now()}
		s.conns[key] = conn
		s.capDevice(key)
	}
	conn.latest++
	conn.generations[conn.latest] = snap
	delete(conn.generations, conn.latest-keptGenerations)
	conn.lastUsed = s.now()
	return conn.nonce + "." + strconv.FormatInt(conn.latest, 10)
}

// sweep forgets idle connections. It runs under the lock on every load,
// which keeps the store bounded without a goroutine of its own.
func (s *connStore) sweep() {
	cutoff := s.now().Add(-connectionTTL)
	for key, conn := range s.conns {
		if conn.lastUsed.Before(cutoff) {
			delete(s.conns, key)
		}
	}
}

// capDevice drops a device's oldest connections beyond the cap. The
// connection being saved is never one of them.
func (s *connStore) capDevice(key connKey) {
	var others []connKey
	for k := range s.conns {
		if k != key && k.user == key.user && k.device == key.device {
			others = append(others, k)
		}
	}
	slices.SortFunc(others, func(a, b connKey) int {
		return s.conns[a].lastUsed.Compare(s.conns[b].lastUsed)
	})
	for len(others) > maxConnectionsPerDevice-1 {
		delete(s.conns, others[0])
		others = others[1:]
	}
}

// connections reports how many connections are held, for tests and
// metrics.
func (s *connStore) connections() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.conns)
}

func randomNonce() string {
	var b [6]byte
	_, _ = rand.Read(b[:])
	return hex.EncodeToString(b[:])
}

// SlidingConnections reports how many sliding sync connections the server
// is holding state for.
func (s *API) SlidingConnections() int { return s.conns.connections() }