From f6f05dfb931147a3621f6f8376375f84d7751e11 Mon Sep 17 00:00:00 2001
From: samliok <samliok@stanford.edu>
Date: Tue, 15 Sep 2026 14:42:38 -0400
Subject: [PATCH] expand is sequence scheduled

---
 common/block_scheduler.go      | 29 +++++++++++++++++------------
 common/block_scheduler_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/common/block_scheduler.go b/common/block_scheduler.go
index 85cd9a1c..06905fb2 100644
--- a/common/block_scheduler.go
+++ b/common/block_scheduler.go
@@ -30,8 +30,10 @@ type BlockDependencyManager struct {
 	scheduler Scheduler
 
 	dependencies []*TaskWithDependents
-	maxDeps      uint64
-	closed       atomic.Bool
+	// scheduledSeqs holds every seq with a task pending, queued or running.
+	scheduledSeqs map[uint64]struct{}
+	maxDeps       uint64
+	closed        atomic.Bool
 }
 
 type TaskWithDependents struct {
@@ -52,9 +54,10 @@ func (t *TaskWithDependents) String() string {
 
 func NewBlockVerificationScheduler(logger Logger, maxDeps uint64, scheduler Scheduler) *BlockDependencyManager {
 	b := &BlockDependencyManager{
-		logger:    logger,
-		maxDeps:   maxDeps,
-		scheduler: scheduler,
+		logger:        logger,
+		maxDeps:       maxDeps,
+		scheduler:     scheduler,
+		scheduledSeqs: make(map[uint64]struct{}),
 	}
 
 	b.logger.Debug("Created BlockVerificationScheduler", zap.Uint64("maxDeps", maxDeps))
@@ -121,17 +124,13 @@ func (bs *BlockDependencyManager) ExecuteEmptyRoundDependents(emptyRound uint64)
 	bs.dependencies = remainingDeps
 }
 
+// IsSequenceScheduled reports whether a task for seq is pending on dependencies, queued or running.
 func (bs *BlockDependencyManager) IsSequenceScheduled(seq uint64) bool {
 	bs.lock.Lock()
 	defer bs.lock.Unlock()
 
-	for _, dep := range bs.dependencies {
-		if dep.blockSeq == seq {
-			return true
-		}
-	}
-
-	return false
+	_, ok := bs.scheduledSeqs[seq]
+	return ok
 }
 
 func (bs *BlockDependencyManager) ScheduleTaskWithDependencies(task Task, blockSeq uint64, prev *Digest, emptyRounds []uint64) error {
@@ -144,6 +143,9 @@ func (bs *BlockDependencyManager) ScheduleTaskWithDependencies(task Task, blockS
 
 	wrappedTask := func() Digest {
 		id := task()
+		bs.lock.Lock()
+		delete(bs.scheduledSeqs, blockSeq)
+		bs.lock.Unlock()
 		bs.ExecuteBlockDependents(id)
 		return id
 	}
@@ -154,6 +156,8 @@ func (bs *BlockDependencyManager) ScheduleTaskWithDependencies(task Task, blockS
 		return fmt.Errorf("%w: %d pending verifications (max %d)", ErrTooManyPendingVerifications, totalSize, bs.maxDeps)
 	}
 
+	bs.scheduledSeqs[blockSeq] = struct{}{}
+
 	if prev == nil && len(emptyRounds) == 0 {
 		bs.logger.Debug("Scheduling block verification task with no dependencies", zap.Uint64("blockSeq", blockSeq))
 		bs.scheduler.Schedule(wrappedTask)
@@ -185,6 +189,7 @@ func (bs *BlockDependencyManager) RemoveOldTasks(seq uint64) {
 	for _, taskWithDeps := range bs.dependencies {
 		if taskWithDeps.blockSeq <= seq {
 			bs.logger.Debug("Removing block verification task as its block seq is less than or equal to finalized seq", zap.Uint64("blockSeq", taskWithDeps.blockSeq), zap.Uint64("finalizedSeq", seq))
+			delete(bs.scheduledSeqs, taskWithDeps.blockSeq)
 			continue
 		}
 		remainingDeps = append(remainingDeps, taskWithDeps)
diff --git a/common/block_scheduler_test.go b/common/block_scheduler_test.go
index 5ade6661..9f617403 100644
--- a/common/block_scheduler_test.go
+++ b/common/block_scheduler_test.go
@@ -223,6 +223,30 @@ func TestBlockVerificationScheduler(t *testing.T) {
 		waitReceive(t, done2)
 	})
 
+	t.Run("IsSequenceScheduled covers queued and running tasks", func(t *testing.T) {
+		scheduler := NewScheduler(noopLogger{}, defaultMaxDeps)
+		bvs := NewBlockVerificationScheduler(noopLogger{}, defaultMaxDeps, scheduler)
+		defer bvs.Close()
+
+		started := make(chan struct{}, 1)
+		release := make(chan struct{})
+		task := func() Digest {
+			started <- struct{}{}
+			<-release
+			return makeDigest(t)
+		}
+
+		require.False(t, bvs.IsSequenceScheduled(3))
+		require.NoError(t, bvs.ScheduleTaskWithDependencies(task, 3, nil, nil))
+
+		// Running with no dependencies still counts as scheduled.
+		waitReceive(t, started)
+		require.True(t, bvs.IsSequenceScheduled(3))
+
+		close(release)
+		require.Eventually(t, func() bool { return !bvs.IsSequenceScheduled(3) }, defaultWaitDuration, 10*time.Millisecond)
+	})
+
 	t.Run("RemoveOldTasks removes tasks with blockSeq <= finalized seq", func(t *testing.T) {
 		scheduler := NewScheduler(noopLogger{}, defaultMaxDeps)
 		bvs := NewBlockVerificationScheduler(noopLogger{}, defaultMaxDeps, scheduler)
@@ -257,6 +281,8 @@ func TestBlockVerificationScheduler(t *testing.T) {
 
 		// Finalize up to seq=6 — this should remove the old task (seq=5) but keep the new one (seq=8).
 		bvs.RemoveOldTasks(6)
+		require.False(t, bvs.IsSequenceScheduled(oldSeq))
+		require.True(t, bvs.IsSequenceScheduled(newSeq))
 
 		// Now resolve the dependency round. Only the "new" task should execute.
 		bvs.ExecuteEmptyRoundDependents(depRound)