feat(sync): split new journal requests from rescheduled

Rescheduled requests only traversed when at least one is ready.

Author: swartzn

rst/sync/internal/workmgr/manager.go

@@ -200,12 +200,23 @@ func NewAndStart(log *zap.Logger, config Config, beeRemoteClient *beeremote.Clie
 
 	allEntriesFound := 0
 	for priority := range priorityLevels {
-		entriesFound, err := m.initScheduler(SubmissionIdPriorityRange(priority))
+		lastSubmissionId, nextRescheduledTime, entriesFound, err := m.initScheduler(SubmissionIdPriorityRange(priority))
 		if err != nil {
 			m.log.Error("failed to initialize scheduler", zap.Error(err))
 			break
 		}
+
 		allEntriesFound += entriesFound
+		if !nextRescheduledTime.IsZero() {
+			m.scheduler.SetNextRescheduledTime(nextRescheduledTime, priority)
+		}
+		if lastSubmissionId != nil {
+			if nextSubmissionId, _, err := IncrementSubmissionId(*lastSubmissionId); err != nil {
+				m.log.Error("failed to initialize scheduler", zap.Error(err))
+			} else {
+				m.scheduler.SetNextSubmissionId(nextSubmissionId, priority)
+			}
+		}
 	}
 	if allEntriesFound > 0 {
 		m.log.Info("discovered work requests from previous run", zap.Int("requests", allEntriesFound))
@@ -285,7 +296,7 @@ func (m *Manager) manage(deferredFuncs []func() error) {
 			workJournal:          m.workJournal,
 			jobStore:             m.jobStore,
 			beeRemoteClient:      m.beeRemoteClient,
-			rescheduleWork:       m.scheduler.AddWorkToken,
+			rescheduleWork:       m.scheduler.RescheduleWork,
 		}
 		m.workerWG.Add(1)
 		go w.run(m.workerCtx, m.workerWG)
@@ -299,17 +310,33 @@ func (m *Manager) manage(deferredFuncs []func() error) {
 		case <-m.mgrCtx.Done():
 			return
 		case allowedTokens := <-m.scheduler.tokensReleased:
+			currentTime := time.Now()
 			for priority, ok := nextPriority(); ok; priority, ok = nextPriority() {
 				tokensDistributed := allowedTokens[priority]
 				if tokensDistributed == 0 {
 					continue
 				}
 
 				start, stop := SubmissionIdPriorityRange(priority)
-				err = m.pullInWork(start, stop, allowedTokens[priority])
-				if err != nil {
+				nextSubmissionId := m.scheduler.GetNextSubmissionId(priority)
+
+				// Add rescheduled work to the activeWork map
+				nextRescheduledTime := m.scheduler.GetNextRescheduledTime(priority)
+				if currentTime.After(nextRescheduledTime) {
+					if _, nextRescheduledTime, err = m.pullInWork(start, nextSubmissionId, &allowedTokens[priority]); err != nil {
+						m.log.Error("failed to pull in new work", zap.Error(err))
+						break
+					} else {
+						m.scheduler.SetNextRescheduledTime(nextRescheduledTime, priority)
+					}
+				}
+
+				// Add new work to the activeWork map
+				if next, _, err := m.pullInWork(nextSubmissionId, stop, &allowedTokens[priority]); err != nil {
 					m.log.Error("failed to pull in new work", zap.Error(err))
 					break
+				} else if next != nil {
+					m.scheduler.SetNextSubmissionId(*next, priority)
 				}
 			}
 		case completion := <-completedWork:
@@ -335,9 +362,13 @@ func (m *Manager) manage(deferredFuncs []func() error) {
 }
 
 // pullInWork moves ready work from the priority range to the activeWork map.
-func (m *Manager) pullInWork(start string, stop string, availableTokens int) error {
-	if availableTokens <= 0 {
-		return nil
+func (m *Manager) pullInWork(start string, stop string, availableTokens *int) (lastSubmissionId *string, nextExecuteAfter time.Time, err error) {
+	if availableTokens == nil {
+		err = fmt.Errorf("availableTokens was unexpectedly nil: this is a bug")
+		return
+	}
+	if *availableTokens <= 0 {
+		return
 	}
 
 	m.activeWorkMu.Lock()
@@ -348,26 +379,29 @@ func (m *Manager) pullInWork(start string, stop string, availableTokens int) err
 		kvstore.WithStopKey(stop),
 	)
 	if err != nil {
-		return fmt.Errorf("unable to get work journal entries: %w", err)
+		err = fmt.Errorf("unable to get work journal entries: %w", err)
+		return
 	}
 	defer cleanupNext()
 
 	item, err := nextItem()
 	if err != nil {
-		return fmt.Errorf("unable to get work journal entry: %w", err)
+		err = fmt.Errorf("unable to get work journal entry: %w", err)
+		return
 	}
 	if item == nil {
-		return nil
+		return
 	}
 
+	lastSubmissionId = new(string)
 	currentTime := time.Now()
-	for item != nil && availableTokens > 0 {
-		submissionID := item.Key
+	for item != nil && *availableTokens > 0 {
+		*lastSubmissionId = item.Key
 		entry := item.Entry.Value
 
 		if currentTime.After(entry.ExecuteAfter) {
 			workId := workIdentifier{
-				submissionID:  submissionID,
+				submissionID:  *lastSubmissionId,
 				jobID:         entry.WorkRequest.JobId,
 				workRequestID: entry.WorkRequest.RequestId,
 			}
@@ -380,11 +414,15 @@ func (m *Manager) pullInWork(start string, stop string, availableTokens int) err
 				activeWork := workAssignment{ctx: workCtx, workIdentifier: workId}
 				m.activeWork[activeWork.workIdentifier] = workContext{ctx: workCtx, cancel: workCtxCancel}
 				m.activeWorkQueue <- activeWork
-				availableTokens -= 1
-				m.scheduler.RemoveWorkToken(submissionID)
+				*availableTokens -= 1
+				m.scheduler.RemoveWorkToken(*lastSubmissionId)
 				priority := priorityIdMap[entry.WorkRequest.GetPriority()]
 				beeSyncActiveQueue.Add(priority, 1)
 			}
+		} else {
+			if nextExecuteAfter.IsZero() || entry.ExecuteAfter.Before(nextExecuteAfter) {
+				nextExecuteAfter = entry.ExecuteAfter
+			}
 		}
 
 		item, err = nextItem()
@@ -393,41 +431,49 @@ func (m *Manager) pullInWork(start string, stop string, availableTokens int) err
 		}
 	}
 
-	return err
+	return
 }
 
-// initScheduler adds tokens for unfinished work requests so they can be handled when the Sync
-// node starts.
-func (m *Manager) initScheduler(start string, stop string) (int, error) {
+// initScheduler adds tokens for unfinished work requests so they can be handled when the Sync node
+// starts. It returns the last entry's submissionId, entries found and an error if there was one.
+func (m *Manager) initScheduler(start string, stop string) (submissionId *string, nextExecuteAfter time.Time, entriesFound int, err error) {
 	nextItem, cleanupNext, err := m.workJournal.GetEntries(
 		kvstore.WithStartingKey(start),
 		kvstore.WithStopKey(stop),
 	)
 	if err != nil {
-		return 0, fmt.Errorf("unable to get work journal entries: %w", err)
+		err = fmt.Errorf("unable to get work journal entries: %w", err)
+		return
 	}
 	defer cleanupNext()
 
 	submission, err := nextItem()
 	if err != nil {
-		return 0, fmt.Errorf("unable to get work journal entry: %w", err)
+		err = fmt.Errorf("unable to get work journal entry: %w", err)
+		return
 	}
 	if submission == nil {
-		return 0, nil
+		return
 	}
 
-	entriesFound := 0
+	submissionId = new(string)
 	for submission != nil {
-		m.scheduler.AddWorkToken(submission.Key)
+		*submissionId = submission.Key
+		entry := submission.Entry.Value
+		m.scheduler.AddWorkToken(*submissionId)
 		entriesFound++
 
+		if nextExecuteAfter.IsZero() || entry.ExecuteAfter.Before(nextExecuteAfter) {
+			nextExecuteAfter = entry.ExecuteAfter
+		}
+
 		submission, err = nextItem()
 		if err != nil {
 			break
 		}
 	}
 
-	return entriesFound, err
+	return
 }
 
 func (m *Manager) SubmitWorkRequest(wr *flex.WorkRequest) (*flex.Work, error) {

rst/sync/internal/workmgr/scheduler.go

@@ -2,6 +2,7 @@ package workmgr
 
 import (
 	"context"
+	"fmt"
 	"math"
 	"strconv"
 	"sync/atomic"
@@ -60,7 +61,9 @@ type scheduler struct {
 	// tokensReleased is a buffered channel releases the priority tokens to the sync manager. This
 	// must be buffered to avoid deadlocking with manager. The channel must hold every sendWork that
 	// can occur before the manager loop reads m.scheduler.tokensReleased.
-	tokensReleased chan [priorityLevels]int
+	tokensReleased       chan [priorityLevels]int
+	nextRescheduledTimes [priorityLevels]*time.Time
+	nextSubmissionIds    [priorityLevels]string
 }
 
 func NewScheduler(ctx context.Context, log *zap.Logger, queue chan workAssignment, fairness gemetricRatio, opts ...schedulerOpt) (s *scheduler, close func() error) {
@@ -148,25 +151,53 @@ func NewScheduler(ctx context.Context, log *zap.Logger, queue chan workAssignmen
 	return
 }
 
+func (s *scheduler) GetNextRescheduledTime(priority int) time.Time {
+	if s.nextRescheduledTimes[priority] == nil {
+		return time.Time{}
+	}
+	return *s.nextRescheduledTimes[priority]
+}
+func (s *scheduler) SetNextRescheduledTime(ExecuteAfter time.Time, priority int) {
+	if s.nextRescheduledTimes[priority] == nil {
+		s.nextRescheduledTimes[priority] = new(time.Time)
+		*s.nextRescheduledTimes[priority] = ExecuteAfter
+	} else if ExecuteAfter.Before(*s.nextRescheduledTimes[priority]) {
+		*s.nextRescheduledTimes[priority] = ExecuteAfter
+	}
+}
+func (s *scheduler) GetNextSubmissionId(priority int) string {
+	return s.nextSubmissionIds[priority]
+}
+func (s *scheduler) SetNextSubmissionId(submissionId string, priority int) {
+	s.nextSubmissionIds[priority] = submissionId
+}
+
+func (s *scheduler) RescheduleWork(submissionId string, ExecuteAfter time.Time) {
+	priority := s.AddWorkToken(submissionId)
+	s.SetNextRescheduledTime(ExecuteAfter, int(priority))
+}
+
 // AddWorkToken(submissionID) tells the scheduler about a submission in the journal that is eligible
 // to be scheduled. It should be called whenever a WR is created, rediscovered on startup, or
 // rescheduled. The scheduler decodes the submission ID to determine the priority and increment that
 // bucket's token count. These counts are used to keep track of pending work at each priority to
 // ensure no priority queue is starved. Tokens represent work that is ready but not yet dispatched
 // to a worker.
-func (s *scheduler) AddWorkToken(submissionId string) {
+func (s *scheduler) AddWorkToken(submissionId string) int32 {
 	priority := submissionIdPriority(submissionId)
 	s.workTokens[priority].Add(1)
 	s.allWorkTokens.Add(1)
+	return priority
 }
 
 // RemoveWorkToken(submissionID) is called once a work assignment has been added to the active work
 // queue (not when it actually completes). This tells the scheduler a request at the given priority
 // has been handed to a worker, allowing it to internally adjust how it assigns new work.
-func (s *scheduler) RemoveWorkToken(submissionId string) {
+func (s *scheduler) RemoveWorkToken(submissionId string) int32 {
 	priority := submissionIdPriority(submissionId)
 	s.workTokens[priority].Add(-1)
 	s.allWorkTokens.Add(-1)
+	return priority
 }
 
 func (s *scheduler) getUpdateStatsFn(
@@ -346,8 +377,22 @@ func CreateSubmissionId(baseKey string, workRequestPriority int32) (string, int3
 	if priority < 0 || priority > priorityLevels-1 {
 		priority = 2
 	}
+
+	workRequestPriority = priority + 1
 	leadByte := baseKey[0] + submissionIdPriorityOffsetTable[priority]
-	return string(leadByte) + baseKey[1:], priority + 1
+	return string(leadByte) + baseKey[1:], workRequestPriority
+}
+
+func IncrementSubmissionId(key string) (string, int32, error) {
+	workRequestPriority := submissionIdPriority(key) + 1
+	value, err := strconv.ParseUint(submissionBaseKey(key), 36, 64)
+	if err != nil {
+		return "", 0, fmt.Errorf("unable to cast last submission ID to an integer '%s': %w", key, err)
+	}
+
+	baseKey := fmt.Sprintf("%013s", strconv.FormatUint(value+1, 36))
+	submissionId, priority := CreateSubmissionId(baseKey, workRequestPriority)
+	return submissionId, priority, nil
 }
 
 func DemoteSubmissionId(key string) (string, int32) {

rst/sync/internal/workmgr/work.go

@@ -130,7 +130,7 @@ type worker struct {
 	workJournal          *kvstore.MapStore[workEntry]
 	jobStore             *kvstore.MapStore[map[string]string]
 	beeRemoteClient      *beeremote.Client
-	rescheduleWork       func(string)
+	rescheduleWork       func(submissionId string, ExecuteAfter time.Time)
 }
 
 func (w *worker) run(ctx context.Context, wg *sync.WaitGroup) {
@@ -296,7 +296,7 @@ func (w *worker) processWork(work workAssignment) {
 		status.SetState(flex.Work_RESCHEDULED)
 		status.SetMessage("waiting for work request to be ready")
 		w.sendWorkResult(work, result.Work)
-		w.rescheduleWork(work.submissionID)
+		w.rescheduleWork(work.submissionID, journalEntry.Value.ExecuteAfter)
 		beeSyncWaitQueue.Add(mappedPriorityId, 1)
 		return
 	}