summercms/modules/conga/commands.go

package conga

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"git.golem15.com/golem15/summercms/modules/backpack"
	"git.golem15.com/golem15/summercms/modules/bonfire"
	"git.golem15.com/golem15/summercms/modules/lagoon"
	"git.golem15.com/golem15/summercms/modules/party"
	"github.com/riverqueue/river"
	"github.com/riverqueue/river/rivertype"
)

// clearBatch is the number of jobs one JobDeleteMany pass removes.
const clearBatch = 10000

// RuntimeCommands returns the queue:work and queue:clear commands for the
// application binary.
func RuntimeCommands(app *backpack.App, plugins []party.Plugin) []bonfire.Command {
	return []bonfire.Command{
		{
			Name:        "queue:work",
			Description: "Run background job workers in the foreground",
			Flags: []bonfire.Flag{{
				Name:        "queue",
				Description: "Queue to work (repeatable; default all known queues)",
				Repeatable:  true,
			}},
			Run: func(ctx context.Context, in bonfire.Input, out bonfire.Output) error {
				return withDB(ctx, app, func() error {
					return work(ctx, app, plugins, in.Flags("queue"), out)
				})
			},
		},
		{
			Name:        "queue:clear",
			Description: "Clear all queued jobs, by deleting all pending jobs.",
			Args: []bonfire.Arg{{
				Name:        "queue",
				Description: "The name of the queue to clear (default \"default\")",
			}},
			Run: func(ctx context.Context, in bonfire.Input, out bonfire.Output) error {
				queue, _ := in.Argument("queue")
				queue = strings.TrimSpace(queue)
				if queue == "" {
					queue = defaultQueue
				}
				return withDB(ctx, app, func() error {
					return clearQueue(ctx, app, queue, out)
				})
			},
		},
	}
}

func work(ctx context.Context, app *backpack.App, plugins []party.Plugin, queues []string, out bonfire.Output) error {
	ctx, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM)
	defer stop()
	w, err := StartWorker(ctx, app, plugins, WorkerOptions{Queues: queues})
	if err != nil {
		return err
	}
	out.Info("worker started on queues: " + strings.Join(w.Queues(), ", "))
	<-ctx.Done()
	stopCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	return w.Stop(stopCtx)
}

// clearQueue deletes the available, scheduled and retryable jobs of queue in
// batches until a pass deletes none. River never deletes running jobs.
func clearQueue(ctx context.Context, app *backpack.App, queue string, out bonfire.Output) error {
	m, err := From(app)
	if err != nil {
		return err
	}
	client, err := m.insertClient()
	if err != nil {
		return err
	}
	out.Info(fmt.Sprintf("Clearing queue %q", queue))
	total := 0
	for {
		res, err := client.JobDeleteMany(ctx, river.NewJobDeleteManyParams().
			Queues(queue).
			States(rivertype.JobStateAvailable, rivertype.JobStateScheduled, rivertype.JobStateRetryable).
			First(clearBatch))
		if err != nil {
			return fmt.Errorf("conga: clear queue %q: %w", queue, err)
		}
		if len(res.Jobs) == 0 {
			break
		}
		total += len(res.Jobs)
	}
	out.Info(fmt.Sprintf("Cleared %d jobs", total))
	return nil
}

// withDB opens and publishes the shared pool for a CLI command and closes it
// when fn returns, the way lagoon's own commands do.
func withDB(ctx context.Context, app *backpack.App, fn func() error) error {
	sqlDB, gdb, err := lagoon.OpenFromApp(ctx, app)
	if err != nil {
		return err
	}
	defer sqlDB.Close()
	if err := lagoon.Publish(app, sqlDB, gdb); err != nil {
		return err
	}
	return fn()
}