bft/common/block_scheduler.go

// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package common

import (
	"errors"
	"fmt"
	"maps"
	"slices"
	"sync"
	"sync/atomic"

	"github.com/luxfi/log"
)

var ErrTooManyPendingVerifications = errors.New("too many blocks being verified to ingest another one")

type Scheduler interface {
	Schedule(task Task)
	Size() int
	Close()
}

// BlockDependencyManager manages block verification tasks with dependencies on previous blocks and empty rounds.
// It schedules tasks when their dependencies are resolved.
type BlockDependencyManager struct {
	lock      sync.Mutex
	logger    Logger
	scheduler Scheduler

	dependencies []*TaskWithDependents
	maxDeps      uint64
	closed       atomic.Bool
}

type TaskWithDependents struct {
	Task Task

	blockSeq    uint64 // the seq of the block being verified
	prevBlock   *Digest
	emptyRounds map[uint64]struct{}
}

func (t *TaskWithDependents) isReady() bool {
	return t.prevBlock == nil && len(t.emptyRounds) == 0
}

func (t *TaskWithDependents) String() string {
	return fmt.Sprintf("BlockVerificationTask{blockSeq: %d, prevBlock: %v, emptyRounds: %v}", t.blockSeq, t.prevBlock, slices.Collect(maps.Keys(t.emptyRounds)))
}

func NewBlockVerificationScheduler(logger Logger, maxDeps uint64, scheduler Scheduler) *BlockDependencyManager {
	b := &BlockDependencyManager{
		logger:    logger,
		maxDeps:   maxDeps,
		scheduler: scheduler,
	}

	b.logger.Debug("Created BlockVerificationScheduler", log.Uint64("maxDeps", maxDeps))
	return b
}

// ExecuteBlockDependents removes the given digest from dependent tasks and schedules any whose dependencies are now resolved.
func (bs *BlockDependencyManager) ExecuteBlockDependents(prev Digest) {
	bs.lock.Lock()
	defer bs.lock.Unlock()

	remainingDeps := make([]*TaskWithDependents, 0, len(bs.dependencies))

	for _, taskWithDeps := range bs.dependencies {
		if taskWithDeps.prevBlock != nil && *taskWithDeps.prevBlock == prev {
			taskWithDeps.prevBlock = nil
		}

		if taskWithDeps.isReady() {
			bs.logger.Debug("Scheduling block verification task as all dependencies are met", log.Stringer("taskID", prev), log.Uint64("Seq", taskWithDeps.blockSeq))
			bs.scheduler.Schedule(taskWithDeps.Task)
			continue
		}

		remainingDeps = append(remainingDeps, taskWithDeps)
	}

	if len(remainingDeps) > 0 {
		bs.logger.Debug("Block scheduler has unsatisfied dependencies",
			log.Stringer("prev", prev),
			log.Stringers("task", remainingDeps),
		)
	}

	bs.dependencies = remainingDeps
}

// ExecuteEmptyRoundDependents removes the given empty round from dependent tasks and schedules any whose dependencies are now resolved.
func (bs *BlockDependencyManager) ExecuteEmptyRoundDependents(emptyRound uint64) {
	bs.lock.Lock()
	defer bs.lock.Unlock()

	remainingDeps := make([]*TaskWithDependents, 0, len(bs.dependencies))

	for _, taskWithDeps := range bs.dependencies {
		delete(taskWithDeps.emptyRounds, emptyRound)

		if taskWithDeps.isReady() {
			bs.logger.Debug("Scheduling block verification task as all dependencies are met", log.Stringer("task", taskWithDeps))
			bs.scheduler.Schedule(taskWithDeps.Task)
			continue
		}

		remainingDeps = append(remainingDeps, taskWithDeps)
	}

	if len(remainingDeps) > 0 {
		bs.logger.Debug("Block scheduler has unsatisfied dependencies",
			log.Uint64("empty round", emptyRound),
			log.Stringers("task", remainingDeps),
		)
	}

	bs.dependencies = remainingDeps
}

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
}

func (bs *BlockDependencyManager) ScheduleTaskWithDependencies(task Task, blockSeq uint64, prev *Digest, emptyRounds []uint64) error {
	bs.lock.Lock()
	defer bs.lock.Unlock()

	if bs.closed.Load() {
		return nil
	}

	wrappedTask := func() Digest {
		id := task()
		bs.ExecuteBlockDependents(id)
		return id
	}

	totalSize := uint64(len(bs.dependencies) + bs.scheduler.Size())
	if totalSize >= bs.maxDeps {
		bs.logger.Warn("Too many blocks being verified to ingest another one", log.Uint64("pendingBlocks", totalSize))
		return fmt.Errorf("%w: %d pending verifications (max %d)", ErrTooManyPendingVerifications, totalSize, bs.maxDeps)
	}

	if prev == nil && len(emptyRounds) == 0 {
		bs.logger.Debug("Scheduling block verification task with no dependencies", log.Uint64("blockSeq", blockSeq))
		bs.scheduler.Schedule(wrappedTask)
		return nil
	}

	bs.logger.Debug("Adding block verification task with dependencies", log.Any("prevBlock", prev), log.Uint64s("emptyRounds", emptyRounds), log.Uint64("blockSeq", blockSeq))
	emptyRoundsSet := make(map[uint64]struct{})
	for _, round := range emptyRounds {
		emptyRoundsSet[round] = struct{}{}
	}

	bs.dependencies = append(bs.dependencies, &TaskWithDependents{
		Task:        wrappedTask,
		prevBlock:   prev,
		emptyRounds: emptyRoundsSet,
		blockSeq:    blockSeq,
	})

	return nil
}

// We can remove all tasks that have an empty notarization dependency for a round that has been finalized.
func (bs *BlockDependencyManager) RemoveOldTasks(seq uint64) {
	bs.lock.Lock()
	defer bs.lock.Unlock()

	remainingDeps := make([]*TaskWithDependents, 0, len(bs.dependencies))
	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", log.Uint64("blockSeq", taskWithDeps.blockSeq), log.Uint64("finalizedSeq", seq))
			continue
		}
		remainingDeps = append(remainingDeps, taskWithDeps)
	}

	bs.dependencies = remainingDeps
}

func (bs *BlockDependencyManager) Close() {
	bs.closed.Store(true)
	bs.scheduler.Close()
}