op-geth/les/servingqueue.go

// Copyright 2018 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

package les

import (
	"sync"

	"github.com/ethereum/go-ethereum/common/mclock"
	"github.com/ethereum/go-ethereum/common/prque"
)

// servingQueue allows running tasks in a limited number of threads and puts the
// waiting tasks in a priority queue
type servingQueue struct {
	tokenCh                 chan runToken
	queueAddCh, queueBestCh chan *servingTask
	stopThreadCh, quit      chan struct{}
	setThreadsCh            chan int

	wg          sync.WaitGroup
	threadCount int          // number of currently running threads
	queue       *prque.Prque // priority queue for waiting or suspended tasks
	best        *servingTask // the highest priority task (not included in the queue)
	suspendBias int64        // priority bias against suspending an already running task
}

// servingTask represents a request serving task. Tasks can be implemented to
// run in multiple steps, allowing the serving queue to suspend execution between
// steps if higher priority tasks are entered. The creator of the task should
// set the following fields:
//
// - priority: greater value means higher priority; values can wrap around the int64 range
// - run: execute a single step; return true if finished
// - after: executed after run finishes or returns an error, receives the total serving time
type servingTask struct {
	sq          *servingQueue
	servingTime uint64
	priority    int64
	biasAdded   bool
	token       runToken
	tokenCh     chan runToken
}

// runToken received by servingTask.start allows the task to run. Closing the
// channel by servingTask.stop signals the thread controller to allow a new task
// to start running.
type runToken chan struct{}

// start blocks until the task can start and returns true if it is allowed to run.
// Returning false means that the task should be cancelled.
func (t *servingTask) start() bool {
	select {
	case t.token = <-t.sq.tokenCh:
	default:
		t.tokenCh = make(chan runToken, 1)
		select {
		case t.sq.queueAddCh <- t:
		case <-t.sq.quit:
			return false
		}
		select {
		case t.token = <-t.tokenCh:
		case <-t.sq.quit:
			return false
		}
	}
	if t.token == nil {
		return false
	}
	t.servingTime -= uint64(mclock.Now())
	return true
}

// done signals the thread controller about the task being finished and returns
// the total serving time of the task in nanoseconds.
func (t *servingTask) done() uint64 {
	t.servingTime += uint64(mclock.Now())
	close(t.token)
	return t.servingTime
}

// waitOrStop can be called during the execution of the task. It blocks if there
// is a higher priority task waiting (a bias is applied in favor of the currently
// running task). Returning true means that the execution can be resumed. False
// means the task should be cancelled.
func (t *servingTask) waitOrStop() bool {
	t.done()
	if !t.biasAdded {
		t.priority += t.sq.suspendBias
		t.biasAdded = true
	}
	return t.start()
}

// newServingQueue returns a new servingQueue
func newServingQueue(suspendBias int64) *servingQueue {
	sq := &servingQueue{
		queue:        prque.New(nil),
		suspendBias:  suspendBias,
		tokenCh:      make(chan runToken),
		queueAddCh:   make(chan *servingTask, 100),
		queueBestCh:  make(chan *servingTask),
		stopThreadCh: make(chan struct{}),
		quit:         make(chan struct{}),
		setThreadsCh: make(chan int, 10),
	}
	sq.wg.Add(2)
	go sq.queueLoop()
	go sq.threadCountLoop()
	return sq
}

// newTask creates a new task with the given priority
func (sq *servingQueue) newTask(priority int64) *servingTask {
	return &servingTask{
		sq:       sq,
		priority: priority,
	}
}

// threadController is started in multiple goroutines and controls the execution
// of tasks. The number of active thread controllers equals the allowed number of
// concurrently running threads. It tries to fetch the highest priority queued
// task first. If there are no queued tasks waiting then it can directly catch
// run tokens from the token channel and allow the corresponding tasks to run
// without entering the priority queue.
func (sq *servingQueue) threadController() {
	for {
		token := make(runToken)
		select {
		case best := <-sq.queueBestCh:
			best.tokenCh <- token
		default:
			select {
			case best := <-sq.queueBestCh:
				best.tokenCh <- token
			case sq.tokenCh <- token:
			case <-sq.stopThreadCh:
				sq.wg.Done()
				return
			case <-sq.quit:
				sq.wg.Done()
				return
			}
		}
		<-token
		select {
		case <-sq.stopThreadCh:
			sq.wg.Done()
			return
		case <-sq.quit:
			sq.wg.Done()
			return
		default:
		}
	}
}

// addTask inserts a task into the priority queue
func (sq *servingQueue) addTask(task *servingTask) {
	if sq.best == nil {
		sq.best = task
	} else if task.priority > sq.best.priority {
		sq.queue.Push(sq.best, sq.best.priority)
		sq.best = task
		return
	} else {
		sq.queue.Push(task, task.priority)
	}
}

// queueLoop is an event loop running in a goroutine. It receives tasks from queueAddCh
// and always tries to send the highest priority task to queueBestCh. Successfully sent
// tasks are removed from the queue.
func (sq *servingQueue) queueLoop() {
	for {
		if sq.best != nil {
			select {
			case task := <-sq.queueAddCh:
				sq.addTask(task)
			case sq.queueBestCh <- sq.best:
				if sq.queue.Size() == 0 {
					sq.best = nil
				} else {
					sq.best, _ = sq.queue.PopItem().(*servingTask)
				}
			case <-sq.quit:
				sq.wg.Done()
				return
			}
		} else {
			select {
			case task := <-sq.queueAddCh:
				sq.addTask(task)
			case <-sq.quit:
				sq.wg.Done()
				return
			}
		}
	}
}

// threadCountLoop is an event loop running in a goroutine. It adjusts the number
// of active thread controller goroutines.
func (sq *servingQueue) threadCountLoop() {
	var threadCountTarget int
	for {
		for threadCountTarget > sq.threadCount {
			sq.wg.Add(1)
			go sq.threadController()
			sq.threadCount++
		}
		if threadCountTarget < sq.threadCount {
			select {
			case threadCountTarget = <-sq.setThreadsCh:
			case sq.stopThreadCh <- struct{}{}:
				sq.threadCount--
			case <-sq.quit:
				sq.wg.Done()
				return
			}
		} else {
			select {
			case threadCountTarget = <-sq.setThreadsCh:
			case <-sq.quit:
				sq.wg.Done()
				return
			}
		}
	}
}

// setThreads sets the allowed processing thread count, suspending tasks as soon as
// possible if necessary.
func (sq *servingQueue) setThreads(threadCount int) {
	select {
	case sq.setThreadsCh <- threadCount:
	case <-sq.quit:
		return
	}
}

// stop stops task processing as soon as possible and shuts down the serving queue.
func (sq *servingQueue) stop() {
	close(sq.quit)
	sq.wg.Wait()
}