github.com/jimmyx0x/go-ethereum@v1.10.28/les/vflux/client/queueiterator.go (about)

     1  // Copyright 2020 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package client
    18  
    19  import (
    20  	"sync"
    21  
    22  	"github.com/ethereum/go-ethereum/p2p/enode"
    23  	"github.com/ethereum/go-ethereum/p2p/nodestate"
    24  )
    25  
    26  // QueueIterator returns nodes from the specified selectable set in the same order as
    27  // they entered the set.
    28  type QueueIterator struct {
    29  	lock sync.Mutex
    30  	cond *sync.Cond
    31  
    32  	ns           *nodestate.NodeStateMachine
    33  	queue        []*enode.Node
    34  	nextNode     *enode.Node
    35  	waitCallback func(bool)
    36  	fifo, closed bool
    37  }
    38  
    39  // NewQueueIterator creates a new QueueIterator. Nodes are selectable if they have all the required
    40  // and none of the disabled flags set. When a node is selected the selectedFlag is set which also
    41  // disables further selectability until it is removed or times out.
    42  func NewQueueIterator(ns *nodestate.NodeStateMachine, requireFlags, disableFlags nodestate.Flags, fifo bool, waitCallback func(bool)) *QueueIterator {
    43  	qi := &QueueIterator{
    44  		ns:           ns,
    45  		fifo:         fifo,
    46  		waitCallback: waitCallback,
    47  	}
    48  	qi.cond = sync.NewCond(&qi.lock)
    49  
    50  	ns.SubscribeState(requireFlags.Or(disableFlags), func(n *enode.Node, oldState, newState nodestate.Flags) {
    51  		oldMatch := oldState.HasAll(requireFlags) && oldState.HasNone(disableFlags)
    52  		newMatch := newState.HasAll(requireFlags) && newState.HasNone(disableFlags)
    53  		if newMatch == oldMatch {
    54  			return
    55  		}
    56  
    57  		qi.lock.Lock()
    58  		defer qi.lock.Unlock()
    59  
    60  		if newMatch {
    61  			qi.queue = append(qi.queue, n)
    62  		} else {
    63  			id := n.ID()
    64  			for i, qn := range qi.queue {
    65  				if qn.ID() == id {
    66  					copy(qi.queue[i:len(qi.queue)-1], qi.queue[i+1:])
    67  					qi.queue = qi.queue[:len(qi.queue)-1]
    68  					break
    69  				}
    70  			}
    71  		}
    72  		qi.cond.Signal()
    73  	})
    74  	return qi
    75  }
    76  
    77  // Next moves to the next selectable node.
    78  func (qi *QueueIterator) Next() bool {
    79  	qi.lock.Lock()
    80  	if !qi.closed && len(qi.queue) == 0 {
    81  		if qi.waitCallback != nil {
    82  			qi.waitCallback(true)
    83  		}
    84  		for !qi.closed && len(qi.queue) == 0 {
    85  			qi.cond.Wait()
    86  		}
    87  		if qi.waitCallback != nil {
    88  			qi.waitCallback(false)
    89  		}
    90  	}
    91  	if qi.closed {
    92  		qi.nextNode = nil
    93  		qi.lock.Unlock()
    94  		return false
    95  	}
    96  	// Move to the next node in queue.
    97  	if qi.fifo {
    98  		qi.nextNode = qi.queue[0]
    99  		copy(qi.queue[:len(qi.queue)-1], qi.queue[1:])
   100  		qi.queue = qi.queue[:len(qi.queue)-1]
   101  	} else {
   102  		qi.nextNode = qi.queue[len(qi.queue)-1]
   103  		qi.queue = qi.queue[:len(qi.queue)-1]
   104  	}
   105  	qi.lock.Unlock()
   106  	return true
   107  }
   108  
   109  // Close ends the iterator.
   110  func (qi *QueueIterator) Close() {
   111  	qi.lock.Lock()
   112  	qi.closed = true
   113  	qi.lock.Unlock()
   114  	qi.cond.Signal()
   115  }
   116  
   117  // Node returns the current node.
   118  func (qi *QueueIterator) Node() *enode.Node {
   119  	qi.lock.Lock()
   120  	defer qi.lock.Unlock()
   121  
   122  	return qi.nextNode
   123  }