ClientFromContext to submit to another queue?

[vsoch]

I'm creating a client and workers as follows - note there is a main queue that a scheduler is using to receive jobs (and then allocate resources) and a cleanup queue that will have a job added after the resources are allocated. For example "This job was just scheduled for a duration of an hour, make sure to run the cancel / cleanup worker in an hour."

// NewQueue starts a new queue with a river client
func NewQueue(ctx context.Context) (*Queue, error) {
	dbPool, err := pgxpool.New(ctx, os.Getenv("DATABASE_URL"))
	if err != nil {
		return nil, err
	}

	// The default strategy now mirrors what fluence with Kubernetes does
	// This can eventually be customizable. We provide the pool to the
	// strategy because it also manages the provisional queue.
	strategy := strategies.EasyBackfill{}
	workers := river.NewWorkers()

	// Each strategy has its own worker type
	strategy.AddWorkers(workers)
	riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{
		// Change the verbosity of the logger here
		Logger: slog.New(&slogutil.SlogMessageOnlyHandler{Level: slog.LevelWarn}),
		Queues: map[string]river.QueueConfig{

			// Default queue handles job allocation
			river.QueueDefault: {MaxWorkers: queueMaxWorkers},

			// Cleanup queue is only for cancel
			"cancel_queue": {MaxWorkers: queueMaxWorkers},
		},
		Workers: workers,
	})
	if err != nil {
		return nil, err
	}

	// Create the queue and setup events for it
	err = riverClient.Start(ctx)
	if err != nil {
		return nil, err
	}

	queue := Queue{
		riverClient:      riverClient,
		Pool:             dbPool,
		Strategy:         strategy,
		Context:          ctx,
		ReservationDepth: depth,
	}
	queue.setupEvents()
	return &queue, nil
}

And note that AddWorkers for the strategy is registering both types:

// AddtWorkers adds the worker for the queue strategy
// job worker: a queue to submit jobs to fluxion
// cleanup worker: a queue to cleanup
func (EasyBackfill) AddWorkers(workers *river.Workers) {
	river.AddWorker(workers, &work.JobWorker{})
	river.AddWorker(workers, &work.CleanupWorker{})
}

In the work function, we define a job as needed a spec with resources,

// The job worker submits jobs to fluxion with match allocate
// or match allocate else reserve, depending on the reservation depth
func (args JobArgs) Kind() string { return "job" }

type JobWorker struct {
	river.WorkerDefaults[JobArgs]
}

type JobArgs struct {

	// Submit Args
	Jobspec   string `json:"jobspec"`
	Podspec   string `json:"podspec"`
	GroupName string `json:"groupName"`
	GroupSize int32  `json:"groupSize"`

	// If true, we are allowed to ask Fluxion for a reservation
	Reservation bool `json:"reservation"`

	// Attributes updated to send back via events
	Nodes   string `json:"nodes"`
	FluxJob int64  `json:"jobid"`
	PodId   string `json:"podid"`
}

// Work performs a call to the scheduler, Fluxion
func (w JobWorker) Work(ctx context.Context, job *river.Job[JobArgs]) error {

       // Assume work is done here to ask the scheduler, and get an allocation
       ...
   
	// Kick off a cleaning job for when everyting should be cancelled
        // This is a duration of 10 seconds for now
	return SubmitCleanup(ctx, pool, 10, int64(fluxID), true, []string{})
}

Note that the context above given to SubmitCleanup is the same one for the worker, so we should be able to create an "in worker" client. Here that is:

// The cleanup workers cleans up a reservation (issuing cancel)
func (args CleanupArgs) Kind() string { return "cleanup" }

type CleanupWorker struct {
	river.WorkerDefaults[CleanupArgs]
}

// SubmitCleanup submits a cleanup job N seconds into the future
func SubmitCleanup(
	ctx context.Context,
	pool *pgxpool.Pool,
	seconds int32,
	fluxID int64,
	inKubernetes bool,
	tags []string,
) error {

	klog.Infof("SUBMIT CLEANUP starting for %d", fluxID)

	client, err := river.ClientFromContextSafely[pgx.Tx](ctx)
	if err != nil {
		return fmt.Errorf("error getting client from context: %w", err)
	}
	tx, err := pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)

	// Create scheduledAt time - N seconds from now
	now := time.Now()
	scheduledAt := now.Add(time.Second * time.Duration(seconds))

	insertOpts := river.InsertOpts{
		MaxAttempts: defaults.MaxAttempts,
		Tags:        tags,
		Queue:       "cleanup_queue",
		ScheduledAt: scheduledAt,
	}
	_, err = client.InsertTx(ctx, tx, CleanupArgs{FluxID: fluxID, Kubernetes: inKubernetes}, &insertOpts)
	if err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return err
	}
	klog.Infof("SUBMIT CLEANUP ending for %d", fluxID)
	return nil
}

The problem is that I can see that function runs through completion (I see the SUBMIT CLEANUP at the end) but I never see any output from the job, where the first thing that happens is a print statement. Is it the case that a worker of type A cannot access a worker of type B (associated with the same client) because it's a different queue? Or there is something off about the context I'm using, or another step to register the job type again? Thanks for the help!

[vsoch]

oh wait I think I see the issue - I hope that's it at least. Brb!

[vsoch]

Yep that was it - the queue name was wrong 🙃

- Queue:       "cleanup_queue",
+ Queue:       "cancel_queue",

Sorry for the noise - it was actually very helpful for me to write those two chunks in one place and quickly see the issue! Moving on to the next step (this project is immensely fun).

[bgentry]

@vsoch Hi, going forward could you hold off on creating issues until you're fairly sure there's a confirmed bug in River and you have an isolated, minimal, and reproducible example to demonstrate it? There are other people who subscribe to the issue tracker here and all these issues result in a lot of broadcasted noise for what is ultimately you looking for community support.

The discussions section is much more useful for that, btw, which is why I keep converting your issues into discussions.