apollo-backend/internal/cmd/worker.go

94 lines
2.1 KiB
Go
Raw Normal View History

package cmd
import (
"context"
"fmt"
"runtime"
"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-10-19 13:37:41 +00:00
"live-activities": worker.NewLiveActivitiesWorker,
2021-10-17 14:17:41 +00:00
"notifications": worker.NewNotificationsWorker,
"stuck-notifications": worker.NewStuckNotificationsWorker,
"subreddits": worker.NewSubredditsWorker,
"trending": worker.NewTrendingWorker,
"users": worker.NewUsersWorker,
}
)
func WorkerCmd(ctx context.Context) *cobra.Command {
var consumers int
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-10-12 16:04:20 +00:00
tag := fmt.Sprintf("worker:%s", queueID)
statsd, err := cmdutil.NewStatsdClient(tag)
if err != nil {
return err
}
defer statsd.Close()
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)
if err != nil {
return err
}
defer db.Close()
redis, err := cmdutil.NewRedisClient(ctx, consumers)
if err != nil {
return err
}
defer redis.Close()
2022-11-01 13:53:09 +00:00
queue, err := cmdutil.NewQueueClient(logger, redis, "worker")
if err != nil {
return err
}
workerFn, ok := queues[queueID]
if !ok {
return fmt.Errorf("invalid queue: %s", queueID)
}
worker := workerFn(ctx, logger, statsd, db, redis, queue, consumers)
2021-09-25 13:19:42 +00:00
if err := worker.Start(); err != nil {
return err
}
<-ctx.Done()
worker.Stop()
return nil
},
}
cmd.Flags().IntVar(&consumers, "consumers", runtime.NumCPU()*64, "The consumers to run")
cmd.Flags().StringVar(&queueID, "queue", "", "The queue to work on")
return cmd
}