owpengram-server/internal/loadharness/seed.go
2026-09-01 12:06:31 +03:00

714 lines
22 KiB
Go

package loadharness
import (
"context"
"crypto/rsa"
"errors"
"fmt"
"sort"
"sync"
"time"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/tgerr"
)
const maxSeedInviteBatch = 200
type SeedConfig struct {
ManifestPath string
SessionKeyPath string
RSAKeyOverride string
DatasetPath string
SeedStatePath string
Concurrency int
OperationTimeout time.Duration
}
type SeedEvent struct {
Phase string
Completed int
Total int
Account int
Err error
}
type SeedResult struct {
PrivateMessages int
Groups int
InvitedMembers int
GroupMessages int
RichStateAccounts int
}
func (c SeedConfig) validate() error {
if c.ManifestPath == "" || c.SessionKeyPath == "" || c.DatasetPath == "" || c.SeedStatePath == "" {
return errors.New("manifest, session-key, dataset and seed-state paths are required")
}
if c.Concurrency <= 0 || c.Concurrency > 64 {
return errors.New("seed concurrency must be between 1 and 64")
}
if c.OperationTimeout <= 0 {
return errors.New("seed operation timeout must be positive")
}
return nil
}
// Seed materializes a Dataset only through authenticated MTProto RPCs. It never
// calls a server-internal handler or store. The journal is persisted before any
// non-idempotent channel operation, so an interrupted run can reconcile using
// the same account's public RPC view before it proceeds.
func Seed(ctx context.Context, cfg SeedConfig, progress func(SeedEvent)) (*SeedResult, error) {
if err := cfg.validate(); err != nil {
return nil, err
}
manifest, err := LoadManifest(cfg.ManifestPath)
if err != nil {
return nil, err
}
dataset, err := LoadDataset(cfg.DatasetPath)
if err != nil {
return nil, err
}
targets, err := seedPrimaryTargets(manifest, dataset.Config.Accounts)
if err != nil {
return nil, err
}
key, err := LoadSessionKey(cfg.SessionKeyPath)
if err != nil {
return nil, err
}
publicKey, err := loadManifestPublicKey(cfg.ManifestPath, manifest.Endpoint, cfg.RSAKeyOverride)
if err != nil {
return nil, err
}
state, err := LoadDatasetSeedState(cfg.SeedStatePath, dataset)
if err != nil {
return nil, err
}
journal := &seedJournal{path: cfg.SeedStatePath, dataset: dataset, state: state}
if err := journal.enableRichState(); err != nil {
return nil, err
}
if err := journal.persist(); err != nil {
return nil, err
}
accounts := make([]int, dataset.Config.Accounts)
for account := range accounts {
accounts[account] = account
}
if err := runSeedAccountPhase(ctx, "private", accounts, cfg.Concurrency, progress, func(ctx context.Context, account int) error {
return seedPrivateMessages(ctx, cfg, manifest, dataset, journal, targets, key, publicKey, account)
}); err != nil {
return nil, err
}
groupsByCreator := make(map[int][]int)
for position, group := range dataset.Groups {
groupsByCreator[group.CreatorAccount] = append(groupsByCreator[group.CreatorAccount], position)
}
creators := make([]int, 0, len(groupsByCreator))
for account := range groupsByCreator {
creators = append(creators, account)
}
sort.Ints(creators)
if err := runSeedAccountPhase(ctx, "groups", creators, cfg.Concurrency, progress, func(ctx context.Context, account int) error {
return withAuthorizedSeedSession(ctx, cfg, manifest, targets[account], key, publicKey, func(ctx context.Context, raw *tg.Client) error {
for _, position := range groupsByCreator[account] {
if err := seedGroup(ctx, cfg, dataset, journal, targets, raw, position); err != nil {
return fmt.Errorf("group %d: %w", dataset.Groups[position].Index, err)
}
}
return nil
})
}); err != nil {
return nil, err
}
historyTasks := datasetHistoryTasks(dataset)
if err := runSeedAccountPhase(ctx, "group-history", accounts, cfg.Concurrency, progress, func(ctx context.Context, account int) error {
return seedGroupHistory(ctx, cfg, manifest, dataset, journal, targets, key, publicKey, account, historyTasks[account])
}); err != nil {
return nil, err
}
if err := runSeedAccountPhase(ctx, "rich-state", accounts, cfg.Concurrency, progress, func(ctx context.Context, account int) error {
return seedRichAccountState(ctx, cfg, manifest, dataset, journal, targets, key, publicKey, account)
}); err != nil {
return nil, err
}
if err := journal.assertComplete(); err != nil {
return nil, err
}
result := &SeedResult{PrivateMessages: len(dataset.PrivateEdges), Groups: len(dataset.Groups), RichStateAccounts: dataset.Config.Accounts}
for _, group := range dataset.Groups {
result.InvitedMembers += len(group.MemberAccounts) - 1
result.GroupMessages += group.HistoryMessages
}
return result, nil
}
func seedPrimaryTargets(manifest *Manifest, accounts int) ([]SessionRecord, error) {
targets := primaryTargets(manifest.Sessions)
if len(targets) < accounts {
return nil, fmt.Errorf("dataset requires %d primary accounts, manifest has account range 0..%d", accounts, len(targets)-1)
}
targets = targets[:accounts]
for account, target := range targets {
if target.AccountIndex != account || target.UserID <= 0 || target.AccessHash == 0 || target.SessionFile == "" {
return nil, fmt.Errorf("manifest has no complete primary session for account %d", account)
}
}
return targets, nil
}
func seedPrivateMessages(
ctx context.Context,
cfg SeedConfig,
manifest *Manifest,
dataset *Dataset,
journal *seedJournal,
targets []SessionRecord,
key [32]byte,
publicKey *rsa.PublicKey,
account int,
) error {
cursor := journal.privateCursor(account)
if cursor == dataset.Config.PrivateFanout {
return nil
}
start := account * dataset.Config.PrivateFanout
edges := dataset.PrivateEdges[start : start+dataset.Config.PrivateFanout]
return withAuthorizedSeedSession(ctx, cfg, manifest, targets[account], key, publicKey, func(ctx context.Context, raw *tg.Client) error {
for i := cursor; i < len(edges); i++ {
edge := edges[i]
target := targets[edge.RecipientAccount]
_, err := rpcWithFloodWaitRetry(ctx, cfg.OperationTimeout, func(rpcCtx context.Context) (tg.UpdatesClass, error) {
return raw.MessagesSendMessage(rpcCtx, &tg.MessagesSendMessageRequest{
Peer: &tg.InputPeerUser{UserID: target.UserID, AccessHash: target.AccessHash},
Message: edge.Marker, RandomID: edge.RandomID,
})
})
if err != nil {
return fmt.Errorf("messages.sendMessage edge %d: %w", i, err)
}
}
return journal.setPrivateCursor(account, len(edges))
})
}
func seedGroup(
ctx context.Context,
cfg SeedConfig,
dataset *Dataset,
journal *seedJournal,
targets []SessionRecord,
raw *tg.Client,
position int,
) error {
group := dataset.Groups[position]
groupState := journal.group(position)
if groupState.ChannelID == 0 && groupState.CreatePending {
channel, found, err := reconcilePendingChannelCreate(ctx, cfg.OperationTimeout, raw, group)
if err != nil {
return err
}
if found {
if err := journal.commitChannel(position, channel.ID, channel.AccessHash); err != nil {
return err
}
} else if err := journal.clearCreatePending(position); err != nil {
return err
}
groupState = journal.group(position)
}
if groupState.ChannelID == 0 {
if err := journal.beginCreate(position); err != nil {
return err
}
updates, err := rpcWithFloodWaitRetry(ctx, cfg.OperationTimeout, func(rpcCtx context.Context) (tg.UpdatesClass, error) {
return raw.ChannelsCreateChannel(rpcCtx, &tg.ChannelsCreateChannelRequest{
Megagroup: true, Title: group.Title, About: group.About,
})
})
if err != nil {
return fmt.Errorf("channels.createChannel pending reconciliation: %w", err)
}
channel, err := createdChannelFromUpdates(updates, group.Title)
if err != nil {
return err
}
if err := journal.commitChannel(position, channel.ID, channel.AccessHash); err != nil {
return err
}
groupState = journal.group(position)
}
invitees := groupInvitees(group)
for groupState.InviteCursor < len(invitees) {
if groupState.InvitePendingEnd > groupState.InviteCursor {
if err := reconcilePendingInvites(ctx, cfg.OperationTimeout, raw, groupState, invitees, targets); err != nil {
return err
}
if err := journal.commitInvite(position, groupState.InvitePendingEnd); err != nil {
return err
}
groupState = journal.group(position)
continue
}
end := min(groupState.InviteCursor+maxSeedInviteBatch, len(invitees))
if err := journal.beginInvite(position, end); err != nil {
return err
}
if err := inviteAccounts(ctx, cfg.OperationTimeout, raw, groupState, invitees[groupState.InviteCursor:end], targets); err != nil {
return fmt.Errorf("channels.inviteToChannel pending reconciliation: %w", err)
}
if err := journal.commitInvite(position, end); err != nil {
return err
}
groupState = journal.group(position)
}
return nil
}
func reconcilePendingChannelCreate(ctx context.Context, timeout time.Duration, raw *tg.Client, group DatasetGroup) (*tg.Channel, bool, error) {
rpcCtx, cancel := context.WithTimeout(ctx, timeout)
dialogs, err := raw.MessagesGetDialogs(rpcCtx, &tg.MessagesGetDialogsRequest{
OffsetPeer: &tg.InputPeerEmpty{}, Limit: 500,
})
cancel()
if err != nil {
return nil, false, fmt.Errorf("reconcile create with messages.getDialogs: %w", err)
}
var chats []tg.ChatClass
switch value := dialogs.(type) {
case *tg.MessagesDialogs:
chats = value.Chats
case *tg.MessagesDialogsSlice:
chats = value.Chats
case *tg.MessagesDialogsNotModified:
return nil, false, errors.New("reconcile create unexpectedly returned dialogsNotModified")
default:
return nil, false, fmt.Errorf("reconcile create messages.getDialogs returned %T", dialogs)
}
matches := make([]*tg.Channel, 0, 1)
for _, chat := range chats {
channel, ok := chat.(*tg.Channel)
if !ok || channel.Title != group.Title || !channel.Megagroup {
continue
}
if channel.ID <= 0 || channel.AccessHash == 0 {
return nil, false, fmt.Errorf("reconciled channel %q has incomplete identity", group.Title)
}
matches = append(matches, channel)
}
if len(matches) > 1 {
return nil, false, fmt.Errorf("ambiguous create produced %d channels named %q", len(matches), group.Title)
}
if len(matches) == 0 {
return nil, false, nil
}
return matches[0], true, nil
}
func createdChannelFromUpdates(updates tg.UpdatesClass, title string) (*tg.Channel, error) {
var chats []tg.ChatClass
switch value := updates.(type) {
case *tg.Updates:
chats = value.Chats
case *tg.UpdatesCombined:
chats = value.Chats
default:
return nil, fmt.Errorf("channels.createChannel returned %T", updates)
}
for _, chat := range chats {
channel, ok := chat.(*tg.Channel)
if ok && channel.Title == title && channel.Megagroup && channel.ID > 0 && channel.AccessHash != 0 {
return channel, nil
}
}
return nil, fmt.Errorf("channels.createChannel response omitted supergroup %q", title)
}
func groupInvitees(group DatasetGroup) []int {
invitees := make([]int, 0, len(group.MemberAccounts)-1)
for _, account := range group.MemberAccounts {
if account != group.CreatorAccount {
invitees = append(invitees, account)
}
}
return invitees
}
func inviteAccounts(
ctx context.Context,
timeout time.Duration,
raw *tg.Client,
groupState DatasetSeedGroupState,
accounts []int,
targets []SessionRecord,
) error {
if len(accounts) == 0 || len(accounts) > maxSeedInviteBatch {
return fmt.Errorf("invalid invite batch size %d", len(accounts))
}
users := make([]tg.InputUserClass, 0, len(accounts))
for _, account := range accounts {
target := targets[account]
users = append(users, &tg.InputUser{UserID: target.UserID, AccessHash: target.AccessHash})
}
result, err := rpcWithFloodWaitRetry(ctx, timeout, func(rpcCtx context.Context) (*tg.MessagesInvitedUsers, error) {
return raw.ChannelsInviteToChannel(rpcCtx, &tg.ChannelsInviteToChannelRequest{
Channel: &tg.InputChannel{ChannelID: groupState.ChannelID, AccessHash: groupState.AccessHash},
Users: users,
})
})
if err != nil {
return err
}
if len(result.MissingInvitees) != 0 {
return fmt.Errorf("server reported %d missing_invitees", len(result.MissingInvitees))
}
return nil
}
func reconcilePendingInvites(
ctx context.Context,
timeout time.Duration,
raw *tg.Client,
groupState DatasetSeedGroupState,
invitees []int,
targets []SessionRecord,
) error {
pending := invitees[groupState.InviteCursor:groupState.InvitePendingEnd]
missing := make([]int, 0, len(pending))
for _, account := range pending {
target := targets[account]
rpcCtx, cancel := context.WithTimeout(ctx, timeout)
_, err := raw.ChannelsGetParticipant(rpcCtx, &tg.ChannelsGetParticipantRequest{
Channel: &tg.InputChannel{ChannelID: groupState.ChannelID, AccessHash: groupState.AccessHash},
Participant: &tg.InputPeerUser{UserID: target.UserID, AccessHash: target.AccessHash},
})
cancel()
switch {
case err == nil:
case tgerr.Is(err, "USER_NOT_PARTICIPANT"):
missing = append(missing, account)
default:
return fmt.Errorf("channels.getParticipant account %d: %w", account, err)
}
}
if len(missing) == 0 {
return nil
}
return inviteAccounts(ctx, timeout, raw, groupState, missing, targets)
}
type datasetHistoryTask struct {
GroupPosition int
MessageIndex int
}
func datasetHistoryTasks(dataset *Dataset) [][]datasetHistoryTask {
tasks := make([][]datasetHistoryTask, dataset.Config.Accounts)
for position, group := range dataset.Groups {
for message := 0; message < group.HistoryMessages; message++ {
account := group.MemberAccounts[message%len(group.MemberAccounts)]
tasks[account] = append(tasks[account], datasetHistoryTask{GroupPosition: position, MessageIndex: message})
}
}
return tasks
}
func seedGroupHistory(
ctx context.Context,
cfg SeedConfig,
manifest *Manifest,
dataset *Dataset,
journal *seedJournal,
targets []SessionRecord,
key [32]byte,
publicKey *rsa.PublicKey,
account int,
tasks []datasetHistoryTask,
) error {
cursor := journal.historyCursor(account)
if cursor == len(tasks) {
return nil
}
return withAuthorizedSeedSession(ctx, cfg, manifest, targets[account], key, publicKey, func(ctx context.Context, raw *tg.Client) error {
for i := cursor; i < len(tasks); i++ {
task := tasks[i]
group := dataset.Groups[task.GroupPosition]
groupState := journal.group(task.GroupPosition)
if groupState.ChannelID == 0 || groupState.InviteCursor != len(group.MemberAccounts)-1 {
return fmt.Errorf("group %d is not fully seeded", group.Index)
}
_, err := rpcWithFloodWaitRetry(ctx, cfg.OperationTimeout, func(rpcCtx context.Context) (tg.UpdatesClass, error) {
return raw.MessagesSendMessage(rpcCtx, &tg.MessagesSendMessageRequest{
Peer: &tg.InputPeerChannel{ChannelID: groupState.ChannelID, AccessHash: groupState.AccessHash},
Message: fmt.Sprintf("[%s group %04d message %04d sender %04d]", dataset.RunID, group.Index, task.MessageIndex+1, account),
RandomID: stableDatasetID(dataset.Config.Seed, "group-history", dataset.Config.Accounts, group.Index, task.MessageIndex),
})
})
if err != nil {
return fmt.Errorf("messages.sendMessage group %d message %d: %w", group.Index, task.MessageIndex, err)
}
}
return journal.setHistoryCursor(account, len(tasks))
})
}
func withAuthorizedSeedSession(
ctx context.Context,
cfg SeedConfig,
manifest *Manifest,
record SessionRecord,
key [32]byte,
publicKey *rsa.PublicKey,
work func(context.Context, *tg.Client) error,
) error {
storage := &EncryptedFileStorage{Path: resolveSessionPath(cfg.ManifestPath, record), Key: key}
client, err := newClient(manifest.Endpoint, publicKey, storage, clientHooks{})
if err != nil {
return err
}
return client.Run(ctx, func(ctx context.Context) error {
statusCtx, cancel := context.WithTimeout(ctx, cfg.OperationTimeout)
status, err := client.Auth().Status(statusCtx)
cancel()
if err != nil {
return fmt.Errorf("authorization status: %w", err)
}
if !status.Authorized || status.User == nil || status.User.ID != record.UserID {
return fmt.Errorf("session account %d is not authorized as expected user", record.AccountIndex)
}
return work(ctx, tg.NewClient(client))
})
}
type seedAccountResult struct {
account int
err error
}
func runSeedAccountPhase(
ctx context.Context,
phase string,
accounts []int,
concurrency int,
progress func(SeedEvent),
work func(context.Context, int) error,
) error {
phaseCtx, cancel := context.WithCancel(ctx)
defer cancel()
jobs := make(chan int)
results := make(chan seedAccountResult, len(accounts))
workers := min(concurrency, len(accounts))
var wg sync.WaitGroup
for i := 0; i < workers; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for account := range jobs {
err := work(phaseCtx, account)
results <- seedAccountResult{account: account, err: err}
if err != nil {
return
}
}
}()
}
go func() {
defer close(jobs)
for _, account := range accounts {
select {
case jobs <- account:
case <-phaseCtx.Done():
return
}
}
}()
go func() {
wg.Wait()
close(results)
}()
completed := 0
var firstErr error
for result := range results {
if result.err == nil {
completed++
} else if firstErr == nil {
firstErr = fmt.Errorf("%s account %d: %w", phase, result.account, result.err)
cancel()
}
if progress != nil {
progress(SeedEvent{Phase: phase, Completed: completed, Total: len(accounts), Account: result.account, Err: result.err})
}
}
if firstErr != nil {
return firstErr
}
if completed != len(accounts) {
if ctx.Err() != nil {
return ctx.Err()
}
return fmt.Errorf("%s completed %d/%d accounts", phase, completed, len(accounts))
}
return nil
}
type seedJournal struct {
mu sync.Mutex
path string
dataset *Dataset
state *DatasetSeedState
}
func (j *seedJournal) persist() error {
j.mu.Lock()
defer j.mu.Unlock()
return WriteDatasetSeedState(j.path, j.dataset, j.state)
}
func (j *seedJournal) privateCursor(account int) int {
j.mu.Lock()
defer j.mu.Unlock()
return j.state.PrivateSentByAccount[account]
}
func (j *seedJournal) historyCursor(account int) int {
j.mu.Lock()
defer j.mu.Unlock()
return j.state.HistorySentByAccount[account]
}
func (j *seedJournal) richStateComplete(account int) bool {
j.mu.Lock()
defer j.mu.Unlock()
return len(j.state.RichStateByAccount) == j.dataset.Config.Accounts && j.state.RichStateByAccount[account]
}
func (j *seedJournal) enableRichState() error {
j.mu.Lock()
defer j.mu.Unlock()
if len(j.state.RichStateByAccount) == j.dataset.Config.Accounts {
return nil
}
if len(j.state.RichStateByAccount) != 0 {
return errors.New("seed rich-state journal has invalid dimensions")
}
j.state.RichStateByAccount = make([]bool, j.dataset.Config.Accounts)
return WriteDatasetSeedState(j.path, j.dataset, j.state)
}
func (j *seedJournal) group(position int) DatasetSeedGroupState {
j.mu.Lock()
defer j.mu.Unlock()
return j.state.Groups[position]
}
func (j *seedJournal) setPrivateCursor(account, cursor int) error {
j.mu.Lock()
defer j.mu.Unlock()
old := j.state.PrivateSentByAccount[account]
j.state.PrivateSentByAccount[account] = cursor
if err := WriteDatasetSeedState(j.path, j.dataset, j.state); err != nil {
j.state.PrivateSentByAccount[account] = old
return err
}
return nil
}
func (j *seedJournal) setHistoryCursor(account, cursor int) error {
j.mu.Lock()
defer j.mu.Unlock()
old := j.state.HistorySentByAccount[account]
j.state.HistorySentByAccount[account] = cursor
if err := WriteDatasetSeedState(j.path, j.dataset, j.state); err != nil {
j.state.HistorySentByAccount[account] = old
return err
}
return nil
}
func (j *seedJournal) setRichStateComplete(account int) error {
j.mu.Lock()
defer j.mu.Unlock()
if len(j.state.RichStateByAccount) != j.dataset.Config.Accounts {
return errors.New("seed rich-state journal is not enabled")
}
old := j.state.RichStateByAccount[account]
j.state.RichStateByAccount[account] = true
if err := WriteDatasetSeedState(j.path, j.dataset, j.state); err != nil {
j.state.RichStateByAccount[account] = old
return err
}
return nil
}
func (j *seedJournal) updateGroup(position int, update func(*DatasetSeedGroupState)) error {
j.mu.Lock()
defer j.mu.Unlock()
old := j.state.Groups[position]
update(&j.state.Groups[position])
if err := WriteDatasetSeedState(j.path, j.dataset, j.state); err != nil {
j.state.Groups[position] = old
return err
}
return nil
}
func (j *seedJournal) beginCreate(position int) error {
return j.updateGroup(position, func(state *DatasetSeedGroupState) { state.CreatePending = true })
}
func (j *seedJournal) clearCreatePending(position int) error {
return j.updateGroup(position, func(state *DatasetSeedGroupState) { state.CreatePending = false })
}
func (j *seedJournal) commitChannel(position int, channelID, accessHash int64) error {
if channelID <= 0 || accessHash == 0 {
return errors.New("cannot commit incomplete channel identity")
}
return j.updateGroup(position, func(state *DatasetSeedGroupState) {
state.ChannelID = channelID
state.AccessHash = accessHash
state.CreatePending = false
})
}
func (j *seedJournal) beginInvite(position, end int) error {
return j.updateGroup(position, func(state *DatasetSeedGroupState) { state.InvitePendingEnd = end })
}
func (j *seedJournal) commitInvite(position, end int) error {
return j.updateGroup(position, func(state *DatasetSeedGroupState) {
state.InviteCursor = end
state.InvitePendingEnd = end
})
}
func (j *seedJournal) assertComplete() error {
j.mu.Lock()
defer j.mu.Unlock()
if err := j.state.Validate(j.dataset); err != nil {
return err
}
historyCounts := datasetHistoryTaskCounts(j.dataset)
for account := 0; account < j.dataset.Config.Accounts; account++ {
if j.state.PrivateSentByAccount[account] != j.dataset.Config.PrivateFanout || j.state.HistorySentByAccount[account] != historyCounts[account] {
return fmt.Errorf("account %d seed is incomplete", account)
}
if len(j.state.RichStateByAccount) != 0 && !j.state.RichStateByAccount[account] {
return fmt.Errorf("account %d rich state is incomplete", account)
}
}
for position, group := range j.dataset.Groups {
state := j.state.Groups[position]
if state.ChannelID == 0 || state.CreatePending || state.InviteCursor != len(group.MemberAccounts)-1 || state.InvitePendingEnd != state.InviteCursor {
return fmt.Errorf("group %d seed is incomplete", group.Index)
}
}
return nil
}