@@ -29,7 +29,7 @@ type kvJobManager struct {
2929// TODO(0x5459): Consider putting `txn` into context?
3030func (tm * kvJobManager ) filter (ctx context.Context , txn kvstore.TxnExt , state core.WdPoStJobState , limit uint32 , f func (* core.WdPoStJob ) bool ) (jobs []* core.WdPoStJob , err error ) {
3131 var it kvstore.Iter
32- it , err = txn .Scan ([] byte (makeWdPoStPrefix (state )))
32+ it , err = txn .Scan (kvstore . Prefix (makeWdPoStPrefix (state )))
3333 if err != nil {
3434 return
3535 }
@@ -134,7 +134,7 @@ func (tm *kvJobManager) Create(ctx context.Context, deadlineIdx uint64, partitio
134134 CreatedAt : uint64 (now ),
135135 UpdatedAt : uint64 (now ),
136136 }
137- return txn .PutJson ([] byte (makeWdPoStKey (core .WdPoStJobReadyToRun , jobID )), & job )
137+ return txn .PutJson (kvstore . Key (makeWdPoStKey (core .WdPoStJobReadyToRun , jobID )), & job )
138138 })
139139
140140 if err == nil {
@@ -163,7 +163,7 @@ func (tm *kvJobManager) AllocateJobs(ctx context.Context, spec core.AllocateWdPo
163163 now := uint64 (time .Now ().Unix ())
164164 for _ , job := range readyToRun {
165165 // Moving ready to run jobs to running jobs
166- if err := txn .Del ([] byte (makeWdPoStKey (core .WdPoStJobReadyToRun , job .ID ))); err != nil {
166+ if err := txn .Del (kvstore . Key (makeWdPoStKey (core .WdPoStJobReadyToRun , job .ID ))); err != nil {
167167 return err
168168 }
169169 job .State = string (core .WdPoStJobRunning )
@@ -172,7 +172,7 @@ func (tm *kvJobManager) AllocateJobs(ctx context.Context, spec core.AllocateWdPo
172172 job .WorkerName = workerName
173173 job .HeartbeatAt = now
174174 job .UpdatedAt = now
175- if err := txn .PutJson ([] byte (makeWdPoStKey (core .WdPoStJobRunning , job .ID )), job ); err != nil {
175+ if err := txn .PutJson (kvstore . Key (makeWdPoStKey (core .WdPoStJobRunning , job .ID )), job ); err != nil {
176176 return err
177177 }
178178 allocatedJobs = append (allocatedJobs , & core.WdPoStAllocatedJob {
@@ -196,16 +196,28 @@ func (tm *kvJobManager) Heartbeat(ctx context.Context, jobIDs []string, workerNa
196196 err := tm .kv .UpdateMustNoConflict (ctx , func (txn kvstore.TxnExt ) error {
197197 for _ , jobID := range jobIDs {
198198 var job core.WdPoStJob
199- if err := txn .Peek ([]byte (makeWdPoStKey (core .WdPoStJobRunning , jobID )), kvstore .LoadJSON (& job )); err != nil {
199+ key , err := txn .PeekAny (
200+ kvstore .LoadJSON (& job ),
201+ kvstore .Key (makeWdPoStKey (core .WdPoStJobRunning , jobID )),
202+ kvstore .Key (makeWdPoStKey (core .WdPoStJobReadyToRun , jobID )),
203+ kvstore .Key (makeWdPoStKey (core .WdPoStJobFinished , jobID )),
204+ )
205+ if err != nil {
206+ return err
207+ }
208+
209+ if err := txn .Del (key ); err != nil {
200210 return err
201211 }
212+
202213 if job .StartedAt == 0 {
203214 job .StartedAt = now
204215 }
205216 job .HeartbeatAt = now
206217 job .WorkerName = workerName
218+ job .State = string (core .WdPoStJobRunning )
207219 job .UpdatedAt = now
208- if err := txn .PutJson ([] byte (makeWdPoStKey (core .WdPoStJobRunning , jobID )), & job ); err != nil {
220+ if err := txn .PutJson (kvstore . Key (makeWdPoStKey (core .WdPoStJobRunning , jobID )), & job ); err != nil {
209221 return err
210222 }
211223 }
@@ -219,12 +231,17 @@ func (tm *kvJobManager) Heartbeat(ctx context.Context, jobIDs []string, workerNa
219231
220232func (tm * kvJobManager ) Finish (ctx context.Context , jobID string , output * stage.WindowPoStOutput , errorReason string ) error {
221233 err := tm .kv .UpdateMustNoConflict (ctx , func (txn kvstore.TxnExt ) error {
222- runningKey := []byte (makeWdPoStKey (core .WdPoStJobRunning , jobID ))
223234 var job core.WdPoStJob
224- if err := txn .Peek (runningKey , kvstore .LoadJSON (& job )); err != nil {
235+ key , err := txn .PeekAny (
236+ kvstore .LoadJSON (& job ),
237+ kvstore .Key (makeWdPoStKey (core .WdPoStJobRunning , jobID )),
238+ kvstore .Key (makeWdPoStKey (core .WdPoStJobReadyToRun , jobID )),
239+ kvstore .Key (makeWdPoStKey (core .WdPoStJobFinished , jobID )),
240+ )
241+ if err != nil {
225242 return err
226243 }
227- if err := txn .Del (runningKey ); err != nil {
244+ if err := txn .Del (key ); err != nil {
228245 return err
229246 }
230247 now := uint64 (time .Now ().Unix ())
@@ -233,7 +250,7 @@ func (tm *kvJobManager) Finish(ctx context.Context, jobID string, output *stage.
233250 job .ErrorReason = errorReason
234251 job .FinishedAt = now
235252 job .UpdatedAt = now
236- return txn .PutJson ([] byte (makeWdPoStKey (core .WdPoStJobFinished , jobID )), & job )
253+ return txn .PutJson (kvstore . Key (makeWdPoStKey (core .WdPoStJobFinished , jobID )), & job )
237254 })
238255
239256 if err == nil {
@@ -260,15 +277,15 @@ func (tm *kvJobManager) MakeJobsDie(ctx context.Context, heartbeatTimeout time.D
260277 }
261278 now := uint64 (time .Now ().Unix ())
262279 for _ , job := range shouldDead {
263- if err := txn .Del ([] byte (makeWdPoStKey (core .WdPoStJobRunning , job .ID ))); err != nil {
280+ if err := txn .Del (kvstore . Key (makeWdPoStKey (core .WdPoStJobRunning , job .ID ))); err != nil {
264281 return err
265282 }
266283 job .State = string (core .WdPoStJobFinished )
267284 job .FinishedAt = now
268285 job .Output = nil
269286 job .ErrorReason = "heartbeat timeout"
270287 job .UpdatedAt = now
271- if err := txn .PutJson ([] byte (makeWdPoStKey (core .WdPoStJobFinished , job .ID )), job ); err != nil {
288+ if err := txn .PutJson (kvstore . Key (makeWdPoStKey (core .WdPoStJobFinished , job .ID )), job ); err != nil {
272289 return err
273290 }
274291 }
@@ -297,7 +314,7 @@ func (tm *kvJobManager) CleanupExpiredJobs(ctx context.Context, jobLifetime time
297314 return err
298315 }
299316 for _ , job := range shouldClean {
300- if err := txn .Del ([] byte (makeWdPoStKey (core .WdPoStJobFinished , job .ID ))); err != nil {
317+ if err := txn .Del (kvstore . Key (makeWdPoStKey (core .WdPoStJobFinished , job .ID ))); err != nil {
301318 return err
302319 }
303320 }
@@ -324,7 +341,7 @@ func (tm *kvJobManager) RetryFailedJobs(ctx context.Context, maxTry, limit uint3
324341 }
325342 now := uint64 (time .Now ().Unix ())
326343 for _ , job := range shouldRetry {
327- err := txn .Del ([] byte (makeWdPoStKey (core .WdPoStJobFinished , job .ID )))
344+ err := txn .Del (kvstore . Key (makeWdPoStKey (core .WdPoStJobFinished , job .ID )))
328345 if err != nil {
329346 return err
330347 }
@@ -334,7 +351,7 @@ func (tm *kvJobManager) RetryFailedJobs(ctx context.Context, maxTry, limit uint3
334351 job .StartedAt = 0
335352 job .FinishedAt = 0
336353 job .UpdatedAt = now
337- if err := txn .PutJson ([] byte (makeWdPoStKey (core .WdPoStJobReadyToRun , job .ID )), job ); err != nil {
354+ if err := txn .PutJson (kvstore . Key (makeWdPoStKey (core .WdPoStJobReadyToRun , job .ID )), job ); err != nil {
338355 return err
339356 }
340357 }
@@ -379,7 +396,7 @@ func (tm *kvJobManager) Reset(ctx context.Context, jobID string) error {
379396 if err := txn .Del (key ); err != nil {
380397 return err
381398 }
382- return txn .PutJson ([] byte (makeWdPoStKey (core .WdPoStJobReadyToRun , jobID )), & job )
399+ return txn .PutJson (kvstore . Key (makeWdPoStKey (core .WdPoStJobReadyToRun , jobID )), & job )
383400 })
384401
385402 if err == nil {
0 commit comments