package main import ( "encoding/json" "fmt" "log" "net/http" "os" "os/signal" "sync" "time" "strconv" acfun_api "git.saveweb.org/saveweb/acfunction_go/pkg" savewebtracker "git.saveweb.org/saveweb/saveweb_tracker/src/saveweb_tracker" "github.com/hashicorp/go-retryablehttp" "github.com/tidwall/gjson" ) var BASE_CONCURRENCY = 10 var WITH_DELAY = true var tasks_chan = make(chan savewebtracker.Task, BASE_CONCURRENCY) var Interrupted = false var WaitClaimWorker sync.WaitGroup var WaitProcesserWorker sync.WaitGroup var Logger = log.New(os.Stdout, "[acdanmaku] ", log.Ldate|log.Ltime|log.Lmsgprefix) var DEBUG = false func init() { if os.Getenv("BASE_CONCURRENCY") != "" { fmt.Println("BASE_CONCURRENCY:", os.Getenv("BASE_CONCURRENCY")) BASE_CONCURRENCY, _ = strconv.Atoi(os.Getenv("BASE_CONCURRENCY")) } if os.Getenv("NO_WITH_DELAY") != "" { fmt.Println("NO_WITH_DELAY:", os.Getenv("NO_WITH_DELAY")) WITH_DELAY = false } if os.Getenv("DEBUG") != "" { DEBUG = true } } // ClaimTask 并把任务放入 task_chan func claimWorker(i int, tracker *savewebtracker.Tracker) { workerName := fmt.Sprintf("[ClaimWorker(%d)]", i) Logger.Println("[START]", workerName) defer Logger.Println("[STOP]", workerName, " exited...") defer WaitClaimWorker.Done() for { if Interrupted { return } task := tracker.ClaimTask(WITH_DELAY) if task == nil { notask_sleep := max( time.Duration(tracker.Project().Client.ClaimTaskDelay)*10*time.Second, time.Duration(10)*time.Second, ) Logger.Println(workerName, "No task to claim, sleep", notask_sleep) time.Sleep(notask_sleep) continue } Logger.Println(workerName, "Claimed task", task.Id) tasks_chan <- *task } } func ProcesserWorker(i int, tracker *savewebtracker.Tracker) { workerName := fmt.Sprintf("[ProcesserWorker(%d)]", i) Logger.Println("[START]", workerName) defer Logger.Println("[STOP]", workerName, " exited...") defer WaitProcesserWorker.Done() for task := range tasks_chan { Logger.Println(workerName, "Processing task", task.Id) // 在这儿处理任务 danmakus, err := acfun_api.GetDanmaku(tracker.HTTP_client, task.Id) if err != nil { tracker.UpdateTask(task.Id, task.Id_type, savewebtracker.StatusFAIL) Logger.Println(workerName, "Failed to get danmaku", task.Id, err) Interrupted = true continue } items := []savewebtracker.Item{} for _, danmaku := range danmakus { r_danmakuId := gjson.Get(danmaku, "danmakuId") if !r_danmakuId.Exists() || r_danmakuId.Type != gjson.Number { Logger.Println(workerName, "danmakuId not found or not a number") tracker.UpdateTask(task.Id, task.Id_type, savewebtracker.StatusFAIL) Interrupted = true continue } danmakuId := r_danmakuId.Int() items = append(items, savewebtracker.Item{ Item_id: fmt.Sprintf("%d", danmakuId), Item_id_type: "int", Item_status: "None", Item_status_type: "None", Payload: danmaku, }) } if len(items) != 0 { resp := tracker.InsertMany(items) Logger.Println(workerName, "InsertMany", "task", task.Id, "->", len(items), "items", resp) } tracker.UpdateTask(task.Id, task.Id_type, savewebtracker.StatusDONE) Logger.Println(workerName, "UpdateTask", task.Id) } } func InterruptHandler() { fmt.Println("Press Ctrl+C to exit") interrupt_c := make(chan os.Signal, 1) signal.Notify(interrupt_c, os.Interrupt) for { s := <-interrupt_c Logger.Println("Interrupted by", s, "signal (Press Ctrl+C again to force exit)") if Interrupted { Logger.Println("Force exit") os.Exit(1) return } Interrupted = true } } func GetRetryableHttpClient(timeout time.Duration, debug bool) *http.Client { retryClient := retryablehttp.NewClient() retryClient.RetryMax = 3 retryClient.RetryWaitMin = 1 * time.Second retryClient.RetryWaitMax = 10 * time.Second retryClient.HTTPClient.Timeout = timeout if !debug { retryClient.Logger = nil } standardClient := retryClient.StandardClient() // *http.Client Logger.Println("standardClient.Timeout:", standardClient.Timeout) return standardClient } func ShowStatus(t *savewebtracker.Tracker) { for { project_json, err := json.Marshal(t.Project()) if err != nil { panic(err) } Logger.Println("Project:", string(project_json)) time.Sleep(60 * time.Second) } } func main() { tracker := savewebtracker.GetTracker("acdanmaku", "0.1", savewebtracker.Archivist()) tracker.PING_client = GetRetryableHttpClient(10*time.Second, DEBUG) tracker.HTTP_client = GetRetryableHttpClient(60*time.Second, DEBUG) tracker.SelectBestTracker() _, err := tracker.FetchProject(10 * time.Second) if err != nil { panic(err) } tracker.StartSelectTrackerBackground().StartFetchProjectBackground() go InterruptHandler() go ShowStatus(tracker) Logger.Println("-- Start --") for i := 0; i < BASE_CONCURRENCY; i++ { go claimWorker(i, tracker) WaitClaimWorker.Add(1) go ProcesserWorker(i, tracker) WaitProcesserWorker.Add(1) } // wait for all claimWorker to finish WaitClaimWorker.Wait() Logger.Println("[STOP] All claimWorker done") // close task_chan close(tasks_chan) Logger.Println("[STOP] task_chan closed") // wait for all task_chan to finish WaitProcesserWorker.Wait() Logger.Println("[STOP] All ProcesserWorker done") Logger.Println("-- All done --") }