apollo-backend/internal/cmd/worker.go

package cmd

import (
	"context"
	"fmt"
	"runtime"

	faktoryworker "github.com/contribsys/faktory_worker_go"
	"github.com/spf13/cobra"

	"github.com/christianselig/apollo-backend/internal/cmdutil"
	"github.com/christianselig/apollo-backend/internal/worker"
)

var (
	queues = map[string]worker.NewWorkerFn{
		"LiveActivityJob":       worker.NewLiveActivitiesWorker,
		"NotificationCheckJob":  worker.NewNotificationsWorker,
		"StuckNotificationsJob": worker.NewStuckNotificationsWorker,
		"SubredditWatcherJob":   worker.NewSubredditsWorker,
		"SubredditTrendingJob":  worker.NewTrendingWorker,
	}
)

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")
			}

			runtime.SetBlockProfileRate(1)
			runtime.SetMutexProfileFraction(1)

			svc := fmt.Sprintf("worker: %s", queueID)
			logger := cmdutil.NewLogger(svc)
			defer func() { _ = logger.Sync() }()

			tag := fmt.Sprintf("worker:%s", queueID)
			statsd, err := cmdutil.NewStatsdClient(tag)
			if err != nil {
				return err
			}
			defer statsd.Close()

			poolSize := consumers / 4

			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()

			workerFn, ok := queues[queueID]
			if !ok {
				return fmt.Errorf("queue does not exist: %s", queueID)
			}
			worker := workerFn(ctx, logger, statsd, db, redis, consumers)

			mgr := faktoryworker.NewManager()
			mgr.Concurrency = consumers
			mgr.Register(queueID, worker.Process)
			return mgr.RunWithContext(ctx)
		},
	}

	cmd.Flags().IntVar(&consumers, "consumers", runtime.NumCPU()*64, "The consumers to run")
	cmd.Flags().StringVar(&queueID, "queue", "", "The queue to work on")

	return cmd
}