Claude Opus 5
@@ -0,0 +1,410 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"91porn-server/common"
|
||||
"91porn-server/common/conf"
|
||||
"91porn-server/common/constant"
|
||||
"91porn-server/common/db"
|
||||
"91porn-server/common/log"
|
||||
"91porn-server/models"
|
||||
"91porn-server/models/Init"
|
||||
"91porn-server/models/commod"
|
||||
"91porn-server/skd/job/dailyAdverCalc"
|
||||
"91porn-server/skd/skdg"
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultConfigPath = "config/skd.json"
|
||||
defaultSeedTag = "abtest-seed"
|
||||
defaultDistrictCode = "ABTEST_SEED"
|
||||
defaultVersion = "1.0.6"
|
||||
defaultStartDate = "2030-01-01"
|
||||
defaultUIDBase = 990000000000000
|
||||
batchSize = 1000
|
||||
)
|
||||
|
||||
type seedConfig struct {
|
||||
ConfigPath string
|
||||
SeedTag string
|
||||
DistrictCode string
|
||||
Version string
|
||||
StartDate string
|
||||
CohortDays int
|
||||
ActiveDays int
|
||||
UsersPerGroup int
|
||||
UIDBase int64
|
||||
Cleanup bool
|
||||
RebuildRetention bool
|
||||
DryRun bool
|
||||
}
|
||||
|
||||
type seedSummary struct {
|
||||
RegDay time.Time
|
||||
AdGroup commod.AdGroup
|
||||
CohortSize int
|
||||
ActiveCounts []int
|
||||
}
|
||||
|
||||
func main() {
|
||||
time.Local, _ = time.LoadLocation("Asia/Shanghai")
|
||||
|
||||
cfg := parseFlags()
|
||||
if err := validateConfig(cfg); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "invalid config: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
log.Init(log.Options{Level: "info"})
|
||||
|
||||
loc, _ := time.LoadLocation("Asia/Shanghai")
|
||||
startDate, err := time.ParseInLocation(time.DateOnly, cfg.StartDate, loc)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "parse start date failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
startDate = common.NormalizeDate(startDate)
|
||||
|
||||
if err := initMongo(cfg.ConfigPath); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "init mongo failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer db.CloseDS()
|
||||
Init.InitMongo()
|
||||
|
||||
users, stats, summaries := buildSeedDocs(cfg, startDate, loc)
|
||||
endActiveDate := startDate.AddDate(0, 0, cfg.CohortDays+cfg.ActiveDays-2)
|
||||
|
||||
printPlan(cfg, startDate, endActiveDate, summaries)
|
||||
if cfg.DryRun {
|
||||
return
|
||||
}
|
||||
|
||||
if cfg.Cleanup {
|
||||
if err := cleanupSeedData(cfg, startDate, endActiveDate); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "cleanup failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
if err := insertInBatches(models.UserTable, users); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "insert users failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if err := insertInBatches(models.UserAdverStat, stats); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "insert user_adver_stat failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if cfg.RebuildRetention {
|
||||
if err := rebuildRetentionRange(startDate, endActiveDate); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "rebuild retention failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("\nseed completed: users=%d stats=%d rebuildRetention=%v\n", len(users), len(stats), cfg.RebuildRetention)
|
||||
}
|
||||
|
||||
func parseFlags() seedConfig {
|
||||
cfg := seedConfig{}
|
||||
flag.StringVar(&cfg.ConfigPath, "config", defaultConfigPath, "path to skd json config")
|
||||
flag.StringVar(&cfg.SeedTag, "seed-tag", defaultSeedTag, "seed tag used in devID/name prefix")
|
||||
flag.StringVar(&cfg.DistrictCode, "district-code", defaultDistrictCode, "district code written into seeded docs")
|
||||
flag.StringVar(&cfg.Version, "version", defaultVersion, "android version written into seeded docs")
|
||||
flag.StringVar(&cfg.StartDate, "start-date", defaultStartDate, "cohort start date in YYYY-MM-DD")
|
||||
flag.IntVar(&cfg.CohortDays, "cohort-days", 5, "number of registration days to generate")
|
||||
flag.IntVar(&cfg.ActiveDays, "active-days", 5, "number of active days to generate for each cohort")
|
||||
flag.IntVar(&cfg.UsersPerGroup, "users-per-group", 30, "base cohort size per ad group per day")
|
||||
flag.Int64Var(&cfg.UIDBase, "uid-base", defaultUIDBase, "starting uid for seeded users")
|
||||
flag.BoolVar(&cfg.Cleanup, "cleanup", true, "cleanup previous seeded docs in the target window")
|
||||
flag.BoolVar(&cfg.RebuildRetention, "rebuild-retention", true, "run CalcRetention for the generated active date range")
|
||||
flag.BoolVar(&cfg.DryRun, "dry-run", false, "print the plan without writing data")
|
||||
flag.Parse()
|
||||
return cfg
|
||||
}
|
||||
|
||||
func validateConfig(cfg seedConfig) error {
|
||||
switch {
|
||||
case cfg.CohortDays <= 0:
|
||||
return fmt.Errorf("cohort-days must be > 0")
|
||||
case cfg.ActiveDays <= 0:
|
||||
return fmt.Errorf("active-days must be > 0")
|
||||
case cfg.UsersPerGroup <= 0:
|
||||
return fmt.Errorf("users-per-group must be > 0")
|
||||
case cfg.UIDBase <= 0:
|
||||
return fmt.Errorf("uid-base must be > 0")
|
||||
case strings.TrimSpace(cfg.SeedTag) == "":
|
||||
return fmt.Errorf("seed-tag must not be empty")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func initMongo(configPath string) error {
|
||||
cfg := skdg.GlobalConfig{}
|
||||
if err := conf.LoadJSON(configPath, &cfg); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
db.InitDS([]db.DBmap{
|
||||
{Key: models.VideoDb, URL: cfg.Mongo.VideoDbUrl},
|
||||
{Key: models.StatDb, URL: cfg.Mongo.StatDbUrl},
|
||||
{Key: models.LogDb, URL: cfg.Mongo.LogDbUrl},
|
||||
}, registerPool())
|
||||
return nil
|
||||
}
|
||||
|
||||
func registerPool() []db.Register {
|
||||
res := make([]db.Register, len(models.RegisterPool))
|
||||
for i, v := range models.RegisterPool {
|
||||
res[i] = db.Register{Key: v.Key, Table: v.Table}
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
func buildSeedDocs(cfg seedConfig, startDate time.Time, loc *time.Location) ([]bson.M, []bson.M, []seedSummary) {
|
||||
var (
|
||||
users []bson.M
|
||||
stats []bson.M
|
||||
summaries []seedSummary
|
||||
nextUID = cfg.UIDBase
|
||||
)
|
||||
|
||||
for dayIdx := 0; dayIdx < cfg.CohortDays; dayIdx++ {
|
||||
regDay := startDate.AddDate(0, 0, dayIdx)
|
||||
for _, group := range []commod.AdGroup{commod.AdGroupA, commod.AdGroupB, commod.AdGroupC} {
|
||||
cohortSize := cohortSizeFor(group, dayIdx, cfg.UsersPerGroup)
|
||||
activeCounts := make([]int, cfg.ActiveDays)
|
||||
cohortUIDs := make([]int64, 0, cohortSize)
|
||||
|
||||
for userIdx := 0; userIdx < cohortSize; userIdx++ {
|
||||
uid := nextUID
|
||||
nextUID++
|
||||
cohortUIDs = append(cohortUIDs, uid)
|
||||
|
||||
createdAt := time.Date(regDay.Year(), regDay.Month(), regDay.Day(), 10, userIdx%50, 0, 0, loc)
|
||||
devID := fmt.Sprintf("%s-%s-%02d-%03d", cfg.SeedTag, group, dayIdx, userIdx)
|
||||
users = append(users, bson.M{
|
||||
"uid": uint64(uid),
|
||||
"name": fmt.Sprintf("%s_%s_d%02d_u%03d", cfg.SeedTag, group, dayIdx, userIdx),
|
||||
"devID": devID,
|
||||
"districtCode": cfg.DistrictCode,
|
||||
"sysType": constant.SysTypeAndroid,
|
||||
"oriVer": cfg.Version,
|
||||
"adGroup": group,
|
||||
"createdAt": createdAt,
|
||||
"updatedAt": createdAt,
|
||||
})
|
||||
}
|
||||
|
||||
for activeOffset := 0; activeOffset < cfg.ActiveDays; activeOffset++ {
|
||||
activeCount := activeUserCount(cohortSize, groupActivityRate(group, activeOffset))
|
||||
activeCounts[activeOffset] = activeCount
|
||||
activeDay := common.NormalizeDate(regDay.AddDate(0, 0, activeOffset))
|
||||
|
||||
for userIdx := 0; userIdx < activeCount; userIdx++ {
|
||||
uid := cohortUIDs[userIdx]
|
||||
adClick := int64(groupWeight(group) + 1 + activeOffset%2)
|
||||
appClick := int64(1)
|
||||
payCount, payTotal := payForUser(group, activeOffset, userIdx)
|
||||
|
||||
stats = append(stats, bson.M{
|
||||
"userId": uid,
|
||||
"date": activeDay,
|
||||
"districtCode": cfg.DistrictCode,
|
||||
"sysType": constant.SysTypeAndroid,
|
||||
"adGroup": group,
|
||||
"adClick": adClick,
|
||||
"appClick": appClick,
|
||||
"totalClick": adClick + appClick,
|
||||
"payCount": payCount,
|
||||
"payTotal": payTotal,
|
||||
"regDay": common.NormalizeDate(regDay),
|
||||
"version": cfg.Version,
|
||||
"createdAt": activeDay.Add(2 * time.Hour),
|
||||
"updatedAt": activeDay.Add(2 * time.Hour),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
summaries = append(summaries, seedSummary{
|
||||
RegDay: common.NormalizeDate(regDay),
|
||||
AdGroup: group,
|
||||
CohortSize: cohortSize,
|
||||
ActiveCounts: activeCounts,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return users, stats, summaries
|
||||
}
|
||||
|
||||
func cohortSizeFor(group commod.AdGroup, dayIdx, base int) int {
|
||||
return base + dayIdx*2 + groupWeight(group)*3
|
||||
}
|
||||
|
||||
func groupWeight(group commod.AdGroup) int {
|
||||
switch group {
|
||||
case commod.AdGroupA:
|
||||
return 0
|
||||
case commod.AdGroupB:
|
||||
return 1
|
||||
default:
|
||||
return 2
|
||||
}
|
||||
}
|
||||
|
||||
func groupActivityRate(group commod.AdGroup, dayOffset int) float64 {
|
||||
rates := map[commod.AdGroup][]float64{
|
||||
commod.AdGroupA: {0.85, 0.60, 0.45, 0.30, 0.20, 0.12, 0.08},
|
||||
commod.AdGroupB: {0.70, 0.50, 0.35, 0.22, 0.14, 0.08, 0.05},
|
||||
commod.AdGroupC: {0.55, 0.35, 0.20, 0.10, 0.06, 0.03, 0.02},
|
||||
}
|
||||
if dayOffset < len(rates[group]) {
|
||||
return rates[group][dayOffset]
|
||||
}
|
||||
last := rates[group][len(rates[group])-1]
|
||||
decay := math.Pow(0.7, float64(dayOffset-len(rates[group])+1))
|
||||
return last * decay
|
||||
}
|
||||
|
||||
func activeUserCount(cohortSize int, rate float64) int {
|
||||
if cohortSize <= 0 || rate <= 0 {
|
||||
return 0
|
||||
}
|
||||
count := int(math.Round(float64(cohortSize) * rate))
|
||||
if count < 0 {
|
||||
return 0
|
||||
}
|
||||
if count > cohortSize {
|
||||
return cohortSize
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func payForUser(group commod.AdGroup, dayOffset, userIdx int) (int64, int64) {
|
||||
switch group {
|
||||
case commod.AdGroupA:
|
||||
if dayOffset <= 2 && userIdx%7 == 0 {
|
||||
return 1, 1999
|
||||
}
|
||||
case commod.AdGroupB:
|
||||
if dayOffset <= 1 && userIdx%9 == 0 {
|
||||
return 1, 1299
|
||||
}
|
||||
case commod.AdGroupC:
|
||||
if dayOffset == 0 && userIdx%12 == 0 {
|
||||
return 1, 699
|
||||
}
|
||||
}
|
||||
return 0, 0
|
||||
}
|
||||
|
||||
func cleanupSeedData(cfg seedConfig, startDate, endActiveDate time.Time) error {
|
||||
userFilter := bson.M{
|
||||
"devID": bson.M{
|
||||
"$regex": fmt.Sprintf("^%s-", regexp.QuoteMeta(cfg.SeedTag)),
|
||||
},
|
||||
}
|
||||
statFilter := bson.M{
|
||||
"districtCode": cfg.DistrictCode,
|
||||
"version": cfg.Version,
|
||||
"regDay": bson.M{
|
||||
"$gte": startDate,
|
||||
"$lt": startDate.AddDate(0, 0, cfg.CohortDays),
|
||||
},
|
||||
"date": bson.M{
|
||||
"$gte": startDate,
|
||||
"$lte": endActiveDate,
|
||||
},
|
||||
}
|
||||
|
||||
if _, err := db.BaseDAO(models.UserTable, nil).DeleteMany(userFilter); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := db.BaseDAO(models.UserAdverStat, nil).DeleteMany(statFilter); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if cfg.RebuildRetention {
|
||||
retentionFilter := bson.M{
|
||||
"date": bson.M{
|
||||
"$gte": startDate,
|
||||
"$lt": startDate.AddDate(0, 0, cfg.CohortDays),
|
||||
},
|
||||
"adGroup": bson.M{
|
||||
"$in": []commod.AdGroup{commod.AdGroupA, commod.AdGroupB, commod.AdGroupC},
|
||||
},
|
||||
}
|
||||
if _, err := db.BaseDAO(models.DailyRetention, nil).DeleteMany(retentionFilter); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func insertInBatches(table string, docs []bson.M) error {
|
||||
if len(docs) == 0 {
|
||||
return nil
|
||||
}
|
||||
dao := db.BaseDAO(table, nil)
|
||||
for start := 0; start < len(docs); start += batchSize {
|
||||
end := start + batchSize
|
||||
if end > len(docs) {
|
||||
end = len(docs)
|
||||
}
|
||||
batch := docs[start:end]
|
||||
writes := make([]mongo.WriteModel, 0, len(batch))
|
||||
for _, doc := range batch {
|
||||
writes = append(writes, mongo.NewInsertOneModel().SetDocument(doc))
|
||||
}
|
||||
if _, err := dao.Bulk(writes, options.BulkWrite().SetOrdered(false)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func rebuildRetentionRange(startDate, endActiveDate time.Time) error {
|
||||
ctx := context.Background()
|
||||
for d := common.NormalizeDate(startDate); !d.After(endActiveDate); d = d.AddDate(0, 0, 1) {
|
||||
if err := dailyAdverCalc.CalcRetention(ctx, d); err != nil {
|
||||
return fmt.Errorf("calc retention for %s failed: %w", d.Format(time.DateOnly), err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func printPlan(cfg seedConfig, startDate, endActiveDate time.Time, summaries []seedSummary) {
|
||||
fmt.Printf("seed tag: %s\n", cfg.SeedTag)
|
||||
fmt.Printf("date range: cohort=%s ~ %s, active=%s ~ %s\n",
|
||||
startDate.Format(time.DateOnly),
|
||||
startDate.AddDate(0, 0, cfg.CohortDays-1).Format(time.DateOnly),
|
||||
startDate.Format(time.DateOnly),
|
||||
endActiveDate.Format(time.DateOnly),
|
||||
)
|
||||
fmt.Printf("base users/group/day: %d, cleanup=%v, rebuildRetention=%v, dryRun=%v\n\n",
|
||||
cfg.UsersPerGroup, cfg.Cleanup, cfg.RebuildRetention, cfg.DryRun,
|
||||
)
|
||||
|
||||
for _, item := range summaries {
|
||||
rates := make([]string, 0, len(item.ActiveCounts))
|
||||
for _, active := range item.ActiveCounts {
|
||||
rates = append(rates, fmt.Sprintf("%.4f", float64(active)/float64(item.CohortSize)))
|
||||
}
|
||||
fmt.Printf("%s group=%s cohort=%d active=%v rates=%v\n",
|
||||
item.RegDay.Format(time.DateOnly), item.AdGroup, item.CohortSize, item.ActiveCounts, rates,
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user