astra/owpengram-server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: defer login bootstrap updates

Browse source
This commit is contained in:
A 2026-07-05 03:28:19 +08:00
parent a246fd5417
commit 4d9f1e271d
18 changed files with 958 additions and 188 deletions
Download patch file Download diff file Expand all files Collapse all files

84
internal/rpc/auth.go
View file

@ -20,8 +20,6 @@ import (
// devCodeLength 是开发固定验证码长度，写入 auth.sentCode 的 type.length。
const devCodeLength = 5
const loginMessagePushDelay = 2 * time.Second
// registerAuth 注册 auth.* RPC handler。
func (r *Router) registerAuth(d *tg.ServerDispatcher) {
d.OnAuthBindTempAuthKey(r.onAuthBindTempAuthKey)
@ -329,7 +327,7 @@ func (r *Router) onAuthSignIn(ctx context.Context, req *tg.AuthSignInRequest) (t
r.setAuthUserCache(id, u.ID, true)
}
r.bindSessionUser(ctx, u.ID)
r.recordAndScheduleLoginMessagePush(ctx, loginMessage)
r.enqueueLoginMessageBootstrap(ctx, loginMessage)
r.pushSignInServiceNotificationToOthers(ctx, u)
return &tg.AuthAuthorization{User: r.tgSelfUser(u)}, nil
}
@ -598,7 +596,7 @@ func (r *Router) onAuthSignUp(ctx context.Context, req *tg.AuthSignUpRequest) (t
r.setAuthUserCache(id, u.ID, true)
}
r.bindSessionUser(ctx, u.ID)
r.recordAndScheduleLoginMessagePush(ctx, loginMessage)
r.enqueueLoginMessageBootstrap(ctx, loginMessage)
return &tg.AuthAuthorization{User: r.tgSelfUser(u)}, nil
}
@ -734,84 +732,6 @@ func (r *Router) pushSignInServiceNotificationToOthers(ctx context.Context, u do
}()
}
func (r *Router) recordAndScheduleLoginMessagePush(ctx context.Context, msg domain.Message) {
authKeyID, hasAuthKeyID := AuthKeyIDFrom(ctx)
sessionID, hasSessionID := SessionIDFrom(ctx)
if !hasAuthKeyID || !hasSessionID || msg.ID == 0 {
return
}
event := domain.UpdateEvent{Type: domain.UpdateEventNewMessage, Pts: 1, PtsCount: 1, Date: msg.Date, Message: msg}
state := domain.UpdateState{Pts: 1, Date: msg.Date, Seq: 0}
if r.deps.Updates != nil {
recorded, st, err := r.deps.Updates.RecordNewMessage(ctx, authKeyID, msg.OwnerUserID, msg)
if err != nil {
r.log.Warn("record login message update", zap.Error(err))
return
}
event = recorded
state = st
}
if r.deps.Sessions == nil {
return
}
// 提前从请求 ctx 取出 rawAuthKeyID（值类型），闭包只捕获该值、不捕获请求 ctx——
// 避免延迟推送的 AfterFunc 在 loginMessagePushDelay 期间延长请求 ctx 链路的存活。
rawAuthKeyID, _ := RawAuthKeyIDFrom(ctx)
time.AfterFunc(loginMessagePushDelay, func() {
pushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
r.pushLoginMessage(pushCtx, rawAuthKeyID, sessionID, event, state)
})
}
func (r *Router) pushLoginMessage(ctx context.Context, rawAuthKeyID [8]byte, sessionID int64, event domain.UpdateEvent, state domain.UpdateState) {
if r.deps.Sessions == nil || event.Message.ID == 0 {
return
}
updates := tgLoginMessageUpdates(event, state)
if updates == nil {
return
}
var err error
if scoped, ok := r.scopedSessions(); ok && rawAuthKeyID != ([8]byte{}) {
err = scoped.PushToSessionForAuthKey(ctx, rawAuthKeyID, sessionID, proto.MessageFromServer, updates)
} else {
err = r.deps.Sessions.PushToSession(ctx, sessionID, proto.MessageFromServer, updates)
}
if err != nil {
r.log.Debug("push login message", zap.Int64("session_id", sessionID), zap.Error(err))
return
}
r.log.Debug("pushed login message",
zap.Int64("session_id", sessionID),
zap.Int("message_id", event.Message.ID),
zap.Int("pts", event.Pts),
zap.Int("seq", state.Seq),
)
}
func tgLoginMessageUpdates(event domain.UpdateEvent, state domain.UpdateState) *tg.Updates {
item := tgMessage(event.Message)
if item == nil {
return nil
}
if state.Date == 0 {
state.Date = event.Date
}
return &tg.Updates{
Updates: []tg.UpdateClass{
&tg.UpdateNewMessage{
Message: item,
Pts: event.Pts,
PtsCount: event.PtsCount,
},
},
Users: []tg.UserClass{tgUser(domain.OfficialSystemUser())},
Date: state.Date,
Seq: state.Seq,
}
}
func (r *Router) tgSignInServiceNotification(ctx context.Context, u domain.User, authKeyID [8]byte) *tg.Updates {
now := r.clock.Now()
client := "Unknown device"

185
internal/rpc/bootstrap_updates.go Normal file
View file

@ -0,0 +1,185 @@
package rpc
import (
"context"
"fmt"
"time"
"go.uber.org/zap"
"telesrv/internal/domain"
"telesrv/internal/postresponse"
)
const (
defaultBootstrapBatch = 100
defaultBootstrapInterval = 200 * time.Millisecond
defaultBootstrapLease = 30 * time.Second
)
type BootstrapUpdateDispatcher struct {
router *Router
log *zap.Logger
batch int
interval time.Duration
maxIdleInterval time.Duration
leaseTimeout time.Duration
}
type BootstrapUpdateOption func(*BootstrapUpdateDispatcher)
func WithBootstrapUpdateBatch(n int) BootstrapUpdateOption {
return func(d *BootstrapUpdateDispatcher) {
if n > 0 {
d.batch = n
}
}
}
func WithBootstrapUpdateInterval(interval time.Duration) BootstrapUpdateOption {
return func(d *BootstrapUpdateDispatcher) {
if interval > 0 {
d.interval = interval
}
}
}
func WithBootstrapUpdateLease(timeout time.Duration) BootstrapUpdateOption {
return func(d *BootstrapUpdateDispatcher) {
if timeout > 0 {
d.leaseTimeout = timeout
}
}
}
func NewBootstrapUpdateDispatcher(router *Router, log *zap.Logger, opts ...BootstrapUpdateOption) *BootstrapUpdateDispatcher {
if log == nil {
log = zap.NewNop()
}
d := &BootstrapUpdateDispatcher{
router: router,
log: log,
batch: defaultBootstrapBatch,
interval: defaultBootstrapInterval,
maxIdleInterval: defaultIdleDispatchMaxInterval,
leaseTimeout: defaultBootstrapLease,
}
for _, opt := range opts {
if opt != nil {
opt(d)
}
}
return d
}
func (d *BootstrapUpdateDispatcher) Run(ctx context.Context) {
if d == nil || d.router == nil {
return
}
runIdleBackoffLoop(ctx, d.interval, d.maxIdleInterval, d.DispatchOnce)
}
func (d *BootstrapUpdateDispatcher) DispatchOnce(ctx context.Context) bool {
return d.router.publishReadyBootstrapUpdates(ctx, d.batch, d.leaseTimeout, d.log) > 0
}
func (r *Router) enqueueLoginMessageBootstrap(ctx context.Context, msg domain.Message) {
if r.deps.BootstrapUpdates == nil || msg.OwnerUserID == 0 || msg.ID == 0 {
return
}
authKeyID, hasAuthKeyID := AuthKeyIDFrom(ctx)
sessionID, hasSessionID := SessionIDFrom(ctx)
if !hasAuthKeyID || !hasSessionID {
return
}
if _, err := r.deps.BootstrapUpdates.EnqueueLoginMessage(ctx, domain.BootstrapUpdateJob{
Kind: domain.BootstrapUpdateJobLoginMessage,
UserID: msg.OwnerUserID,
AuthKeyID: authKeyID,
SessionID: sessionID,
MessageBoxID: msg.ID,
Status: domain.BootstrapUpdateJobPending,
}); err != nil {
r.log.Warn("enqueue bootstrap login message", zap.Int64("user_id", msg.OwnerUserID), zap.Int("message_id", msg.ID), zap.Error(err))
}
}
func (r *Router) registerBootstrapAfterBaseline(ctx context.Context, userID int64) {
if r.deps.BootstrapUpdates == nil || userID == 0 {
return
}
authKeyID, hasAuthKeyID := AuthKeyIDFrom(ctx)
sessionID, hasSessionID := SessionIDFrom(ctx)
if !hasAuthKeyID || !hasSessionID {
return
}
postresponse.Register(ctx, func() {
cbCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ready, err := r.deps.BootstrapUpdates.MarkReadyForSession(cbCtx, userID, authKeyID, sessionID)
if err != nil {
r.log.Warn("mark bootstrap updates ready", zap.Int64("user_id", userID), zap.Int64("session_id", sessionID), zap.Error(err))
return
}
if ready == 0 {
return
}
r.publishReadyBootstrapUpdates(cbCtx, ready, defaultBootstrapLease, r.log.Named("bootstrap"))
})
}
func (r *Router) publishReadyBootstrapUpdates(ctx context.Context, batch int, leaseTimeout time.Duration, log *zap.Logger) int {
if r.deps.BootstrapUpdates == nil || r.deps.Updates == nil || r.deps.Messages == nil {
return 0
}
if log == nil {
log = zap.NewNop()
}
jobs, err := r.deps.BootstrapUpdates.ClaimReady(ctx, batch, leaseTimeout)
if err != nil {
log.Warn("claim bootstrap updates", zap.Error(err))
return 0
}
for _, job := range jobs {
if err := r.publishBootstrapUpdateJob(ctx, job); err != nil {
log.Warn("publish bootstrap update",
zap.Int64("job_id", job.ID),
zap.Int64("user_id", job.UserID),
zap.String("kind", string(job.Kind)),
zap.Error(err),
)
_ = r.deps.BootstrapUpdates.MarkFailed(ctx, job.ID, err.Error())
}
}
return len(jobs)
}
func (r *Router) publishBootstrapUpdateJob(ctx context.Context, job domain.BootstrapUpdateJob) error {
switch job.Kind {
case domain.BootstrapUpdateJobLoginMessage:
return r.publishBootstrapLoginMessage(ctx, job)
default:
return fmt.Errorf("unknown bootstrap update kind %q", job.Kind)
}
}
func (r *Router) publishBootstrapLoginMessage(ctx context.Context, job domain.BootstrapUpdateJob) error {
list, err := r.deps.Messages.GetMessages(ctx, job.UserID, []int{job.MessageBoxID})
if err != nil {
return fmt.Errorf("load bootstrap login message: %w", err)
}
if len(list.Messages) == 0 {
return fmt.Errorf("bootstrap login message missing: box_id=%d", job.MessageBoxID)
}
msg := list.Messages[0]
if msg.OwnerUserID != job.UserID || msg.ID != job.MessageBoxID {
return fmt.Errorf("bootstrap login message mismatch: owner=%d id=%d", msg.OwnerUserID, msg.ID)
}
if _, _, err := r.deps.Updates.PublishNewMessage(ctx, job.UserID, msg); err != nil {
return fmt.Errorf("publish bootstrap login message update: %w", err)
}
if err := r.deps.BootstrapUpdates.MarkPublished(ctx, job.ID); err != nil {
return err
}
return nil
}

60
internal/rpc/deps.go
View file

@ -315,6 +315,7 @@ type UpdatesService interface {
GetDifference(ctx context.Context, authKeyID [8]byte, userID int64, from domain.UpdateState) (domain.UpdateDifference, error)
ClearAuthKey(ctx context.Context, authKeyID [8]byte) error
RecordNewMessage(ctx context.Context, authKeyID [8]byte, userID int64, msg domain.Message) (domain.UpdateEvent, domain.UpdateState, error)
PublishNewMessage(ctx context.Context, userID int64, msg domain.Message) (domain.UpdateEvent, domain.UpdateState, error)
RecordStory(ctx context.Context, authKeyID [8]byte, userID int64, story domain.Story, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordStoryFanout(ctx context.Context, userID int64, story domain.Story) (domain.UpdateEvent, domain.UpdateState, error)
RecordReadStories(ctx context.Context, authKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
@ -674,35 +675,36 @@ type AIComposeService interface {
// Deps 按业务域注入服务接口。各域的 handler 注册见对应文件（auth.go / users.go / updates.go）。
type Deps struct {
Auth AuthService
Account AccountService
Privacy PrivacyService
Help HelpService
AICompose AIComposeService
Users UsersService
Updates UpdatesService
Contacts ContactsService
Dialogs DialogsService
Messages MessagesService
Stories StoriesService
Channels ChannelsService
Files FilesService
Bots BotsService
Polls PollsService
Phone PhoneService
GroupCalls GroupCallsService
SFU sfu.Service
TURN turnsrv.Service
LangPack LangPackService
Sessions SessionBinder
Inline store.InlineRegistryStore
Limiter RateLimiter
Metrics Metrics
SecretChats SecretChatService
Stars StarsService
Gifts GiftsService
Passkey PasskeyService
Themes ThemeService
Auth AuthService
Account AccountService
Privacy PrivacyService
Help HelpService
AICompose AIComposeService
Users UsersService
Updates UpdatesService
BootstrapUpdates store.BootstrapUpdateJobStore
Contacts ContactsService
Dialogs DialogsService
Messages MessagesService
Stories StoriesService
Channels ChannelsService
Files FilesService
Bots BotsService
Polls PollsService
Phone PhoneService
GroupCalls GroupCallsService
SFU sfu.Service
TURN turnsrv.Service
LangPack LangPackService
Sessions SessionBinder
Inline store.InlineRegistryStore
Limiter RateLimiter
Metrics Metrics
SecretChats SecretChatService
Stars StarsService
Gifts GiftsService
Passkey PasskeyService
Themes ThemeService
}
// ThemeService 抽象自定义云主题(app/themes):创建/更新/查询主题 + 维护每用户已安装列表。

13
internal/rpc/rpc_testkit_updates_test.go
View file

@ -71,6 +71,19 @@ func (s *captureUpdates) RecordNewMessage(_ context.Context, authKeyID [8]byte,
return event, s.state, nil
}
func (s *captureUpdates) PublishNewMessage(_ context.Context, userID int64, msg domain.Message) (domain.UpdateEvent, domain.UpdateState, error) {
s.userID = userID
s.date = msg.Date
event := domain.UpdateEvent{Type: domain.UpdateEventNewMessage, Pts: s.state.Pts + 1, PtsCount: 1, Date: msg.Date, Message: msg}
s.events = append(s.events, event)
st := s.state
st.Pts = event.Pts
if st.Date == 0 {
st.Date = msg.Date
}
return event, st, nil
}
func (s *captureUpdates) RecordStory(_ context.Context, authKeyID [8]byte, userID int64, story domain.Story, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
s.excludeSessionID = excludeSessionID
return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{

3
internal/rpc/updates.go
View file

@ -34,6 +34,7 @@ func (r *Router) onUpdatesGetState(ctx context.Context) (*tg.UpdatesState, error
return nil, internalErr()
}
r.markSessionReceivesUpdates(ctx, userID)
r.registerBootstrapAfterBaseline(ctx, userID)
// 密聊 qts 是设备级、独立于账号级 pts 引擎：注入当前设备已分配的 qts（无密聊设备为 0）。
out := tgUpdateState(st)
out.Qts = r.deviceEncryptedQts(ctx)
@ -76,6 +77,7 @@ func (r *Router) onUpdatesGetDifference(ctx context.Context, req *tg.UpdatesGetD
// 密聊握手/已读状态事件（无 qts）：按未投递标记补回 OtherUpdates。
stateUpdates, statePeerUserIDs, stateEventIDs := r.encryptedStateUpdates(ctx, userID)
if len(st.Events) == 0 && len(st.ChannelNudges) == 0 && len(encMsgs) == 0 && len(stateUpdates) == 0 {
r.registerBootstrapAfterBaseline(ctx, userID)
return &tg.UpdatesDifferenceEmpty{Date: st.State.Date, Seq: st.State.Seq}, nil
}
st.Events = r.enrichUpdateEvents(ctx, userID, st.Events)
@ -87,6 +89,7 @@ func (r *Router) onUpdatesGetDifference(ctx context.Context, req *tg.UpdatesGetD
_ = r.deps.SecretChats.MarkStateEventsDelivered(ctx, deviceKey, stateEventIDs)
}
}
r.registerBootstrapAfterBaseline(ctx, userID)
return diff, nil
}

109
internal/rpc/updates_rpc_test.go
View file

@ -7,19 +7,31 @@ import (
"time"
"github.com/gotd/td/clock"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"go.uber.org/zap/zaptest"
appchannels "telesrv/internal/app/channels"
appupdates "telesrv/internal/app/updates"
"telesrv/internal/domain"
"telesrv/internal/postresponse"
"telesrv/internal/store/memory"
)
func TestPushLoginMessageSendsUpdateNewMessage(t *testing.T) {
sessions := &captureSessions{}
r := New(Config{}, Deps{Sessions: sessions}, zaptest.NewLogger(t), clock.System)
func TestBootstrapLoginMessagePublishesNewMessageAfterReady(t *testing.T) {
bootstrap := memory.NewBootstrapUpdateJobStore()
updates := &captureUpdates{state: domain.UpdateState{Pts: 3, Date: 1700000000}}
messages := &captureMessages{
list: domain.MessageList{Messages: []domain.Message{{
ID: 99,
OwnerUserID: 1000000001,
Peer: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID},
From: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID},
Date: 1700000100,
Body: "Login code: 12345",
Entities: []domain.MessageEntity{{Type: domain.MessageEntityBold, Offset: 0, Length: 11}},
}}},
}
r := New(Config{}, Deps{BootstrapUpdates: bootstrap, Updates: updates, Messages: messages}, zaptest.NewLogger(t), clock.System)
msg := domain.Message{
ID: 99,
OwnerUserID: 1000000001,
@ -29,39 +41,78 @@ func TestPushLoginMessageSendsUpdateNewMessage(t *testing.T) {
Body: "Login code: 12345",
Entities: []domain.MessageEntity{{Type: domain.MessageEntityBold, Offset: 0, Length: 11}},
}
event := domain.UpdateEvent{Type: domain.UpdateEventNewMessage, Pts: 4, PtsCount: 1, Date: msg.Date, Message: msg}
state := domain.UpdateState{Pts: 4, Date: 1700000100, Seq: 2}
authKeyID := [8]byte{1, 2, 3}
sessionID := int64(55)
ctx := WithSessionID(WithAuthKeyID(context.Background(), authKeyID), sessionID)
r.enqueueLoginMessageBootstrap(ctx, msg)
claimed := r.publishReadyBootstrapUpdates(context.Background(), 10, time.Second, zaptest.NewLogger(t))
if claimed != 0 || len(updates.events) != 0 {
t.Fatalf("published before ready = claimed %d events %d, want none", claimed, len(updates.events))
}
r.pushLoginMessage(context.Background(), [8]byte{}, 55, event, state)
ready, err := bootstrap.MarkReadyForSession(context.Background(), msg.OwnerUserID, authKeyID, sessionID)
if err != nil {
t.Fatalf("mark ready: %v", err)
}
if ready != 1 {
t.Fatalf("ready jobs = %d, want 1", ready)
}
claimed = r.publishReadyBootstrapUpdates(context.Background(), 10, time.Second, zaptest.NewLogger(t))
if claimed != 1 {
t.Fatalf("claimed jobs = %d, want 1", claimed)
}
if len(updates.events) != 1 {
t.Fatalf("published events = %d, want 1", len(updates.events))
}
event := updates.events[0]
if event.Type != domain.UpdateEventNewMessage || event.Message.ID != msg.ID || event.Pts != 4 {
t.Fatalf("event = %+v, want login message pts 4", event)
}
}
gotSession := sessions.snapshot()
if gotSession.sessionID != 55 || gotSession.messageType != proto.MessageFromServer {
t.Fatalf("push target = session %d type %v, want session 55 MessageFromServer", gotSession.sessionID, gotSession.messageType)
func TestUpdatesGetStatePublishesBootstrapAfterRPCResult(t *testing.T) {
bootstrap := memory.NewBootstrapUpdateJobStore()
updates := &captureUpdates{state: domain.UpdateState{Pts: 0, Date: 1700000000}}
msg := domain.Message{
ID: 1,
OwnerUserID: 1780243001,
Peer: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID},
From: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID},
Date: 1700000100,
Body: "Login code: 12345",
}
got, ok := gotSession.message.(*tg.Updates)
if !ok {
t.Fatalf("pushed message = %T, want *tg.Updates", gotSession.message)
messages := &captureMessages{list: domain.MessageList{Messages: []domain.Message{msg}}}
r := New(Config{}, Deps{BootstrapUpdates: bootstrap, Updates: updates, Messages: messages}, zaptest.NewLogger(t), clock.System)
authKeyID := [8]byte{9, 8, 7}
sessionID := int64(5723482677041206318)
ctx := postresponse.WithCallbacks(
WithUserID(
WithSessionID(WithAuthKeyID(context.Background(), authKeyID), sessionID),
msg.OwnerUserID,
),
)
r.enqueueLoginMessageBootstrap(ctx, msg)
state, err := r.onUpdatesGetState(ctx)
if err != nil {
t.Fatalf("getState: %v", err)
}
if len(got.Updates) != 1 || len(got.Users) != 1 {
t.Fatalf("updates payload = %+v, want one update and official user", got)
if state.Pts != 0 {
t.Fatalf("state pts = %d, want 0 before bootstrap publish", state.Pts)
}
update, ok := got.Updates[0].(*tg.UpdateNewMessage)
if !ok {
t.Fatalf("update = %T, want *tg.UpdateNewMessage", got.Updates[0])
if claimed := r.publishReadyBootstrapUpdates(context.Background(), 10, time.Second, zaptest.NewLogger(t)); claimed != 0 {
t.Fatalf("bootstrap published before post-response callback: %d", claimed)
}
if update.Pts != event.Pts || update.PtsCount != event.PtsCount || got.Seq != state.Seq {
t.Fatalf("update state = pts %d pts_count %d seq %d, want %d/%d/%d", update.Pts, update.PtsCount, got.Seq, event.Pts, event.PtsCount, state.Seq)
if len(updates.events) != 0 {
t.Fatalf("events before post-response = %d, want 0", len(updates.events))
}
message, ok := update.Message.(*tg.Message)
if !ok {
t.Fatalf("update message = %T, want *tg.Message", update.Message)
postresponse.Run(ctx)
if len(updates.events) != 1 {
t.Fatalf("events after post-response = %d, want 1", len(updates.events))
}
if message.ID != msg.ID || message.Message != msg.Body {
t.Fatalf("message = %+v, want login message %+v", message, msg)
}
user, ok := got.Users[0].(*tg.User)
if !ok || user.ID != domain.OfficialSystemUserID || !user.Verified || !user.Support {
t.Fatalf("user = %#v, want verified official system user", got.Users[0])
if got := updates.events[0]; got.Type != domain.UpdateEventNewMessage || got.Message.ID != msg.ID {
t.Fatalf("event after post-response = %+v, want login message", got)
}
}