clientpool.go 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871
  1. // Copyright 2019 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. package les
  17. import (
  18. "bytes"
  19. "encoding/binary"
  20. "fmt"
  21. "io"
  22. "math"
  23. "sync"
  24. "time"
  25. "github.com/ethereum/go-ethereum/common"
  26. "github.com/ethereum/go-ethereum/common/mclock"
  27. "github.com/ethereum/go-ethereum/common/prque"
  28. "github.com/ethereum/go-ethereum/ethdb"
  29. "github.com/ethereum/go-ethereum/log"
  30. "github.com/ethereum/go-ethereum/p2p/enode"
  31. "github.com/ethereum/go-ethereum/rlp"
  32. lru "github.com/hashicorp/golang-lru"
  33. )
  34. const (
  35. negBalanceExpTC = time.Hour // time constant for exponentially reducing negative balance
  36. fixedPointMultiplier = 0x1000000 // constant to convert logarithms to fixed point format
  37. lazyQueueRefresh = time.Second * 10 // refresh period of the connected queue
  38. persistCumulativeTimeRefresh = time.Minute * 5 // refresh period of the cumulative running time persistence
  39. posBalanceCacheLimit = 8192 // the maximum number of cached items in positive balance queue
  40. negBalanceCacheLimit = 8192 // the maximum number of cached items in negative balance queue
  41. defaultConnectedBias = time.Minute * 3 // the default connectedBias used in clientPool
  42. )
  43. // clientPool implements a client database that assigns a priority to each client
  44. // based on a positive and negative balance. Positive balance is externally assigned
  45. // to prioritized clients and is decreased with connection time and processed
  46. // requests (unless the price factors are zero). If the positive balance is zero
  47. // then negative balance is accumulated.
  48. //
  49. // Balance tracking and priority calculation for connected clients is done by
  50. // balanceTracker. connectedQueue ensures that clients with the lowest positive or
  51. // highest negative balance get evicted when the total capacity allowance is full
  52. // and new clients with a better balance want to connect.
  53. //
  54. // Already connected nodes receive a small bias in their favor in order to avoid
  55. // accepting and instantly kicking out clients. In theory, we try to ensure that
  56. // each client can have several minutes of connection time.
  57. //
  58. // Balances of disconnected clients are stored in nodeDB including positive balance
  59. // and negative banalce. Negative balance is transformed into a logarithmic form
  60. // with a constantly shifting linear offset in order to implement an exponential
  61. // decrease. Besides nodeDB will have a background thread to check the negative
  62. // balance of disconnected client. If the balance is low enough, then the record
  63. // will be dropped.
  64. type clientPool struct {
  65. ndb *nodeDB
  66. lock sync.Mutex
  67. clock mclock.Clock
  68. stopCh chan struct{}
  69. closed bool
  70. removePeer func(enode.ID)
  71. connectedMap map[enode.ID]*clientInfo
  72. connectedQueue *prque.LazyQueue
  73. defaultPosFactors, defaultNegFactors priceFactors
  74. connLimit int // The maximum number of connections that clientpool can support
  75. capLimit uint64 // The maximum cumulative capacity that clientpool can support
  76. connectedCap uint64 // The sum of the capacity of the current clientpool connected
  77. priorityConnected uint64 // The sum of the capacity of currently connected priority clients
  78. freeClientCap uint64 // The capacity value of each free client
  79. startTime mclock.AbsTime // The timestamp at which the clientpool started running
  80. cumulativeTime int64 // The cumulative running time of clientpool at the start point.
  81. connectedBias time.Duration // The connection bias. 0: Disable connection bias(used in testing)
  82. }
  83. // clientPoolPeer represents a client peer in the pool.
  84. // Positive balances are assigned to node key while negative balances are assigned
  85. // to freeClientId. Currently network IP address without port is used because
  86. // clients have a limited access to IP addresses while new node keys can be easily
  87. // generated so it would be useless to assign a negative value to them.
  88. type clientPoolPeer interface {
  89. ID() enode.ID
  90. freeClientId() string
  91. updateCapacity(uint64)
  92. freezeClient()
  93. }
  94. // clientInfo represents a connected client
  95. type clientInfo struct {
  96. address string
  97. id enode.ID
  98. connectedAt mclock.AbsTime
  99. capacity uint64
  100. priority bool
  101. pool *clientPool
  102. peer clientPoolPeer
  103. queueIndex int // position in connectedQueue
  104. balanceTracker balanceTracker
  105. posFactors, negFactors priceFactors
  106. balanceMetaInfo string
  107. }
  108. // connSetIndex callback updates clientInfo item index in connectedQueue
  109. func connSetIndex(a interface{}, index int) {
  110. a.(*clientInfo).queueIndex = index
  111. }
  112. // connPriority callback returns actual priority of clientInfo item in connectedQueue
  113. func connPriority(a interface{}, now mclock.AbsTime) int64 {
  114. c := a.(*clientInfo)
  115. return c.balanceTracker.getPriority(now)
  116. }
  117. // connMaxPriority callback returns estimated maximum priority of clientInfo item in connectedQueue
  118. func connMaxPriority(a interface{}, until mclock.AbsTime) int64 {
  119. c := a.(*clientInfo)
  120. pri := c.balanceTracker.estimatedPriority(until, true)
  121. c.balanceTracker.addCallback(balanceCallbackQueue, pri+1, func() {
  122. c.pool.lock.Lock()
  123. if c.queueIndex != -1 {
  124. c.pool.connectedQueue.Update(c.queueIndex)
  125. }
  126. c.pool.lock.Unlock()
  127. })
  128. return pri
  129. }
  130. // priceFactors determine the pricing policy (may apply either to positive or
  131. // negative balances which may have different factors).
  132. // - timeFactor is cost unit per nanosecond of connection time
  133. // - capacityFactor is cost unit per nanosecond of connection time per 1000000 capacity
  134. // - requestFactor is cost unit per request "realCost" unit
  135. type priceFactors struct {
  136. timeFactor, capacityFactor, requestFactor float64
  137. }
  138. // newClientPool creates a new client pool
  139. func newClientPool(db ethdb.Database, freeClientCap uint64, clock mclock.Clock, removePeer func(enode.ID)) *clientPool {
  140. ndb := newNodeDB(db, clock)
  141. pool := &clientPool{
  142. ndb: ndb,
  143. clock: clock,
  144. connectedMap: make(map[enode.ID]*clientInfo),
  145. connectedQueue: prque.NewLazyQueue(connSetIndex, connPriority, connMaxPriority, clock, lazyQueueRefresh),
  146. freeClientCap: freeClientCap,
  147. removePeer: removePeer,
  148. startTime: clock.Now(),
  149. cumulativeTime: ndb.getCumulativeTime(),
  150. stopCh: make(chan struct{}),
  151. connectedBias: defaultConnectedBias,
  152. }
  153. // If the negative balance of free client is even lower than 1,
  154. // delete this entry.
  155. ndb.nbEvictCallBack = func(now mclock.AbsTime, b negBalance) bool {
  156. balance := math.Exp(float64(b.logValue-pool.logOffset(now)) / fixedPointMultiplier)
  157. return balance <= 1
  158. }
  159. go func() {
  160. for {
  161. select {
  162. case <-clock.After(lazyQueueRefresh):
  163. pool.lock.Lock()
  164. pool.connectedQueue.Refresh()
  165. pool.lock.Unlock()
  166. case <-clock.After(persistCumulativeTimeRefresh):
  167. pool.ndb.setCumulativeTime(pool.logOffset(clock.Now()))
  168. case <-pool.stopCh:
  169. return
  170. }
  171. }
  172. }()
  173. return pool
  174. }
  175. // stop shuts the client pool down
  176. func (f *clientPool) stop() {
  177. close(f.stopCh)
  178. f.lock.Lock()
  179. f.closed = true
  180. f.lock.Unlock()
  181. f.ndb.setCumulativeTime(f.logOffset(f.clock.Now()))
  182. f.ndb.close()
  183. }
  184. // connect should be called after a successful handshake. If the connection was
  185. // rejected, there is no need to call disconnect.
  186. func (f *clientPool) connect(peer clientPoolPeer, capacity uint64) bool {
  187. f.lock.Lock()
  188. defer f.lock.Unlock()
  189. // Short circuit if clientPool is already closed.
  190. if f.closed {
  191. return false
  192. }
  193. // Dedup connected peers.
  194. id, freeID := peer.ID(), peer.freeClientId()
  195. if _, ok := f.connectedMap[id]; ok {
  196. clientRejectedMeter.Mark(1)
  197. log.Debug("Client already connected", "address", freeID, "id", id.String())
  198. return false
  199. }
  200. // Create a clientInfo but do not add it yet
  201. var (
  202. posBalance uint64
  203. negBalance uint64
  204. now = f.clock.Now()
  205. )
  206. pb := f.ndb.getOrNewPB(id)
  207. posBalance = pb.value
  208. nb := f.ndb.getOrNewNB(freeID)
  209. if nb.logValue != 0 {
  210. negBalance = uint64(math.Exp(float64(nb.logValue-f.logOffset(now))/fixedPointMultiplier) * float64(time.Second))
  211. }
  212. e := &clientInfo{
  213. pool: f,
  214. peer: peer,
  215. address: freeID,
  216. queueIndex: -1,
  217. id: id,
  218. connectedAt: now,
  219. priority: posBalance != 0,
  220. posFactors: f.defaultPosFactors,
  221. negFactors: f.defaultNegFactors,
  222. balanceMetaInfo: pb.meta,
  223. }
  224. // If the client is a free client, assign with a low free capacity,
  225. // Otherwise assign with the given value(priority client)
  226. if !e.priority || capacity == 0 {
  227. capacity = f.freeClientCap
  228. }
  229. e.capacity = capacity
  230. // Starts a balance tracker
  231. e.balanceTracker.init(f.clock, capacity)
  232. e.balanceTracker.setBalance(posBalance, negBalance)
  233. e.updatePriceFactors()
  234. // If the number of clients already connected in the clientpool exceeds its
  235. // capacity, evict some clients with lowest priority.
  236. //
  237. // If the priority of the newly added client is lower than the priority of
  238. // all connected clients, the client is rejected.
  239. newCapacity := f.connectedCap + capacity
  240. newCount := f.connectedQueue.Size() + 1
  241. if newCapacity > f.capLimit || newCount > f.connLimit {
  242. var (
  243. kickList []*clientInfo
  244. kickPriority int64
  245. )
  246. f.connectedQueue.MultiPop(func(data interface{}, priority int64) bool {
  247. c := data.(*clientInfo)
  248. kickList = append(kickList, c)
  249. kickPriority = priority
  250. newCapacity -= c.capacity
  251. newCount--
  252. return newCapacity > f.capLimit || newCount > f.connLimit
  253. })
  254. if newCapacity > f.capLimit || newCount > f.connLimit || (e.balanceTracker.estimatedPriority(now+mclock.AbsTime(f.connectedBias), false)-kickPriority) > 0 {
  255. for _, c := range kickList {
  256. f.connectedQueue.Push(c)
  257. }
  258. clientRejectedMeter.Mark(1)
  259. log.Debug("Client rejected", "address", freeID, "id", id.String())
  260. return false
  261. }
  262. // accept new client, drop old ones
  263. for _, c := range kickList {
  264. f.dropClient(c, now, true)
  265. }
  266. }
  267. // Register new client to connection queue.
  268. f.connectedMap[id] = e
  269. f.connectedQueue.Push(e)
  270. f.connectedCap += e.capacity
  271. // If the current client is a paid client, monitor the status of client,
  272. // downgrade it to normal client if positive balance is used up.
  273. if e.priority {
  274. f.priorityConnected += capacity
  275. e.balanceTracker.addCallback(balanceCallbackZero, 0, func() { f.balanceExhausted(id) })
  276. }
  277. // If the capacity of client is not the default value(free capacity), notify
  278. // it to update capacity.
  279. if e.capacity != f.freeClientCap {
  280. e.peer.updateCapacity(e.capacity)
  281. }
  282. totalConnectedGauge.Update(int64(f.connectedCap))
  283. clientConnectedMeter.Mark(1)
  284. log.Debug("Client accepted", "address", freeID)
  285. return true
  286. }
  287. // disconnect should be called when a connection is terminated. If the disconnection
  288. // was initiated by the pool itself using disconnectFn then calling disconnect is
  289. // not necessary but permitted.
  290. func (f *clientPool) disconnect(p clientPoolPeer) {
  291. f.lock.Lock()
  292. defer f.lock.Unlock()
  293. // Short circuit if client pool is already closed.
  294. if f.closed {
  295. return
  296. }
  297. // Short circuit if the peer hasn't been registered.
  298. e := f.connectedMap[p.ID()]
  299. if e == nil {
  300. log.Debug("Client not connected", "address", p.freeClientId(), "id", p.ID().String())
  301. return
  302. }
  303. f.dropClient(e, f.clock.Now(), false)
  304. }
  305. // forClients iterates through a list of clients, calling the callback for each one.
  306. // If a client is not connected then clientInfo is nil. If the specified list is empty
  307. // then the callback is called for all connected clients.
  308. func (f *clientPool) forClients(ids []enode.ID, callback func(*clientInfo, enode.ID) error) error {
  309. f.lock.Lock()
  310. defer f.lock.Unlock()
  311. if len(ids) > 0 {
  312. for _, id := range ids {
  313. if err := callback(f.connectedMap[id], id); err != nil {
  314. return err
  315. }
  316. }
  317. } else {
  318. for _, c := range f.connectedMap {
  319. if err := callback(c, c.id); err != nil {
  320. return err
  321. }
  322. }
  323. }
  324. return nil
  325. }
  326. // setDefaultFactors sets the default price factors applied to subsequently connected clients
  327. func (f *clientPool) setDefaultFactors(posFactors, negFactors priceFactors) {
  328. f.lock.Lock()
  329. defer f.lock.Unlock()
  330. f.defaultPosFactors = posFactors
  331. f.defaultNegFactors = negFactors
  332. }
  333. // setConnectedBias sets the connection bias, which is applied to already connected clients
  334. // So that already connected client won't be kicked out very soon and we can ensure all
  335. // connected clients can have enough time to request or sync some data.
  336. func (f *clientPool) setConnectedBias(bias time.Duration) {
  337. f.lock.Lock()
  338. defer f.lock.Unlock()
  339. f.connectedBias = bias
  340. }
  341. // dropClient removes a client from the connected queue and finalizes its balance.
  342. // If kick is true then it also initiates the disconnection.
  343. func (f *clientPool) dropClient(e *clientInfo, now mclock.AbsTime, kick bool) {
  344. if _, ok := f.connectedMap[e.id]; !ok {
  345. return
  346. }
  347. f.finalizeBalance(e, now)
  348. f.connectedQueue.Remove(e.queueIndex)
  349. delete(f.connectedMap, e.id)
  350. f.connectedCap -= e.capacity
  351. if e.priority {
  352. f.priorityConnected -= e.capacity
  353. }
  354. totalConnectedGauge.Update(int64(f.connectedCap))
  355. if kick {
  356. clientKickedMeter.Mark(1)
  357. log.Debug("Client kicked out", "address", e.address)
  358. f.removePeer(e.id)
  359. } else {
  360. clientDisconnectedMeter.Mark(1)
  361. log.Debug("Client disconnected", "address", e.address)
  362. }
  363. }
  364. // capacityInfo returns the total capacity allowance, the total capacity of connected
  365. // clients and the total capacity of connected and prioritized clients
  366. func (f *clientPool) capacityInfo() (uint64, uint64, uint64) {
  367. f.lock.Lock()
  368. defer f.lock.Unlock()
  369. return f.capLimit, f.connectedCap, f.priorityConnected
  370. }
  371. // finalizeBalance stops the balance tracker, retrieves the final balances and
  372. // stores them in posBalanceQueue and negBalanceQueue
  373. func (f *clientPool) finalizeBalance(c *clientInfo, now mclock.AbsTime) {
  374. c.balanceTracker.stop(now)
  375. pos, neg := c.balanceTracker.getBalance(now)
  376. pb, nb := f.ndb.getOrNewPB(c.id), f.ndb.getOrNewNB(c.address)
  377. pb.value = pos
  378. f.ndb.setPB(c.id, pb)
  379. neg /= uint64(time.Second) // Convert the expanse to second level.
  380. if neg > 1 {
  381. nb.logValue = int64(math.Log(float64(neg))*fixedPointMultiplier) + f.logOffset(now)
  382. f.ndb.setNB(c.address, nb)
  383. } else {
  384. f.ndb.delNB(c.address) // Negative balance is small enough, drop it directly.
  385. }
  386. }
  387. // balanceExhausted callback is called by balanceTracker when positive balance is exhausted.
  388. // It revokes priority status and also reduces the client capacity if necessary.
  389. func (f *clientPool) balanceExhausted(id enode.ID) {
  390. f.lock.Lock()
  391. defer f.lock.Unlock()
  392. c := f.connectedMap[id]
  393. if c == nil || !c.priority {
  394. return
  395. }
  396. if c.priority {
  397. f.priorityConnected -= c.capacity
  398. }
  399. c.priority = false
  400. if c.capacity != f.freeClientCap {
  401. f.connectedCap += f.freeClientCap - c.capacity
  402. totalConnectedGauge.Update(int64(f.connectedCap))
  403. c.capacity = f.freeClientCap
  404. c.balanceTracker.setCapacity(c.capacity)
  405. c.peer.updateCapacity(c.capacity)
  406. }
  407. pb := f.ndb.getOrNewPB(id)
  408. pb.value = 0
  409. f.ndb.setPB(id, pb)
  410. }
  411. // setConnLimit sets the maximum number and total capacity of connected clients,
  412. // dropping some of them if necessary.
  413. func (f *clientPool) setLimits(totalConn int, totalCap uint64) {
  414. f.lock.Lock()
  415. defer f.lock.Unlock()
  416. f.connLimit = totalConn
  417. f.capLimit = totalCap
  418. if f.connectedCap > f.capLimit || f.connectedQueue.Size() > f.connLimit {
  419. f.connectedQueue.MultiPop(func(data interface{}, priority int64) bool {
  420. f.dropClient(data.(*clientInfo), mclock.Now(), true)
  421. return f.connectedCap > f.capLimit || f.connectedQueue.Size() > f.connLimit
  422. })
  423. }
  424. }
  425. // setCapacity sets the assigned capacity of a connected client
  426. func (f *clientPool) setCapacity(c *clientInfo, capacity uint64) error {
  427. if f.connectedMap[c.id] != c {
  428. return fmt.Errorf("client %064x is not connected", c.id[:])
  429. }
  430. if c.capacity == capacity {
  431. return nil
  432. }
  433. if !c.priority {
  434. return errNoPriority
  435. }
  436. oldCapacity := c.capacity
  437. c.capacity = capacity
  438. f.connectedCap += capacity - oldCapacity
  439. c.balanceTracker.setCapacity(capacity)
  440. f.connectedQueue.Update(c.queueIndex)
  441. if f.connectedCap > f.capLimit {
  442. var kickList []*clientInfo
  443. kick := true
  444. f.connectedQueue.MultiPop(func(data interface{}, priority int64) bool {
  445. client := data.(*clientInfo)
  446. kickList = append(kickList, client)
  447. f.connectedCap -= client.capacity
  448. if client == c {
  449. kick = false
  450. }
  451. return kick && (f.connectedCap > f.capLimit)
  452. })
  453. if kick {
  454. now := mclock.Now()
  455. for _, c := range kickList {
  456. f.dropClient(c, now, true)
  457. }
  458. } else {
  459. c.capacity = oldCapacity
  460. c.balanceTracker.setCapacity(oldCapacity)
  461. for _, c := range kickList {
  462. f.connectedCap += c.capacity
  463. f.connectedQueue.Push(c)
  464. }
  465. return errNoPriority
  466. }
  467. }
  468. totalConnectedGauge.Update(int64(f.connectedCap))
  469. f.priorityConnected += capacity - oldCapacity
  470. c.updatePriceFactors()
  471. c.peer.updateCapacity(c.capacity)
  472. return nil
  473. }
  474. // requestCost feeds request cost after serving a request from the given peer.
  475. func (f *clientPool) requestCost(p *clientPeer, cost uint64) {
  476. f.lock.Lock()
  477. defer f.lock.Unlock()
  478. info, exist := f.connectedMap[p.ID()]
  479. if !exist || f.closed {
  480. return
  481. }
  482. info.balanceTracker.requestCost(cost)
  483. }
  484. // logOffset calculates the time-dependent offset for the logarithmic
  485. // representation of negative balance
  486. //
  487. // From another point of view, the result returned by the function represents
  488. // the total time that the clientpool is cumulatively running(total_hours/multiplier).
  489. func (f *clientPool) logOffset(now mclock.AbsTime) int64 {
  490. // Note: fixedPointMultiplier acts as a multiplier here; the reason for dividing the divisor
  491. // is to avoid int64 overflow. We assume that int64(negBalanceExpTC) >> fixedPointMultiplier.
  492. cumulativeTime := int64((time.Duration(now - f.startTime)) / (negBalanceExpTC / fixedPointMultiplier))
  493. return f.cumulativeTime + cumulativeTime
  494. }
  495. // setClientPriceFactors sets the pricing factors for an individual connected client
  496. func (c *clientInfo) updatePriceFactors() {
  497. c.balanceTracker.setFactors(true, c.negFactors.timeFactor+float64(c.capacity)*c.negFactors.capacityFactor/1000000, c.negFactors.requestFactor)
  498. c.balanceTracker.setFactors(false, c.posFactors.timeFactor+float64(c.capacity)*c.posFactors.capacityFactor/1000000, c.posFactors.requestFactor)
  499. }
  500. // getPosBalance retrieves a single positive balance entry from cache or the database
  501. func (f *clientPool) getPosBalance(id enode.ID) posBalance {
  502. f.lock.Lock()
  503. defer f.lock.Unlock()
  504. return f.ndb.getOrNewPB(id)
  505. }
  506. // addBalance updates the balance of a client (either overwrites it or adds to it).
  507. // It also updates the balance meta info string.
  508. func (f *clientPool) addBalance(id enode.ID, amount int64, meta string) (uint64, uint64, error) {
  509. f.lock.Lock()
  510. defer f.lock.Unlock()
  511. pb := f.ndb.getOrNewPB(id)
  512. var negBalance uint64
  513. c := f.connectedMap[id]
  514. if c != nil {
  515. pb.value, negBalance = c.balanceTracker.getBalance(f.clock.Now())
  516. }
  517. oldBalance := pb.value
  518. if amount > 0 {
  519. if amount > maxBalance || pb.value > maxBalance-uint64(amount) {
  520. return oldBalance, oldBalance, errBalanceOverflow
  521. }
  522. pb.value += uint64(amount)
  523. } else {
  524. if uint64(-amount) > pb.value {
  525. pb.value = 0
  526. } else {
  527. pb.value -= uint64(-amount)
  528. }
  529. }
  530. pb.meta = meta
  531. f.ndb.setPB(id, pb)
  532. if c != nil {
  533. c.balanceTracker.setBalance(pb.value, negBalance)
  534. if !c.priority && pb.value > 0 {
  535. // The capacity should be adjusted based on the requirement,
  536. // but we have no idea about the new capacity, need a second
  537. // call to update it.
  538. c.priority = true
  539. f.priorityConnected += c.capacity
  540. c.balanceTracker.addCallback(balanceCallbackZero, 0, func() { f.balanceExhausted(id) })
  541. }
  542. // if balance is set to zero then reverting to non-priority status
  543. // is handled by the balanceExhausted callback
  544. c.balanceMetaInfo = meta
  545. }
  546. return oldBalance, pb.value, nil
  547. }
  548. // posBalance represents a recently accessed positive balance entry
  549. type posBalance struct {
  550. value uint64
  551. meta string
  552. }
  553. // EncodeRLP implements rlp.Encoder
  554. func (e *posBalance) EncodeRLP(w io.Writer) error {
  555. return rlp.Encode(w, []interface{}{e.value, e.meta})
  556. }
  557. // DecodeRLP implements rlp.Decoder
  558. func (e *posBalance) DecodeRLP(s *rlp.Stream) error {
  559. var entry struct {
  560. Value uint64
  561. Meta string
  562. }
  563. if err := s.Decode(&entry); err != nil {
  564. return err
  565. }
  566. e.value = entry.Value
  567. e.meta = entry.Meta
  568. return nil
  569. }
  570. // negBalance represents a negative balance entry of a disconnected client
  571. type negBalance struct{ logValue int64 }
  572. // EncodeRLP implements rlp.Encoder
  573. func (e *negBalance) EncodeRLP(w io.Writer) error {
  574. return rlp.Encode(w, []interface{}{uint64(e.logValue)})
  575. }
  576. // DecodeRLP implements rlp.Decoder
  577. func (e *negBalance) DecodeRLP(s *rlp.Stream) error {
  578. var entry struct {
  579. LogValue uint64
  580. }
  581. if err := s.Decode(&entry); err != nil {
  582. return err
  583. }
  584. e.logValue = int64(entry.LogValue)
  585. return nil
  586. }
  587. const (
  588. // nodeDBVersion is the version identifier of the node data in db
  589. //
  590. // Changelog:
  591. // * Replace `lastTotal` with `meta` in positive balance: version 0=>1
  592. nodeDBVersion = 1
  593. // dbCleanupCycle is the cycle of db for useless data cleanup
  594. dbCleanupCycle = time.Hour
  595. )
  596. var (
  597. positiveBalancePrefix = []byte("pb:") // dbVersion(uint16 big endian) + positiveBalancePrefix + id -> balance
  598. negativeBalancePrefix = []byte("nb:") // dbVersion(uint16 big endian) + negativeBalancePrefix + ip -> balance
  599. cumulativeRunningTimeKey = []byte("cumulativeTime:") // dbVersion(uint16 big endian) + cumulativeRunningTimeKey -> cumulativeTime
  600. )
  601. type nodeDB struct {
  602. db ethdb.Database
  603. pcache *lru.Cache
  604. ncache *lru.Cache
  605. auxbuf []byte // 37-byte auxiliary buffer for key encoding
  606. verbuf [2]byte // 2-byte auxiliary buffer for db version
  607. nbEvictCallBack func(mclock.AbsTime, negBalance) bool // Callback to determine whether the negative balance can be evicted.
  608. clock mclock.Clock
  609. closeCh chan struct{}
  610. cleanupHook func() // Test hook used for testing
  611. }
  612. func newNodeDB(db ethdb.Database, clock mclock.Clock) *nodeDB {
  613. pcache, _ := lru.New(posBalanceCacheLimit)
  614. ncache, _ := lru.New(negBalanceCacheLimit)
  615. ndb := &nodeDB{
  616. db: db,
  617. pcache: pcache,
  618. ncache: ncache,
  619. auxbuf: make([]byte, 37),
  620. clock: clock,
  621. closeCh: make(chan struct{}),
  622. }
  623. binary.BigEndian.PutUint16(ndb.verbuf[:], uint16(nodeDBVersion))
  624. go ndb.expirer()
  625. return ndb
  626. }
  627. func (db *nodeDB) close() {
  628. close(db.closeCh)
  629. }
  630. func (db *nodeDB) getPrefix(neg bool) []byte {
  631. prefix := positiveBalancePrefix
  632. if neg {
  633. prefix = negativeBalancePrefix
  634. }
  635. return append(db.verbuf[:], prefix...)
  636. }
  637. func (db *nodeDB) key(id []byte, neg bool) []byte {
  638. prefix := positiveBalancePrefix
  639. if neg {
  640. prefix = negativeBalancePrefix
  641. }
  642. if len(prefix)+len(db.verbuf)+len(id) > len(db.auxbuf) {
  643. db.auxbuf = append(db.auxbuf, make([]byte, len(prefix)+len(db.verbuf)+len(id)-len(db.auxbuf))...)
  644. }
  645. copy(db.auxbuf[:len(db.verbuf)], db.verbuf[:])
  646. copy(db.auxbuf[len(db.verbuf):len(db.verbuf)+len(prefix)], prefix)
  647. copy(db.auxbuf[len(prefix)+len(db.verbuf):len(prefix)+len(db.verbuf)+len(id)], id)
  648. return db.auxbuf[:len(prefix)+len(db.verbuf)+len(id)]
  649. }
  650. func (db *nodeDB) getCumulativeTime() int64 {
  651. blob, err := db.db.Get(append(cumulativeRunningTimeKey, db.verbuf[:]...))
  652. if err != nil || len(blob) == 0 {
  653. return 0
  654. }
  655. return int64(binary.BigEndian.Uint64(blob))
  656. }
  657. func (db *nodeDB) setCumulativeTime(v int64) {
  658. binary.BigEndian.PutUint64(db.auxbuf[:8], uint64(v))
  659. db.db.Put(append(cumulativeRunningTimeKey, db.verbuf[:]...), db.auxbuf[:8])
  660. }
  661. func (db *nodeDB) getOrNewPB(id enode.ID) posBalance {
  662. key := db.key(id.Bytes(), false)
  663. item, exist := db.pcache.Get(string(key))
  664. if exist {
  665. return item.(posBalance)
  666. }
  667. var balance posBalance
  668. if enc, err := db.db.Get(key); err == nil {
  669. if err := rlp.DecodeBytes(enc, &balance); err != nil {
  670. log.Error("Failed to decode positive balance", "err", err)
  671. }
  672. }
  673. db.pcache.Add(string(key), balance)
  674. return balance
  675. }
  676. func (db *nodeDB) setPB(id enode.ID, b posBalance) {
  677. if b.value == 0 && len(b.meta) == 0 {
  678. db.delPB(id)
  679. return
  680. }
  681. key := db.key(id.Bytes(), false)
  682. enc, err := rlp.EncodeToBytes(&(b))
  683. if err != nil {
  684. log.Error("Failed to encode positive balance", "err", err)
  685. return
  686. }
  687. db.db.Put(key, enc)
  688. db.pcache.Add(string(key), b)
  689. }
  690. func (db *nodeDB) delPB(id enode.ID) {
  691. key := db.key(id.Bytes(), false)
  692. db.db.Delete(key)
  693. db.pcache.Remove(string(key))
  694. }
  695. // getPosBalanceIDs returns a lexicographically ordered list of IDs of accounts
  696. // with a positive balance
  697. func (db *nodeDB) getPosBalanceIDs(start, stop enode.ID, maxCount int) (result []enode.ID) {
  698. if maxCount <= 0 {
  699. return
  700. }
  701. prefix := db.getPrefix(false)
  702. it := db.db.NewIterator(prefix, start.Bytes())
  703. defer it.Release()
  704. for i := len(stop[:]) - 1; i >= 0; i-- {
  705. stop[i]--
  706. if stop[i] != 255 {
  707. break
  708. }
  709. }
  710. stopKey := db.key(stop.Bytes(), false)
  711. keyLen := len(stopKey)
  712. for it.Next() {
  713. var id enode.ID
  714. if len(it.Key()) != keyLen || bytes.Compare(it.Key(), stopKey) == 1 {
  715. return
  716. }
  717. copy(id[:], it.Key()[keyLen-len(id):])
  718. result = append(result, id)
  719. if len(result) == maxCount {
  720. return
  721. }
  722. }
  723. return
  724. }
  725. func (db *nodeDB) getOrNewNB(id string) negBalance {
  726. key := db.key([]byte(id), true)
  727. item, exist := db.ncache.Get(string(key))
  728. if exist {
  729. return item.(negBalance)
  730. }
  731. var balance negBalance
  732. if enc, err := db.db.Get(key); err == nil {
  733. if err := rlp.DecodeBytes(enc, &balance); err != nil {
  734. log.Error("Failed to decode negative balance", "err", err)
  735. }
  736. }
  737. db.ncache.Add(string(key), balance)
  738. return balance
  739. }
  740. func (db *nodeDB) setNB(id string, b negBalance) {
  741. key := db.key([]byte(id), true)
  742. enc, err := rlp.EncodeToBytes(&(b))
  743. if err != nil {
  744. log.Error("Failed to encode negative balance", "err", err)
  745. return
  746. }
  747. db.db.Put(key, enc)
  748. db.ncache.Add(string(key), b)
  749. }
  750. func (db *nodeDB) delNB(id string) {
  751. key := db.key([]byte(id), true)
  752. db.db.Delete(key)
  753. db.ncache.Remove(string(key))
  754. }
  755. func (db *nodeDB) expirer() {
  756. for {
  757. select {
  758. case <-db.clock.After(dbCleanupCycle):
  759. db.expireNodes()
  760. case <-db.closeCh:
  761. return
  762. }
  763. }
  764. }
  765. // expireNodes iterates the whole node db and checks whether the negative balance
  766. // entry can deleted.
  767. //
  768. // The rationale behind this is: server doesn't need to keep the negative balance
  769. // records if they are low enough.
  770. func (db *nodeDB) expireNodes() {
  771. var (
  772. visited int
  773. deleted int
  774. start = time.Now()
  775. prefix = db.getPrefix(true)
  776. )
  777. iter := db.db.NewIterator(prefix, nil)
  778. for iter.Next() {
  779. visited += 1
  780. var balance negBalance
  781. if err := rlp.DecodeBytes(iter.Value(), &balance); err != nil {
  782. log.Error("Failed to decode negative balance", "err", err)
  783. continue
  784. }
  785. if db.nbEvictCallBack != nil && db.nbEvictCallBack(db.clock.Now(), balance) {
  786. deleted += 1
  787. db.db.Delete(iter.Key())
  788. }
  789. }
  790. // Invoke testing hook if it's not nil.
  791. if db.cleanupHook != nil {
  792. db.cleanupHook()
  793. }
  794. log.Debug("Expire nodes", "visited", visited, "deleted", deleted, "elapsed", common.PrettyDuration(time.Since(start)))
  795. }