Method Work

internal/rivermiddleware/middleware.go:30–79

(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error)

func (*ResumableMiddleware) IsMiddleware() bool { return true }

func (*ResumableMiddleware) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error {
	metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx)
	if !hasMetadataUpdates {
		return errors.New("expected to find metadata updates in context, but didn't")
	}

	state := &ResumableState{
		AllStepNames:  make(map[string]struct{}),
		Cursors:       make(map[string]json.RawMessage),
		ResumeMatched: true,
		ResumeStep:    gjson.GetBytes(job.Metadata, rivercommon.MetadataKeyResumableStep).Str,
	}
	if state.ResumeStep != "" {
		state.ResumeMatched = false
	}

	hadCursors := false
	if cursorJSON := gjson.GetBytes(job.Metadata, rivercommon.MetadataKeyResumableCursor); cursorJSON.Exists() && cursorJSON.Type == gjson.JSON {
		if err := json.Unmarshal([]byte(cursorJSON.Raw), &state.Cursors); err != nil {
			return fmt.Errorf("river: unmarshal resumable cursors: %w", err)
		}
		hadCursors = len(state.Cursors) > 0
	}

	ctx = context.WithValue(ctx, ResumableContextKey{}, state)

	err := doInner(ctx)
	if err == nil {
		switch {
		case state.Err != nil:
			err = state.Err
		case state.ResumeStep != "" && !state.ResumeMatched:
			err = fmt.Errorf("river: resumable step %q not found in Worker", state.ResumeStep)
		}
	}

	if err != nil && state.CompletedStep != "" {
		if len(state.Cursors) > 0 {
			metadataUpdates[rivercommon.MetadataKeyResumableCursor] = state.Cursors
		} else if hadCursors {
			// All cursors were consumed (their steps completed). Write null
			// to clear the stale cursor metadata, since the metadata merge
			// is additive and wouldn't remove the old key on its own.
			metadataUpdates[rivercommon.MetadataKeyResumableCursor] = nil
		}
		metadataUpdates[rivercommon.MetadataKeyResumableStep] = state.CompletedStep
	}

	return err
}

// ResumableState holds the state for a resumable job execution. It is stored in
// the context and accessed by ResumableStep and ResumableStepCursor.