merged from gramsrv upstream
This commit is contained in:
onysd
parent
79c64ee916
commit
21a0856587
651 changed files with 54774 additions and 4590 deletions
232
internal/loadharness/rich_state.go
Normal file
232
internal/loadharness/rich_state.go
Normal file
|
|
@ -0,0 +1,232 @@
|
|||
package loadharness
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rsa"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/iamxvbaba/td/tg"
|
||||
)
|
||||
|
||||
type richAccountStatePlan struct {
|
||||
PinnedPeerAccount int
|
||||
ReadPeerAccount int
|
||||
ReadGroupPosition int
|
||||
DraftMarker string
|
||||
}
|
||||
|
||||
func planRichAccountState(dataset *Dataset, account int) (richAccountStatePlan, error) {
|
||||
if dataset == nil || account < 0 || account >= dataset.Config.Accounts {
|
||||
return richAccountStatePlan{}, errors.New("invalid rich-state account")
|
||||
}
|
||||
if dataset.Config.PrivateFanout < 1 {
|
||||
return richAccountStatePlan{}, errors.New("rich startup dataset requires at least one private peer per account")
|
||||
}
|
||||
groupPosition := -1
|
||||
for position, group := range dataset.Groups {
|
||||
if group.HistoryMessages > 0 && datasetGroupHasAccount(group, account) {
|
||||
groupPosition = position
|
||||
break
|
||||
}
|
||||
}
|
||||
if groupPosition < 0 {
|
||||
return richAccountStatePlan{}, errors.New("rich startup dataset requires a non-empty supergroup per account")
|
||||
}
|
||||
return richAccountStatePlan{
|
||||
PinnedPeerAccount: (account + 1) % dataset.Config.Accounts,
|
||||
ReadPeerAccount: (account - 1 + dataset.Config.Accounts) % dataset.Config.Accounts,
|
||||
ReadGroupPosition: groupPosition,
|
||||
DraftMarker: fmt.Sprintf("[%s draft account %04d]", dataset.RunID, account),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// seedRichAccountState creates synchronized dialog state exclusively through
|
||||
// public MTProto RPCs. Every mutation is absolute and then read back through
|
||||
// messages.getPeerDialogs before the resumable journal is committed.
|
||||
func seedRichAccountState(
|
||||
ctx context.Context,
|
||||
cfg SeedConfig,
|
||||
manifest *Manifest,
|
||||
dataset *Dataset,
|
||||
journal *seedJournal,
|
||||
targets []SessionRecord,
|
||||
key [32]byte,
|
||||
publicKey *rsa.PublicKey,
|
||||
account int,
|
||||
) error {
|
||||
if journal.richStateComplete(account) {
|
||||
return nil
|
||||
}
|
||||
plan, err := planRichAccountState(dataset, account)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return withAuthorizedSeedSession(ctx, cfg, manifest, targets[account], key, publicKey, func(ctx context.Context, raw *tg.Client) error {
|
||||
pinnedTarget := targets[plan.PinnedPeerAccount]
|
||||
readTarget := targets[plan.ReadPeerAccount]
|
||||
pinnedPeer := &tg.InputPeerUser{UserID: pinnedTarget.UserID, AccessHash: pinnedTarget.AccessHash}
|
||||
readPeer := &tg.InputPeerUser{UserID: readTarget.UserID, AccessHash: readTarget.AccessHash}
|
||||
groupIdentity := journal.group(plan.ReadGroupPosition)
|
||||
if groupIdentity.ChannelID <= 0 || groupIdentity.AccessHash == 0 {
|
||||
return errors.New("rich-state group identity is incomplete")
|
||||
}
|
||||
channelPeer := &tg.InputPeerChannel{ChannelID: groupIdentity.ChannelID, AccessHash: groupIdentity.AccessHash}
|
||||
|
||||
before, err := getRichStateDialogs(ctx, cfg.OperationTimeout, raw, pinnedPeer, readPeer, channelPeer)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read rich-state cursors: %w", err)
|
||||
}
|
||||
readPrivate, ok := before[clientPeerKey{typ: "user", id: readTarget.UserID}]
|
||||
if !ok || readPrivate.TopMessage <= 0 {
|
||||
return errors.New("read private peer omitted its top message")
|
||||
}
|
||||
readChannel, ok := before[clientPeerKey{typ: "channel", id: groupIdentity.ChannelID}]
|
||||
if !ok || readChannel.TopMessage <= 0 {
|
||||
return errors.New("read channel omitted its top message")
|
||||
}
|
||||
|
||||
pinned, err := rpcWithFloodWaitRetry(ctx, cfg.OperationTimeout, func(rpcCtx context.Context) (bool, error) {
|
||||
return raw.MessagesToggleDialogPin(rpcCtx, &tg.MessagesToggleDialogPinRequest{
|
||||
Pinned: true, Peer: &tg.InputDialogPeer{Peer: pinnedPeer},
|
||||
})
|
||||
})
|
||||
if err != nil || !pinned {
|
||||
return rpcBooleanError("messages.toggleDialogPin", pinned, err)
|
||||
}
|
||||
saved, err := rpcWithFloodWaitRetry(ctx, cfg.OperationTimeout, func(rpcCtx context.Context) (bool, error) {
|
||||
return raw.MessagesSaveDraft(rpcCtx, &tg.MessagesSaveDraftRequest{Peer: pinnedPeer, Message: plan.DraftMarker})
|
||||
})
|
||||
if err != nil || !saved {
|
||||
return rpcBooleanError("messages.saveDraft", saved, err)
|
||||
}
|
||||
if _, err := rpcWithFloodWaitRetry(ctx, cfg.OperationTimeout, func(rpcCtx context.Context) (*tg.MessagesAffectedMessages, error) {
|
||||
return raw.MessagesReadHistory(rpcCtx, &tg.MessagesReadHistoryRequest{Peer: readPeer, MaxID: readPrivate.TopMessage})
|
||||
}); err != nil {
|
||||
return fmt.Errorf("messages.readHistory: %w", err)
|
||||
}
|
||||
channelRead, err := rpcWithFloodWaitRetry(ctx, cfg.OperationTimeout, func(rpcCtx context.Context) (bool, error) {
|
||||
return raw.ChannelsReadHistory(rpcCtx, &tg.ChannelsReadHistoryRequest{
|
||||
Channel: &tg.InputChannel{ChannelID: groupIdentity.ChannelID, AccessHash: groupIdentity.AccessHash},
|
||||
MaxID: readChannel.TopMessage,
|
||||
})
|
||||
})
|
||||
if err != nil || !channelRead {
|
||||
return rpcBooleanError("channels.readHistory", channelRead, err)
|
||||
}
|
||||
|
||||
after, err := getRichStateDialogs(ctx, cfg.OperationTimeout, raw, pinnedPeer, readPeer, channelPeer)
|
||||
if err != nil {
|
||||
return fmt.Errorf("verify rich-state dialogs: %w", err)
|
||||
}
|
||||
pinnedDialog, ok := after[clientPeerKey{typ: "user", id: pinnedTarget.UserID}]
|
||||
if !ok || !pinnedDialog.Pinned || !pinnedDialog.HasDraft || pinnedDialog.DraftText != plan.DraftMarker {
|
||||
return errors.New("pinned private dialog or exact draft was not persisted")
|
||||
}
|
||||
readPrivate = after[clientPeerKey{typ: "user", id: readTarget.UserID}]
|
||||
if readPrivate.ReadInboxMaxID < before[clientPeerKey{typ: "user", id: readTarget.UserID}].TopMessage || readPrivate.UnreadCount != 0 {
|
||||
return errors.New("private read boundary did not converge")
|
||||
}
|
||||
readChannel = after[clientPeerKey{typ: "channel", id: groupIdentity.ChannelID}]
|
||||
if readChannel.ReadInboxMaxID < before[clientPeerKey{typ: "channel", id: groupIdentity.ChannelID}].TopMessage || readChannel.UnreadCount != 0 {
|
||||
return errors.New("channel read boundary did not converge")
|
||||
}
|
||||
return journal.setRichStateComplete(account)
|
||||
})
|
||||
}
|
||||
|
||||
func rpcBooleanError(operation string, result bool, err error) error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("%s result=%v: %w", operation, result, err)
|
||||
}
|
||||
return fmt.Errorf("%s returned false", operation)
|
||||
}
|
||||
|
||||
func getRichStateDialogs(
|
||||
ctx context.Context,
|
||||
timeout time.Duration,
|
||||
raw *tg.Client,
|
||||
peers ...tg.InputPeerClass,
|
||||
) (map[clientPeerKey]ClientDialogState, error) {
|
||||
requests := make([]tg.InputDialogPeerClass, 0, len(peers))
|
||||
seen := make(map[clientPeerKey]struct{}, len(peers))
|
||||
for _, peer := range peers {
|
||||
key, ok := clientPeerFromInput(peer)
|
||||
if !ok {
|
||||
return nil, errors.New("invalid rich-state input peer")
|
||||
}
|
||||
if _, exists := seen[key]; exists {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
requests = append(requests, &tg.InputDialogPeer{Peer: peer})
|
||||
}
|
||||
response, err := rpcWithFloodWaitRetry(ctx, timeout, func(rpcCtx context.Context) (*tg.MessagesPeerDialogs, error) {
|
||||
return raw.MessagesGetPeerDialogs(rpcCtx, requests)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dialogs := make(map[clientPeerKey]ClientDialogState, len(response.Dialogs))
|
||||
if _, err := mergeDialogPage(dialogs, response.Dialogs, response.Messages, response.Chats, response.Users, false); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(dialogs) != len(requests) {
|
||||
return nil, fmt.Errorf("messages.getPeerDialogs returned %d/%d dialogs", len(dialogs), len(requests))
|
||||
}
|
||||
return dialogs, nil
|
||||
}
|
||||
|
||||
func clientPeerFromInput(peer tg.InputPeerClass) (clientPeerKey, bool) {
|
||||
switch value := peer.(type) {
|
||||
case *tg.InputPeerUser:
|
||||
return clientPeerKey{typ: "user", id: value.UserID}, value.UserID > 0 && value.AccessHash != 0
|
||||
case *tg.InputPeerChannel:
|
||||
return clientPeerKey{typ: "channel", id: value.ChannelID}, value.ChannelID > 0 && value.AccessHash != 0
|
||||
default:
|
||||
return clientPeerKey{}, false
|
||||
}
|
||||
}
|
||||
|
||||
func validateSeededRichDialogs(
|
||||
dataset *Dataset,
|
||||
seedState *DatasetSeedState,
|
||||
targets []SessionRecord,
|
||||
account int,
|
||||
dialogs []ClientDialogState,
|
||||
requireReadBoundaries bool,
|
||||
) error {
|
||||
if len(seedState.RichStateByAccount) == 0 {
|
||||
return nil
|
||||
}
|
||||
if len(seedState.RichStateByAccount) != dataset.Config.Accounts || !seedState.RichStateByAccount[account] {
|
||||
return errors.New("rich-state seed is incomplete")
|
||||
}
|
||||
plan, err := planRichAccountState(dataset, account)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
byPeer := make(map[clientPeerKey]ClientDialogState, len(dialogs))
|
||||
for _, dialog := range dialogs {
|
||||
byPeer[clientPeerKey{typ: dialog.PeerType, id: dialog.PeerID}] = dialog
|
||||
}
|
||||
pinned := byPeer[clientPeerKey{typ: "user", id: targets[plan.PinnedPeerAccount].UserID}]
|
||||
if !pinned.Pinned || !pinned.HasDraft || pinned.DraftText != plan.DraftMarker {
|
||||
return fmt.Errorf("seeded pinned dialog state mismatch: present=%v pinned=%v has_draft=%v draft_matches=%v got_draft=%q want_draft=%q",
|
||||
pinned.PeerID != 0, pinned.Pinned, pinned.HasDraft, pinned.DraftText == plan.DraftMarker, pinned.DraftText, plan.DraftMarker)
|
||||
}
|
||||
if !requireReadBoundaries {
|
||||
return nil
|
||||
}
|
||||
readPrivate := byPeer[clientPeerKey{typ: "user", id: targets[plan.ReadPeerAccount].UserID}]
|
||||
if readPrivate.TopMessage <= 0 || readPrivate.ReadInboxMaxID < readPrivate.TopMessage || readPrivate.UnreadCount != 0 {
|
||||
return errors.New("seeded private read boundary is stale")
|
||||
}
|
||||
group := seedState.Groups[plan.ReadGroupPosition]
|
||||
readChannel := byPeer[clientPeerKey{typ: "channel", id: group.ChannelID}]
|
||||
if readChannel.TopMessage <= 0 || readChannel.ReadInboxMaxID < readChannel.TopMessage || readChannel.UnreadCount != 0 {
|
||||
return errors.New("seeded channel read boundary is stale")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue