714 lines
22 KiB
Go
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
|
|
}
|