acfunction_go/cmd/acdanmaku/acdanmaku.go

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:      task.Id,
				Item_status_type: "int",
				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.2.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 --")

}