owpengram-server/internal/rpc/updates.go

package rpc

import (
	"context"
	"time"

	"github.com/gotd/td/bin"
	"github.com/gotd/td/tg"

	"telesrv/internal/domain"
)

const updatesTooLongNudgeDelay = 300 * time.Millisecond
const legacyUpdatesGetDifferenceTypeID uint32 = 0x25939651

// registerUpdates 注册 updates.* RPC handler。
func (r *Router) registerUpdates(d *tg.ServerDispatcher) {
	d.OnUpdatesGetState(r.onUpdatesGetState)
	d.OnUpdatesGetDifference(r.onUpdatesGetDifference)
}

// onUpdatesGetState 处理 updates.getState（第一阶段返回零状态）。
func (r *Router) onUpdatesGetState(ctx context.Context) (*tg.UpdatesState, error) {
	id, _ := AuthKeyIDFrom(ctx)
	userID, _, err := r.currentUserID(ctx)
	if err != nil {
		return nil, internalErr()
	}
	if r.deps.Updates == nil {
		r.markSessionReceivesUpdates(ctx, userID)
		return &tg.UpdatesState{Date: int(r.clock.Now().Unix())}, nil
	}
	st, err := r.deps.Updates.GetState(ctx, id, userID)
	if err != nil {
		return nil, internalErr()
	}
	current, err := r.deps.Updates.CurrentState(ctx, userID)
	if err != nil {
		return nil, internalErr()
	}
	r.markSessionReceivesUpdates(ctx, userID)
	if current.Pts > st.Pts {
		r.scheduleCurrentSessionDifferenceNudge(ctx)
	}
	return ptr(tgUpdateState(st)), nil
}

func (r *Router) onUpdatesGetDifference(ctx context.Context, req *tg.UpdatesGetDifferenceRequest) (tg.UpdatesDifferenceClass, error) {
	id, _ := AuthKeyIDFrom(ctx)
	userID, _, err := r.currentUserID(ctx)
	if err != nil {
		return nil, internalErr()
	}
	if r.deps.Updates == nil {
		now := int(r.clock.Now().Unix())
		r.markSessionReceivesUpdates(ctx, userID)
		return &tg.UpdatesDifferenceEmpty{Date: now}, nil
	}
	st, err := r.deps.Updates.GetDifference(ctx, id, userID, domain.UpdateState{
		Pts:  req.Pts,
		Qts:  req.Qts,
		Date: req.Date,
	})
	if err != nil {
		return nil, internalErr()
	}
	r.markSessionReceivesUpdates(ctx, userID)
	st.ChannelNudges = r.accountChannelDifferenceNudges(ctx, userID, req.Date)
	if len(st.Events) == 0 && len(st.ChannelNudges) == 0 {
		return &tg.UpdatesDifferenceEmpty{Date: st.State.Date, Seq: st.State.Seq}, nil
	}
	st.Events = r.enrichUpdateEvents(ctx, userID, st.Events)
	return tgUpdatesDifference(st), nil
}

func (r *Router) handleLegacyUpdatesGetDifference(ctx context.Context, b *bin.Buffer) (bin.Encoder, error) {
	if err := b.ConsumeID(legacyUpdatesGetDifferenceTypeID); err != nil {
		return nil, err
	}
	var flags bin.Fields
	if err := flags.Decode(b); err != nil {
		return nil, err
	}
	pts, err := b.Int()
	if err != nil {
		return nil, err
	}
	var ptsTotalLimit int
	if flags.Has(0) {
		ptsTotalLimit, err = b.Int()
		if err != nil {
			return nil, err
		}
	}
	date, err := b.Int()
	if err != nil {
		return nil, err
	}
	qts, err := b.Int()
	if err != nil {
		return nil, err
	}
	return r.onUpdatesGetDifference(ctx, &tg.UpdatesGetDifferenceRequest{
		Flags:         flags,
		Pts:           pts,
		PtsTotalLimit: ptsTotalLimit,
		Date:          date,
		Qts:           qts,
	})
}

func (r *Router) accountChannelDifferenceNudges(ctx context.Context, userID int64, sinceDate int) []domain.ChannelDifferenceNudge {
	if r.deps.Channels == nil || userID == 0 || sinceDate <= 0 {
		return nil
	}
	dirty, err := r.deps.Channels.DirtyActiveChannelsForUser(ctx, userID, sinceDate, 0, domain.MaxChannelDifferenceLimit)
	if err != nil || len(dirty) == 0 {
		return nil
	}
	out := make([]domain.ChannelDifferenceNudge, 0, len(dirty))
	for _, item := range dirty {
		if item.ChannelID == 0 {
			continue
		}
		out = append(out, domain.ChannelDifferenceNudge{ChannelID: item.ChannelID, Pts: item.Pts})
	}
	return out
}

func (r *Router) markSessionReceivesUpdates(ctx context.Context, userID int64) {
	if r.deps.Sessions == nil {
		return
	}
	r.syncSessionChannelMemberships(ctx, userID)
	sessionID, ok := SessionIDFrom(ctx)
	if !ok {
		return
	}
	if scoped, ok := r.scopedSessions(); ok {
		if rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx); ok {
			scoped.SetReceivesUpdatesForAuthKey(rawAuthKeyID, sessionID, true)
			return
		}
	}
	r.deps.Sessions.SetReceivesUpdates(sessionID, true)
}

func (r *Router) scheduleCurrentSessionDifferenceNudge(ctx context.Context) {
	pushCtx := context.Background()
	if sessionID, ok := SessionIDFrom(ctx); ok {
		pushCtx = WithSessionID(pushCtx, sessionID)
	}
	if rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx); ok {
		pushCtx = WithRawAuthKeyID(pushCtx, rawAuthKeyID)
	}
	if authKeyID, ok := AuthKeyIDFrom(ctx); ok {
		pushCtx = WithAuthKeyID(pushCtx, authKeyID)
	}
	time.AfterFunc(updatesTooLongNudgeDelay, func() {
		r.pushCurrentSessionMessage(pushCtx, "push updatesTooLong after getState", &tg.UpdatesTooLong{})
	})
}

func ptr[T any](v T) *T { return &v }