2021-07-13 17:14:25 +00:00
|
|
|
package cmd
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"fmt"
|
|
|
|
"runtime"
|
|
|
|
|
2022-11-02 04:49:20 +00:00
|
|
|
faktoryworker "github.com/contribsys/faktory_worker_go"
|
2021-07-13 17:14:25 +00:00
|
|
|
"github.com/spf13/cobra"
|
|
|
|
|
|
|
|
"github.com/christianselig/apollo-backend/internal/cmdutil"
|
|
|
|
"github.com/christianselig/apollo-backend/internal/worker"
|
|
|
|
)
|
|
|
|
|
|
|
|
var (
|
|
|
|
queues = map[string]worker.NewWorkerFn{
|
2022-11-02 04:49:20 +00:00
|
|
|
"LiveActivityJob": worker.NewLiveActivitiesWorker,
|
|
|
|
"NotificationCheckJob": worker.NewNotificationsWorker,
|
|
|
|
"StuckNotificationsJob": worker.NewStuckNotificationsWorker,
|
|
|
|
"SubredditWatcherJob": worker.NewSubredditsWorker,
|
|
|
|
"SubredditTrendingJob": worker.NewTrendingWorker,
|
2021-07-13 17:14:25 +00:00
|
|
|
}
|
|
|
|
)
|
|
|
|
|
|
|
|
func WorkerCmd(ctx context.Context) *cobra.Command {
|
2022-10-31 21:33:41 +00:00
|
|
|
var consumers int
|
2021-07-13 17:14:25 +00:00
|
|
|
var queueID string
|
|
|
|
|
|
|
|
cmd := &cobra.Command{
|
|
|
|
Use: "worker",
|
|
|
|
Args: cobra.ExactArgs(0),
|
|
|
|
Short: "Work through job queues.",
|
|
|
|
RunE: func(cmd *cobra.Command, args []string) error {
|
|
|
|
if queueID == "" {
|
|
|
|
return fmt.Errorf("need a queue to work on")
|
|
|
|
}
|
|
|
|
|
2022-11-01 23:36:40 +00:00
|
|
|
runtime.SetBlockProfileRate(1)
|
|
|
|
runtime.SetMutexProfileFraction(1)
|
|
|
|
|
2022-05-23 18:29:15 +00:00
|
|
|
svc := fmt.Sprintf("worker: %s", queueID)
|
|
|
|
logger := cmdutil.NewLogger(svc)
|
2022-05-23 18:17:25 +00:00
|
|
|
defer func() { _ = logger.Sync() }()
|
2021-07-13 17:14:25 +00:00
|
|
|
|
2021-10-12 16:04:20 +00:00
|
|
|
tag := fmt.Sprintf("worker:%s", queueID)
|
|
|
|
statsd, err := cmdutil.NewStatsdClient(tag)
|
2021-07-13 17:14:25 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
defer statsd.Close()
|
|
|
|
|
2022-10-31 21:33:41 +00:00
|
|
|
poolSize := consumers / 4
|
2021-07-15 17:47:22 +00:00
|
|
|
|
2021-07-23 01:11:27 +00:00
|
|
|
db, err := cmdutil.NewDatabasePool(ctx, poolSize)
|
2021-07-13 17:14:25 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
defer db.Close()
|
|
|
|
|
2022-10-31 21:33:41 +00:00
|
|
|
redis, err := cmdutil.NewRedisClient(ctx, consumers)
|
2021-07-13 17:14:25 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
defer redis.Close()
|
|
|
|
|
|
|
|
workerFn, ok := queues[queueID]
|
|
|
|
if !ok {
|
2022-11-02 04:49:20 +00:00
|
|
|
return fmt.Errorf("queue does not exist: %s", queueID)
|
2021-07-13 17:14:25 +00:00
|
|
|
}
|
2022-11-02 04:49:20 +00:00
|
|
|
worker := workerFn(ctx, logger, statsd, db, redis, consumers)
|
2021-07-13 17:14:25 +00:00
|
|
|
|
2022-11-02 04:49:20 +00:00
|
|
|
mgr := faktoryworker.NewManager()
|
|
|
|
mgr.Concurrency = consumers
|
|
|
|
mgr.Register(queueID, worker.Process)
|
|
|
|
return mgr.RunWithContext(ctx)
|
2021-07-13 17:14:25 +00:00
|
|
|
},
|
|
|
|
}
|
|
|
|
|
2022-10-31 21:33:41 +00:00
|
|
|
cmd.Flags().IntVar(&consumers, "consumers", runtime.NumCPU()*64, "The consumers to run")
|
2021-07-13 17:14:25 +00:00
|
|
|
cmd.Flags().StringVar(&queueID, "queue", "", "The queue to work on")
|
|
|
|
|
|
|
|
return cmd
|
|
|
|
}
|