apollo-backend/internal/cmd/worker.go
2021-07-15 10:51:34 -04:00

81 lines
1.6 KiB
Go

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{
"notifications": worker.NewNotificationsWorker,
}
)
func WorkerCmd(ctx context.Context) *cobra.Command {
var multiplier 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")
}
logger := cmdutil.NewLogrusLogger(false)
statsd, err := cmdutil.NewStatsdClient()
if err != nil {
return err
}
defer statsd.Close()
db, err := cmdutil.NewDatabasePool(ctx)
if err != nil {
return err
}
defer db.Close()
redis, err := cmdutil.NewRedisClient(ctx)
if err != nil {
return err
}
defer redis.Close()
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)
}
consumers := runtime.NumCPU() * multiplier
//consumers = 1
worker := workerFn(logger, statsd, db, redis, queue, consumers)
worker.Start()
<-ctx.Done()
worker.Stop()
return nil
},
}
cmd.Flags().IntVar(&multiplier, "multiplier", 12, "The multiplier (by CPUs) to run")
cmd.Flags().StringVar(&queueID, "queue", "", "The queue to work on")
return cmd
}