376 lines
12 KiB
Go
376 lines
12 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"sort"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
func runPrepare(ctx context.Context, cfg *Config, manifestPath string) error {
|
|
apiKey, err := cfg.apiKey()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
password, err := cfg.password()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
verificationCode, err := cfg.verificationCode()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
manifest := newManifest(cfg)
|
|
if existing, loadErr := loadManifest(manifestPath); loadErr == nil {
|
|
manifest, err = reconcileManifest(cfg, existing)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
log.Printf("resuming manifest with %d existing accounts, target=%d", len(existing.Accounts), len(manifest.Accounts))
|
|
}
|
|
if err = applyConfigEndpoints(cfg, manifest); err != nil {
|
|
return err
|
|
}
|
|
api := newAPIClient(cfg, apiKey)
|
|
if err = prepareAccounts(ctx, cfg, manifest, api, password, verificationCode); err != nil {
|
|
return err
|
|
}
|
|
if err = saveManifest(manifestPath, manifest); err != nil {
|
|
return err
|
|
}
|
|
if err = prepareFriendships(ctx, cfg, manifest, api); err != nil {
|
|
return err
|
|
}
|
|
if err = saveManifest(manifestPath, manifest); err != nil {
|
|
return err
|
|
}
|
|
if err = prepareGroup(ctx, cfg, manifest, api, manifestPath); err != nil {
|
|
return err
|
|
}
|
|
if err = saveManifest(manifestPath, manifest); err != nil {
|
|
return err
|
|
}
|
|
return runVerify(ctx, cfg, manifest, api)
|
|
}
|
|
|
|
func prepareAccounts(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient, password, verificationCode string) error {
|
|
var created atomic.Int64
|
|
var existing atomic.Int64
|
|
var mu sync.Mutex
|
|
err := parallelAccounts(ctx, manifest.Accounts, cfg.PrepareWorkers, func(account AccountRecord) error {
|
|
if account.UserID != "" && account.Token != "" && account.TokenExpires.After(time.Now().Add(5*time.Minute)) {
|
|
existing.Add(1)
|
|
return nil
|
|
}
|
|
client, wasCreated, err := connectForPrepareWithRetry(ctx, cfg, account, api, password, verificationCode)
|
|
if err != nil {
|
|
return fmt.Errorf("account %s: %w", account.Username, err)
|
|
}
|
|
updated := client.Account()
|
|
client.Close()
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
for i := range manifest.Accounts {
|
|
if manifest.Accounts[i].Machine == account.Machine && manifest.Accounts[i].Index == account.Index {
|
|
manifest.Accounts[i].UserID = updated.UserID
|
|
manifest.Accounts[i].Token = updated.Token
|
|
manifest.Accounts[i].TokenExpires = updated.TokenExpires
|
|
break
|
|
}
|
|
}
|
|
if wasCreated {
|
|
created.Add(1)
|
|
} else {
|
|
existing.Add(1)
|
|
}
|
|
return nil
|
|
})
|
|
if err == nil {
|
|
log.Printf("accounts ready=%d created=%d existing=%d", len(manifest.Accounts), created.Load(), existing.Load())
|
|
}
|
|
return err
|
|
}
|
|
|
|
func connectForPrepareWithRetry(ctx context.Context, cfg *Config, account AccountRecord, api *apiClient,
|
|
password, verificationCode string) (*WSClient, bool, error) {
|
|
var lastErr error
|
|
for attempt := 1; attempt <= 6; attempt++ {
|
|
client, created, err := connectForPrepare(ctx, cfg, account, api, password, verificationCode)
|
|
if err == nil {
|
|
return client, created, nil
|
|
}
|
|
lastErr = err
|
|
if ctx.Err() != nil {
|
|
return nil, false, ctx.Err()
|
|
}
|
|
delay := time.Duration(attempt*5) * time.Second
|
|
log.Printf("account=%s prepare attempt=%d failed; retry_in=%s: %v", account.Username, attempt, delay, err)
|
|
if err = waitContext(ctx, delay); err != nil {
|
|
return nil, false, err
|
|
}
|
|
}
|
|
return nil, false, fmt.Errorf("account %s prepare failed after retries: %w", account.Username, lastErr)
|
|
}
|
|
|
|
func prepareFriendships(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
|
|
accounts := manifestAccountMap(manifest)
|
|
pairs := make([]AccountRecord, 0, cfg.UsersPerMachine)
|
|
for i := 1; i <= cfg.UsersPerMachine; i++ {
|
|
pairs = append(pairs, *accounts[accountKey("a", i)])
|
|
}
|
|
var mu sync.Mutex
|
|
err := parallelAccounts(ctx, pairs, cfg.PrepareWorkers, func(a AccountRecord) error {
|
|
return retryPrepareOperation(ctx, fmt.Sprintf("friend-pair-%04d", a.Index), func() error {
|
|
b := *accounts[accountKey("b", a.Index)]
|
|
if a.UserID == "" || b.UserID == "" {
|
|
return fmt.Errorf("pair %d has missing user ID", a.Index)
|
|
}
|
|
conversation, conversationErr := api.ensureP2P(ctx, a, b)
|
|
if conversationErr == nil && conversation.Relationship == "friend" && conversation.Topic != "" {
|
|
mu.Lock()
|
|
accounts[accountKey("a", a.Index)].P2PTopic = conversation.Topic
|
|
accounts[accountKey("b", a.Index)].P2PTopic = conversation.Topic
|
|
mu.Unlock()
|
|
return nil
|
|
}
|
|
if conversationErr != nil {
|
|
var apiErr *apiError
|
|
if !errors.As(conversationErr, &apiErr) || apiErr.Status != http.StatusForbidden {
|
|
return fmt.Errorf("check friend pair %d: %w", a.Index, conversationErr)
|
|
}
|
|
}
|
|
request, requestErr := api.createFriendRequest(ctx, a, b)
|
|
if requestErr == nil && request.Status == "pending" {
|
|
if _, requestErr = api.acceptFriendRequest(ctx, b, request.RequestID); requestErr != nil {
|
|
return fmt.Errorf("accept pair %d: %w", a.Index, requestErr)
|
|
}
|
|
} else if requestErr != nil {
|
|
var apiErr *apiError
|
|
if !errors.As(requestErr, &apiErr) || apiErr.Status != http.StatusConflict {
|
|
return fmt.Errorf("friend pair %d: %w", a.Index, requestErr)
|
|
}
|
|
pendingID, findErr := api.findPendingFriendRequest(ctx, b, a.UserID)
|
|
if findErr != nil {
|
|
return fmt.Errorf("find pending friend pair %d: %w", a.Index, findErr)
|
|
}
|
|
if pendingID != "" {
|
|
if _, acceptErr := api.acceptFriendRequest(ctx, b, pendingID); acceptErr != nil {
|
|
return fmt.Errorf("resume friend pair %d: %w", a.Index, acceptErr)
|
|
}
|
|
}
|
|
}
|
|
conversation, err := api.ensureP2P(ctx, a, b)
|
|
if err != nil {
|
|
return fmt.Errorf("ensure P2P pair %d: %w", a.Index, err)
|
|
}
|
|
if conversation.Relationship != "friend" || conversation.Topic == "" {
|
|
return fmt.Errorf("pair %d relationship=%s topic=%q", a.Index, conversation.Relationship, conversation.Topic)
|
|
}
|
|
mu.Lock()
|
|
accounts[accountKey("a", a.Index)].P2PTopic = conversation.Topic
|
|
accounts[accountKey("b", a.Index)].P2PTopic = conversation.Topic
|
|
mu.Unlock()
|
|
return nil
|
|
})
|
|
})
|
|
if err == nil {
|
|
log.Printf("friendships ready=%d", cfg.UsersPerMachine)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func prepareGroup(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient, manifestPath string) error {
|
|
if err := parallelAccounts(ctx, manifest.Accounts, cfg.PrepareWorkers, func(account AccountRecord) error {
|
|
return retryPrepareOperation(ctx, "privacy-"+account.Username, func() error {
|
|
return api.disableGroupInviteVerification(ctx, account)
|
|
})
|
|
}); err != nil {
|
|
return fmt.Errorf("prepare group invite privacy: %w", err)
|
|
}
|
|
accounts := manifestAccountMap(manifest)
|
|
owner := *accounts[accountKey("a", 1)]
|
|
if manifest.GroupTopic == "" {
|
|
group, err := api.createGroup(ctx, owner)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
manifest.GroupTopic = group.Topic
|
|
if err = saveManifest(manifestPath, manifest); err != nil {
|
|
return fmt.Errorf("checkpoint created group: %w", err)
|
|
}
|
|
log.Printf("created load-test group topic_hash=%s", topicHash(group.Topic))
|
|
}
|
|
memberIDs := make([]string, 0, len(manifest.Accounts)-1)
|
|
for _, account := range manifest.Accounts {
|
|
if account.UserID != owner.UserID {
|
|
memberIDs = append(memberIDs, account.UserID)
|
|
}
|
|
}
|
|
sort.Strings(memberIDs)
|
|
for offset, batch := 0, 1; offset < len(memberIDs); offset, batch = offset+100, batch+1 {
|
|
end := offset + 100
|
|
if end > len(memberIDs) {
|
|
end = len(memberIDs)
|
|
}
|
|
if err := ensureGroupMembers(ctx, api, owner, manifest.GroupTopic, batch, memberIDs[offset:end]); err != nil {
|
|
return fmt.Errorf("invite group batch %d: %w", batch, err)
|
|
}
|
|
}
|
|
count, err := api.countGroupMembers(ctx, owner, manifest.GroupTopic)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if count != len(manifest.Accounts) {
|
|
return fmt.Errorf("group has %d members, expected %d", count, len(manifest.Accounts))
|
|
}
|
|
log.Printf("group ready members=%d topic_hash=%s", count, topicHash(manifest.GroupTopic))
|
|
return nil
|
|
}
|
|
|
|
func ensureGroupMembers(ctx context.Context, api *apiClient, owner AccountRecord, topic string, batch int, desired []string) error {
|
|
var lastErr error
|
|
for attempt := 1; attempt <= 12; attempt++ {
|
|
members, err := api.groupMemberIDs(ctx, owner, topic)
|
|
if err != nil {
|
|
lastErr = fmt.Errorf("list current members: %w", err)
|
|
} else {
|
|
missing := make([]string, 0, len(desired))
|
|
for _, userID := range desired {
|
|
if !members[userID] {
|
|
missing = append(missing, userID)
|
|
}
|
|
}
|
|
if len(missing) == 0 {
|
|
return nil
|
|
}
|
|
response, inviteErr := api.inviteGroup(ctx, owner, topic, batch, missing)
|
|
if inviteErr != nil {
|
|
lastErr = inviteErr
|
|
} else {
|
|
lastErr = nil
|
|
for _, item := range response.Items {
|
|
if item.Status != "joined" && item.Status != "already_member" {
|
|
lastErr = fmt.Errorf("invite user %s status=%s reason=%s", item.UserID, item.Status, item.Reason)
|
|
break
|
|
}
|
|
}
|
|
if lastErr == nil {
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
if attempt == 12 {
|
|
break
|
|
}
|
|
delay := time.Duration(attempt*5) * time.Second
|
|
log.Printf("operation=group-invite-%03d attempt=%d failed; retry_in=%s: %v", batch, attempt, delay, lastErr)
|
|
if err := waitContext(ctx, delay); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return fmt.Errorf("missing members not recovered after retries: %w", lastErr)
|
|
}
|
|
|
|
func retryPrepareOperation(ctx context.Context, label string, operation func() error) error {
|
|
return retryOperation(ctx, label, 6, operation)
|
|
}
|
|
|
|
func retryOperation(ctx context.Context, label string, attempts int, operation func() error) error {
|
|
var lastErr error
|
|
for attempt := 1; attempt <= attempts; attempt++ {
|
|
if err := operation(); err == nil {
|
|
return nil
|
|
} else {
|
|
lastErr = err
|
|
}
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
if attempt == attempts {
|
|
break
|
|
}
|
|
delay := time.Duration(attempt*5) * time.Second
|
|
log.Printf("operation=%s attempt=%d failed; retry_in=%s: %v", label, attempt, delay, lastErr)
|
|
if err := waitContext(ctx, delay); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return fmt.Errorf("operation %s failed after retries: %w", label, lastErr)
|
|
}
|
|
|
|
func runVerify(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
|
|
if len(manifest.Accounts) != cfg.UsersPerMachine*2 {
|
|
return fmt.Errorf("manifest accounts=%d expected=%d", len(manifest.Accounts), cfg.UsersPerMachine*2)
|
|
}
|
|
for _, account := range manifest.Accounts {
|
|
if account.UserID == "" || account.Token == "" || account.P2PTopic == "" {
|
|
return fmt.Errorf("account %s is incomplete", account.Username)
|
|
}
|
|
}
|
|
if manifest.GroupTopic == "" {
|
|
return errors.New("manifest group topic is empty")
|
|
}
|
|
owner := manifestAccountMap(manifest)[accountKey("a", 1)]
|
|
count, err := api.countGroupMembers(ctx, *owner, manifest.GroupTopic)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if count != len(manifest.Accounts) {
|
|
return fmt.Errorf("group member count=%d expected=%d", count, len(manifest.Accounts))
|
|
}
|
|
log.Printf("verify passed accounts=%d pairs=%d group_members=%d", len(manifest.Accounts), cfg.UsersPerMachine, count)
|
|
return nil
|
|
}
|
|
|
|
func parallelAccounts(ctx context.Context, accounts []AccountRecord, workers int, fn func(AccountRecord) error) error {
|
|
if workers <= 0 {
|
|
workers = 1
|
|
}
|
|
workCtx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
jobs := make(chan AccountRecord)
|
|
errCh := make(chan error, 1)
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < workers; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for account := range jobs {
|
|
if err := fn(account); err != nil {
|
|
select {
|
|
case errCh <- err:
|
|
default:
|
|
}
|
|
cancel()
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
go func() {
|
|
defer close(jobs)
|
|
for _, account := range accounts {
|
|
select {
|
|
case jobs <- account:
|
|
case <-workCtx.Done():
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
wg.Wait()
|
|
select {
|
|
case err := <-errCh:
|
|
return err
|
|
default:
|
|
return ctx.Err()
|
|
}
|
|
}
|