junhong_cmp_fiber/scripts/benchmark/generate_cards.go

// +build ignore

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

// IotCard 简化的卡模型
type IotCard struct {
	ID                uint       `gorm:"primaryKey"`
	ICCID             string     `gorm:"column:iccid;uniqueIndex:idx_iot_card_iccid,where:deleted_at IS NULL"`
	CardCategory      string     `gorm:"column:card_category;default:normal"`
	CarrierID         uint       `gorm:"column:carrier_id"`
	Status            int        `gorm:"column:status;default:1"`
	ActivationStatus  int        `gorm:"column:activation_status;default:0"`
	RealNameStatus    int        `gorm:"column:real_name_status;default:0"`
	NetworkStatus     int        `gorm:"column:network_status;default:0"`
	EnablePolling     bool       `gorm:"column:enable_polling;default:true"`
	Creator           uint       `gorm:"column:creator"`
	Updater           uint       `gorm:"column:updater"`
	CreatedAt         time.Time
	UpdatedAt         time.Time
	DeletedAt         *time.Time `gorm:"index"`
}

func (IotCard) TableName() string {
	return "tb_iot_card"
}

var (
	totalCards   = flag.Int("total", 10000000, "要生成的卡数量")
	batchSize    = flag.Int("batch", 10000, "每批插入数量")
	workers      = flag.Int("workers", 10, "并行 worker 数量")
	startICCID   = flag.String("start", "898600000", "起始 ICCID 前缀（9位，总长度不超过20位）")
	clearOld     = flag.Bool("clear", false, "是否清空现有测试卡")

	insertedCount int64
	startTime     time.Time
)

func main() {
	flag.Parse()

	fmt.Println("=== 生成测试卡数据 ===")
	fmt.Printf("目标数量: %d 张\n", *totalCards)
	fmt.Printf("批次大小: %d\n", *batchSize)
	fmt.Printf("并行数: %d\n", *workers)
	fmt.Println("")

	// 连接数据库
	dsn := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable",
		os.Getenv("JUNHONG_DATABASE_HOST"),
		os.Getenv("JUNHONG_DATABASE_PORT"),
		os.Getenv("JUNHONG_DATABASE_USER"),
		os.Getenv("JUNHONG_DATABASE_PASSWORD"),
		os.Getenv("JUNHONG_DATABASE_DBNAME"),
	)

	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	})
	if err != nil {
		log.Fatalf("连接数据库失败: %v", err)
	}

	// 配置连接池
	sqlDB, _ := db.DB()
	sqlDB.SetMaxOpenConns(50)
	sqlDB.SetMaxIdleConns(25)

	fmt.Println("✓ 数据库连接成功")

	// 检查现有卡数量
	var existingCount int64
	db.Model(&IotCard{}).Count(&existingCount)
	fmt.Printf("现有卡数量: %d\n", existingCount)

	if *clearOld {
		fmt.Println("清空现有测试卡...")
		// 只删除 ICCID 以 898600000 开头的测试卡
		db.Exec("DELETE FROM tb_iot_card WHERE iccid LIKE '898600000%'")
		fmt.Println("✓ 清空完成")
	}

	// 开始生成
	startTime = time.Now()
	ctx := context.Background()

	// 创建任务通道
	taskCh := make(chan int, *workers*2)
	var wg sync.WaitGroup

	// 启动 worker
	for i := 0; i < *workers; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			worker(ctx, db, workerID, taskCh)
		}(i)
	}

	// 分发任务
	batches := *totalCards / *batchSize
	for i := 0; i < batches; i++ {
		taskCh <- i
	}
	close(taskCh)

	// 等待完成
	wg.Wait()

	elapsed := time.Since(startTime)
	fmt.Println("")
	fmt.Println("=== 生成完成 ===")
	fmt.Printf("总插入: %d 张\n", atomic.LoadInt64(&insertedCount))
	fmt.Printf("耗时: %v\n", elapsed)
	fmt.Printf("速度: %.0f 张/秒\n", float64(atomic.LoadInt64(&insertedCount))/elapsed.Seconds())

	// 验证
	var finalCount int64
	db.Model(&IotCard{}).Count(&finalCount)
	fmt.Printf("数据库总卡数: %d\n", finalCount)
}

func worker(ctx context.Context, db *gorm.DB, workerID int, taskCh <-chan int) {
	rng := rand.New(rand.NewSource(time.Now().UnixNano() + int64(workerID)))

	for batchIndex := range taskCh {
		cards := generateBatch(rng, *startICCID, batchIndex, *batchSize)

		// 批量插入
		err := db.WithContext(ctx).CreateInBatches(cards, 1000).Error
		if err != nil {
			log.Printf("Worker %d 插入失败: %v", workerID, err)
			continue
		}

		count := atomic.AddInt64(&insertedCount, int64(len(cards)))

		// 进度报告
		if count%100000 == 0 {
			elapsed := time.Since(startTime).Seconds()
			speed := float64(count) / elapsed
			eta := float64(*totalCards-int(count)) / speed
			fmt.Printf("进度: %d/%d (%.1f%%) | 速度: %.0f/秒 | ETA: %.0f秒\n",
				count, *totalCards, float64(count)*100/float64(*totalCards), speed, eta)
		}
	}
}

func generateBatch(rng *rand.Rand, iccidPrefix string, batchIndex int, size int) []IotCard {
	cards := make([]IotCard, size)
	now := time.Now()

	for i := 0; i < size; i++ {
		// 使用前缀 + 序号生成 ICCID（总长度 20 位）
		// 例如: 898600000 (9位) + 00000000001 (11位) = 20 位
		cardIndex := batchIndex*size + i
		iccid := fmt.Sprintf("%s%011d", iccidPrefix, cardIndex)

		// 随机分配状态（匹配轮询配置条件）
		// 实名状态: 0=未实名, 1=实名中, 2=已实名
		// 网络状态: 0=停机, 1=正常
		// 配置匹配逻辑:
		//   - not_real_name: RealNameStatus == 0 或 1
		//   - real_name: RealNameStatus == 2 && NetworkStatus != 1
		//   - activated: RealNameStatus == 2 && NetworkStatus == 1
		r := rng.Float64()
		var realNameStatus, activationStatus, networkStatus int
		if r < 0.10 {
			// 10% 未实名 -> 匹配 not_real_name 配置
			realNameStatus = 0
			activationStatus = 0
			networkStatus = 0
		} else if r < 0.30 {
			// 20% 已实名未激活 -> 匹配 real_name 配置
			realNameStatus = 2
			activationStatus = 0
			networkStatus = 0
		} else {
			// 70% 已激活 -> 匹配 activated 配置（流量+套餐检查）
			realNameStatus = 2
			activationStatus = 1
			networkStatus = 1
		}

		// 随机卡类型
		cardCategory := "normal"
		if rng.Float64() < 0.05 {
			cardCategory = "industry"
		}

		cards[i] = IotCard{
			ICCID:            iccid,
			CardCategory:     cardCategory,
			CarrierID:        uint(rng.Intn(3) + 1), // 1-3 运营商
			Status:           1,
			ActivationStatus: activationStatus,
			RealNameStatus:   realNameStatus,
			NetworkStatus:    networkStatus,
			EnablePolling:    true,
			Creator:          1,
			Updater:          1,
			CreatedAt:        now,
			UpdatedAt:        now,
		}
	}

	return cards
}