@@ -43,6 +43,8 @@ type Pool struct {
4343 queueManager QueueManager
4444
4545 stopCh chan struct {}
46+ cancelCtx context.Context // Cancelled on Stop to interrupt in-flight store operations
47+ cancelFunc context.CancelFunc // Cancels cancelCtx
4648 wg sync.WaitGroup
4749 mu sync.Mutex
4850 running bool
@@ -125,6 +127,7 @@ func (p *Pool) Start(_ context.Context) error {
125127 return nil
126128 }
127129 p .running = true
130+ p .cancelCtx , p .cancelFunc = context .WithCancel (context .Background ())
128131
129132 p .logger .Info ("worker pool starting" ,
130133 log .String ("worker_id" , p .workerID .String ()),
@@ -165,8 +168,11 @@ func (p *Pool) Stop(ctx context.Context) error {
165168
166169 p .logger .Info ("worker pool stopping" , log .String ("worker_id" , p .workerID .String ()))
167170
168- // Signal all workers to stop.
171+ // Signal all workers to stop and cancel in-flight store operations .
169172 close (p .stopCh )
173+ if p .cancelFunc != nil {
174+ p .cancelFunc ()
175+ }
170176
171177 // Wait for completion or context deadline.
172178 done := make (chan struct {})
@@ -195,11 +201,16 @@ func (p *Pool) dequeueLoop() {
195201 select {
196202 case <- p .stopCh :
197203 return
204+ case <- p .cancelCtx .Done ():
205+ return
198206 default :
199207 }
200208
201- jobs , err := p .store .DequeueJobs (context . Background () , p .queues , 1 )
209+ jobs , err := p .store .DequeueJobs (p . cancelCtx , p .queues , 1 )
202210 if err != nil {
211+ if p .cancelCtx .Err () != nil {
212+ return // Clean exit during shutdown
213+ }
203214 p .logger .Error ("dequeue error" , log .String ("error" , err .Error ()))
204215 p .sleep ()
205216 continue
@@ -217,7 +228,10 @@ func (p *Pool) dequeueLoop() {
217228 // Rate limited — return the job to pending with a small delay.
218229 j .State = job .StatePending
219230 j .RunAt = time .Now ().Add (p .pollInterval )
220- if updateErr := p .store .UpdateJob (context .Background (), j ); updateErr != nil {
231+ if updateErr := p .store .UpdateJob (p .cancelCtx , j ); updateErr != nil {
232+ if p .cancelCtx .Err () != nil {
233+ return
234+ }
221235 p .logger .Error ("failed to re-enqueue rate-limited job" ,
222236 log .String ("job_id" , j .ID .String ()),
223237 log .String ("error" , updateErr .Error ()),
@@ -227,9 +241,9 @@ func (p *Pool) dequeueLoop() {
227241 continue
228242 }
229243
230- p .extensions .EmitJobStarted (context . Background () , j )
244+ p .extensions .EmitJobStarted (p . cancelCtx , j )
231245
232- ctx , cancel := context .WithCancel (context . Background () )
246+ ctx , cancel := context .WithCancel (p . cancelCtx )
233247 p .trackJob (j .ID .String (), cancel )
234248
235249 execErr := p .executor .Execute (ctx , j )
@@ -282,7 +296,7 @@ func (p *Pool) sendHeartbeats() {
282296 p .logger .Warn ("heartbeat: invalid job id" , log .String ("job_id" , jobIDStr ))
283297 continue
284298 }
285- if err := p .store .HeartbeatJob (context . Background () , parsedID , p .workerID ); err != nil {
299+ if err := p .store .HeartbeatJob (p . cancelCtx , parsedID , p .workerID ); err != nil {
286300 p .logger .Warn ("heartbeat failed" ,
287301 log .String ("job_id" , jobIDStr ),
288302 log .String ("error" , err .Error ()),
@@ -309,7 +323,7 @@ func (p *Pool) reaperLoop() {
309323}
310324
311325func (p * Pool ) reapStaleJobs () {
312- stale , err := p .store .ReapStaleJobs (context . Background () , p .staleJobThreshold )
326+ stale , err := p .store .ReapStaleJobs (p . cancelCtx , p .staleJobThreshold )
313327 if err != nil {
314328 p .logger .Error ("reap stale jobs error" , log .String ("error" , err .Error ()))
315329 return
@@ -322,7 +336,7 @@ func (p *Pool) reapStaleJobs() {
322336 j .HeartbeatAt = nil
323337 j .StartedAt = nil
324338
325- if updateErr := p .store .UpdateJob (context . Background () , j ); updateErr != nil {
339+ if updateErr := p .store .UpdateJob (p . cancelCtx , j ); updateErr != nil {
326340 p .logger .Error ("reap: failed to reset stale job" ,
327341 log .String ("job_id" , j .ID .String ()),
328342 log .String ("error" , updateErr .Error ()),
@@ -341,6 +355,7 @@ func (p *Pool) sleep() {
341355 select {
342356 case <- time .After (p .pollInterval ):
343357 case <- p .stopCh :
358+ case <- p .cancelCtx .Done ():
344359 }
345360}
346361
0 commit comments