package main import ( "flag" "os" "os/signal" "syscall" "github.com/tonywangcn/tiktok-crawler/pkg/log" "github.com/tonywangcn/tiktok-crawler/workers" ) var workerMap map[string]func(count int) = map[string]func(count int){ "tiktok": workers.RunTikTokWorker, "tiktok-exp": workers.RunTikTokExploreWorker, } func main() { shutdown := make(chan int) sigChan := make(chan os.Signal, 1) signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGINT) go func() { <-sigChan log.Info("Shutting down...") shutdown <- 1 }() var worker string var isServer bool var count int flag.BoolVar(&isServer, "s", false, "Run as server") flag.StringVar(&worker, "w", "tiktok", "Choose the worker to run") flag.IntVar(&count, "c", 20, "Number of workers to run") flag.Parse() if isServer { go Server() } else { f, ok := workerMap[worker] if !ok { panic("worker not found") } f(count) } <-shutdown }