im_server/loadtest/im-pressure/prepare.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net/http"
	"sort"
	"sync"
	"sync/atomic"
)

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
	}
	manifest := newManifest(cfg)
	if existing, loadErr := loadManifest(manifestPath); loadErr == nil {
		manifest = existing
		log.Printf("resuming manifest with %d accounts", len(manifest.Accounts))
	}
	if err = applyConfigEndpoints(cfg, manifest); err != nil {
		return err
	}
	if err = prepareAccounts(ctx, cfg, manifest, apiKey, password); err != nil {
		return err
	}
	if err = saveManifest(manifestPath, manifest); err != nil {
		return err
	}
	api := newAPIClient(cfg, apiKey)
	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); 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, apiKey, password 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 {
		client, wasCreated, err := connectForPrepare(ctx, cfg, account, apiKey, password)
		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 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 {
		b := *accounts[accountKey("b", a.Index)]
		if a.UserID == "" || b.UserID == "" {
			return fmt.Errorf("pair %d has missing user ID", a.Index)
		}
		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) error {
	if err := parallelAccounts(ctx, manifest.Accounts, cfg.PrepareWorkers, func(account AccountRecord) 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
		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)
		}
		response, err := api.inviteGroup(ctx, owner, manifest.GroupTopic, batch, memberIDs[offset:end])
		if err != nil {
			return fmt.Errorf("invite group batch %d: %w", batch, err)
		}
		for _, item := range response.Items {
			if item.Status != "joined" && item.Status != "already_member" {
				return fmt.Errorf("invite user %s status=%s reason=%s", item.UserID, item.Status, item.Reason)
			}
		}
	}
	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 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()
	}
}