@@ -2,8 +2,16 @@ package queue
22
33import (
44 "context"
5+ "encoding/json"
56 "errors"
7+ "fmt"
68 "sync"
9+
10+ "github.com/jackc/pgx/v5"
11+ "github.com/jackc/pgx/v5/pgxpool"
12+ "github.com/riverqueue/river"
13+ "github.com/riverqueue/river/riverdriver/riverpgxv5"
14+ "github.com/riverqueue/river/rivermigrate"
715)
816
917// RiverConfig configures the River (Postgres-backed) Queue backend.
@@ -12,7 +20,7 @@ type RiverConfig struct {
1220 // engine — no additional infrastructure required.
1321 DatabaseURL string
1422
15- // NumWorkers is the size of the worker pool that drains jobs .
23+ // NumWorkers is the max concurrent jobs drained from the default queue .
1624 NumWorkers int
1725}
1826
@@ -22,45 +30,181 @@ type RiverConfig struct {
2230// reuses the Postgres database the engine already depends on. No Redis, no
2331// new ops surface.
2432//
25- // The actual River client wiring will be added when we integrate the
26- // `river` + `riverpgxv5` packages. For scaffolding we keep the shape of the
27- // implementation so the rest of the engine can depend on it.
33+ // Implementation shape: River expects typed job args per kind, but the
34+ // engine's Queue interface speaks a single generic Job shape. We bridge by
35+ // registering a single River worker for one "envelope" kind. The envelope
36+ // carries the engine's JobKind + JSON payload; the worker dispatches to the
37+ // registered handler at work time.
2838type River struct {
29- cfg RiverConfig
39+ cfg RiverConfig
40+ pool * pgxpool.Pool
41+
3042 mu sync.RWMutex
3143 handlers map [JobKind ]Handler
44+
45+ client * river.Client [pgx.Tx ]
46+ once sync.Once
47+ initEr error
48+ }
49+
50+ // envelopeArgs is the single job type we register with River. The real kind
51+ // lives in DomainKind; River only sees "vle_envelope".
52+ type envelopeArgs struct {
53+ DomainKind JobKind `json:"domain_kind"`
54+ Payload json.RawMessage `json:"payload"`
55+ }
56+
57+ // Kind is River's own kind identifier. It must be stable and unique within
58+ // the database; it intentionally differs from the engine's JobKind space so
59+ // the two can't collide if River ever grows a second typed job.
60+ func (envelopeArgs ) Kind () string { return "vle_envelope" }
61+
62+ // envelopeWorker dispatches an envelope to the Handler registered on the
63+ // owning *River for that DomainKind.
64+ type envelopeWorker struct {
65+ river.WorkerDefaults [envelopeArgs ]
66+ q * River
67+ }
68+
69+ func (w * envelopeWorker ) Work (ctx context.Context , job * river.Job [envelopeArgs ]) error {
70+ w .q .mu .RLock ()
71+ h , ok := w .q .handlers [job .Args .DomainKind ]
72+ w .q .mu .RUnlock ()
73+ if ! ok {
74+ return fmt .Errorf ("%w: %q" , ErrUnknownKind , job .Args .DomainKind )
75+ }
76+ return h (ctx , Job {Kind : job .Args .DomainKind , Payload : job .Args .Payload })
3277}
3378
34- // NewRiver constructs a new River-backed Queue.
79+ // NewRiver constructs a new River-backed Queue. It opens its own pgxpool
80+ // against cfg.DatabaseURL (intentionally separate from the engine's main
81+ // pool so queue traffic and data-plane traffic can't starve each other) and
82+ // runs River's schema migrations idempotently at startup.
3583func NewRiver (cfg RiverConfig ) (* River , error ) {
3684 if cfg .DatabaseURL == "" {
3785 return nil , errors .New ("river: database_url is required" )
3886 }
3987 if cfg .NumWorkers <= 0 {
4088 cfg .NumWorkers = 10
4189 }
90+
91+ // We open the pool eagerly so NewRiver surfaces DB misconfiguration at
92+ // boot rather than on the first enqueue. The pool is closed by Close().
93+ ctx , cancel := context .WithCancel (context .Background ())
94+ defer cancel ()
95+
96+ pool , err := pgxpool .New (ctx , cfg .DatabaseURL )
97+ if err != nil {
98+ return nil , fmt .Errorf ("river: open pool: %w" , err )
99+ }
100+
101+ // Apply River's own migrations (creates river_job + related tables).
102+ // This is idempotent and fast when up-to-date.
103+ migrator , err := rivermigrate .New (riverpgxv5 .New (pool ), nil )
104+ if err != nil {
105+ pool .Close ()
106+ return nil , fmt .Errorf ("river: migrator: %w" , err )
107+ }
108+ if _ , err := migrator .Migrate (ctx , rivermigrate .DirectionUp , nil ); err != nil {
109+ pool .Close ()
110+ return nil , fmt .Errorf ("river: migrate: %w" , err )
111+ }
112+
42113 return & River {
43114 cfg : cfg ,
115+ pool : pool ,
44116 handlers : map [JobKind ]Handler {},
45117 }, nil
46118}
47119
120+ // Register binds a handler to a domain JobKind. It's safe to call before or
121+ // after the client starts — the envelope worker looks handlers up at work
122+ // time, not at registration time.
48123func (r * River ) Register (kind JobKind , h Handler ) {
49124 r .mu .Lock ()
50125 defer r .mu .Unlock ()
51126 r .handlers [kind ] = h
52127}
53128
129+ // Enqueue inserts a job. If the client hasn't been built yet (no prior
130+ // Enqueue or Start call), this triggers lazy client construction.
54131func (r * River ) Enqueue (ctx context.Context , j Job ) error {
55- // TODO(phase-1): insert into river_job via riverpgxv5.Client.Insert.
56- return errors .New ("river: Enqueue not yet implemented" )
132+ if err := r .ensureClient (); err != nil {
133+ return err
134+ }
135+ args := envelopeArgs {DomainKind : j .Kind , Payload : j .Payload }
136+
137+ opts := & river.InsertOpts {}
138+ if ! j .RunAt .IsZero () {
139+ opts .ScheduledAt = j .RunAt
140+ }
141+ if j .MaxRetries > 0 {
142+ opts .MaxAttempts = j .MaxRetries
143+ }
144+ // DedupeKey intentionally left on the table; River's native uniqueness
145+ // story has moved through several shapes (UniqueOpts, ByArgs, etc.) and
146+ // wiring it in here should be a deliberate call, not a guess. When we
147+ // need it we'll plumb UniqueOpts with ByKey matching DedupeKey.
148+
149+ if _ , err := r .client .Insert (ctx , args , opts ); err != nil {
150+ return fmt .Errorf ("river: insert: %w" , err )
151+ }
152+ return nil
57153}
58154
155+ // Start begins draining jobs. Blocks until ctx is cancelled, then performs a
156+ // graceful shutdown (finishes in-flight jobs, stops fetching new ones).
59157func (r * River ) Start (ctx context.Context ) error {
60- // TODO(phase-1): start river.Client workers. Each registered JobKind
61- // maps to a river.Worker[T] that calls the handler.
158+ if err := r .ensureClient (); err != nil {
159+ return err
160+ }
161+ if err := r .client .Start (ctx ); err != nil {
162+ return fmt .Errorf ("river: start: %w" , err )
163+ }
62164 <- ctx .Done ()
165+ // Give in-flight jobs a bounded window to finish. We detach from ctx
166+ // (which is already done) so Stop doesn't immediately cancel what it's
167+ // trying to drain.
168+ stopCtx , cancel := context .WithCancel (context .Background ())
169+ defer cancel ()
170+ if err := r .client .Stop (stopCtx ); err != nil && ! errors .Is (err , context .Canceled ) {
171+ return fmt .Errorf ("river: stop: %w" , err )
172+ }
173+ return nil
174+ }
175+
176+ // Close releases the pool. Safe to call after Start returns. It does not
177+ // stop in-flight work — Start handles graceful shutdown on ctx cancel.
178+ func (r * River ) Close () error {
179+ if r .pool != nil {
180+ r .pool .Close ()
181+ }
63182 return nil
64183}
65184
66- func (r * River ) Close () error { return nil }
185+ // ensureClient lazily constructs the river.Client. We defer this until first
186+ // use so tests that just want to Register handlers (or engines that queue
187+ // without working) can construct a *River without spinning up workers.
188+ //
189+ // The client's Workers are registered once at build time — our single
190+ // envelopeWorker dispatches by looking up the handler map on each Work call,
191+ // so new Register() calls after client construction still take effect.
192+ func (r * River ) ensureClient () error {
193+ r .once .Do (func () {
194+ workers := river .NewWorkers ()
195+ river .AddWorker (workers , & envelopeWorker {q : r })
196+
197+ client , err := river .NewClient [pgx.Tx ](riverpgxv5 .New (r .pool ), & river.Config {
198+ Workers : workers ,
199+ Queues : map [string ]river.QueueConfig {
200+ river .QueueDefault : {MaxWorkers : r .cfg .NumWorkers },
201+ },
202+ })
203+ if err != nil {
204+ r .initEr = fmt .Errorf ("river: new client: %w" , err )
205+ return
206+ }
207+ r .client = client
208+ })
209+ return r .initEr
210+ }
0 commit comments