github.com/core-coin/go-core/v2@v2.1.9/les/lespay/client/fillset.go (about) 1 // Copyright 2020 by the Authors 2 // This file is part of the go-core library. 3 // 4 // The go-core 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-core 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-core library. If not, see <http://www.gnu.org/licenses/>. 16 17 package client 18 19 import ( 20 "sync" 21 22 "github.com/core-coin/go-core/v2/p2p/enode" 23 "github.com/core-coin/go-core/v2/p2p/nodestate" 24 ) 25 26 // FillSet tries to read nodes from an input iterator and add them to a node set by 27 // setting the specified node state flag(s) until the size of the set reaches the target. 28 // Note that other mechanisms (like other FillSet instances reading from different inputs) 29 // can also set the same flag(s) and FillSet will always care about the total number of 30 // nodes having those flags. 31 type FillSet struct { 32 lock sync.Mutex 33 cond *sync.Cond 34 ns *nodestate.NodeStateMachine 35 input enode.Iterator 36 closed bool 37 flags nodestate.Flags 38 count, target int 39 } 40 41 // NewFillSet creates a new FillSet 42 func NewFillSet(ns *nodestate.NodeStateMachine, input enode.Iterator, flags nodestate.Flags) *FillSet { 43 fs := &FillSet{ 44 ns: ns, 45 input: input, 46 flags: flags, 47 } 48 fs.cond = sync.NewCond(&fs.lock) 49 50 ns.SubscribeState(flags, func(n *enode.Node, oldState, newState nodestate.Flags) { 51 fs.lock.Lock() 52 if oldState.Equals(flags) { 53 fs.count-- 54 } 55 if newState.Equals(flags) { 56 fs.count++ 57 } 58 if fs.target > fs.count { 59 fs.cond.Signal() 60 } 61 fs.lock.Unlock() 62 }) 63 64 go fs.readLoop() 65 return fs 66 } 67 68 // readLoop keeps reading nodes from the input and setting the specified flags for them 69 // whenever the node set size is under the current target 70 func (fs *FillSet) readLoop() { 71 for { 72 fs.lock.Lock() 73 for fs.target <= fs.count && !fs.closed { 74 fs.cond.Wait() 75 } 76 77 fs.lock.Unlock() 78 if !fs.input.Next() { 79 return 80 } 81 fs.ns.SetState(fs.input.Node(), fs.flags, nodestate.Flags{}, 0) 82 } 83 } 84 85 // SetTarget sets the current target for node set size. If the previous target was not 86 // reached and FillSet was still waiting for the next node from the input then the next 87 // incoming node will be added to the set regardless of the target. This ensures that 88 // all nodes coming from the input are eventually added to the set. 89 func (fs *FillSet) SetTarget(target int) { 90 fs.lock.Lock() 91 defer fs.lock.Unlock() 92 93 fs.target = target 94 if fs.target > fs.count { 95 fs.cond.Signal() 96 } 97 } 98 99 // Close shuts FillSet down and closes the input iterator 100 func (fs *FillSet) Close() { 101 fs.lock.Lock() 102 defer fs.lock.Unlock() 103 104 fs.closed = true 105 fs.input.Close() 106 fs.cond.Signal() 107 }