import tables, heapqueue import chronos type PeerType* = enum None, Incoming, Outgoing PeerFlags = enum Acquired, DeleteOnRelease PeerItem[T] = object data: T peerType: PeerType flags: set[PeerFlags] index: int PeerIndex = object data: int cmp: proc(a, b: PeerIndex): bool {.closure, gcsafe.} PeerScoreCheckCallback*[T] = proc(peer: T): bool {.gcsafe, raises: [Defect].} PeerPool*[A, B] = ref object incNotEmptyEvent: AsyncEvent outNotEmptyEvent: AsyncEvent incNotFullEvent: AsyncEvent outNotFullEvent: AsyncEvent incQueue: HeapQueue[PeerIndex] outQueue: HeapQueue[PeerIndex] registry: Table[B, PeerIndex] storage: seq[PeerItem[A]] cmp: proc(a, b: PeerIndex): bool {.closure, gcsafe.} scoreCheck: PeerScoreCheckCallback[A] maxPeersCount: int maxIncPeersCount: int maxOutPeersCount: int curIncPeersCount: int curOutPeersCount: int acqIncPeersCount: int acqOutPeersCount: int PeerPoolError* = object of CatchableError proc `<`*(a, b: PeerIndex): bool = result = a.cmp(b, a) proc fireNotEmptyEvent[A, B](pool: PeerPool[A, B], item: PeerItem[A]) {.inline.} = if item.peerType == PeerType.Incoming: pool.incNotEmptyEvent.fire() elif item.peerType == PeerType.Outgoing: pool.outNotEmptyEvent.fire() proc fireNotFullEvent[A, B](pool: PeerPool[A, B], item: PeerItem[A]) {.inline.} = if item.peerType == PeerType.Incoming: pool.incNotFullEvent.fire() elif item.peerType == PeerType.Outgoing: pool.outNotFullEvent.fire() iterator pairs*[A, B](pool: PeerPool[A, B]): (B, A) = for peerId, peerIdx in pool.registry: yield (peerId, pool.storage[peerIdx.data].data) proc waitNotEmptyEvent[A, B](pool: PeerPool[A, B], filter: set[PeerType]) {.async.} = if filter == {PeerType.Incoming, PeerType.Outgoing} or filter == {}: var fut1 = pool.incNotEmptyEvent.wait() var fut2 = pool.outNotEmptyEvent.wait() try: discard await one(fut1, fut2) if fut1.finished: if not(fut2.finished): fut2.cancel() pool.incNotEmptyEvent.clear() else: if not(fut1.finished): fut1.cancel() pool.outNotEmptyEvent.clear() except CancelledError: if not(fut1.finished): fut1.cancel() if not(fut2.finished): fut2.cancel() raise elif PeerType.Incoming in filter: await pool.incNotEmptyEvent.wait() pool.incNotEmptyEvent.clear() elif PeerType.Outgoing in filter: await pool.outNotEmptyEvent.wait() pool.outNotEmptyEvent.clear() proc waitNotFullEvent[A, B](pool: PeerPool[A, B], peerType: PeerType) {.async.} = if peerType == PeerType.Incoming: await pool.incNotFullEvent.wait() pool.incNotFullEvent.clear() elif peerType == PeerType.Outgoing: await pool.outNotFullEvent.wait() pool.outNotFullEvent.clear() template getItem[A, B](pool: PeerPool[A, B], filter: set[PeerType]): ptr PeerItem[A] = doAssert((len(pool.outQueue) > 0) or (len(pool.incQueue) > 0)) var pindex: int if filter == {PeerType.Incoming, PeerType.Outgoing}: if len(pool.outQueue) > 0 and len(pool.incQueue) > 0: # Don't think `<` is actually `<` here. if pool.incQueue[0] < pool.outQueue[0]: inc(pool.acqIncPeersCount) pindex = pool.incQueue.pop().data else: inc(pool.acqOutPeersCount) pindex = pool.outQueue.pop().data else: if len(pool.outQueue) > 0: inc(pool.acqOutPeersCount) pindex = pool.outQueue.pop().data else: inc(pool.acqIncPeersCount) pindex = pool.incQueue.pop().data else: if PeerType.Outgoing in filter: inc(pool.acqOutPeersCount) pindex = pool.outQueue.pop().data elif PeerType.Incoming in filter: inc(pool.acqIncPeersCount) pindex = pool.incQueue.pop().data addr(pool.storage[pindex]) proc newPeerPool*[A, B](maxPeers = -1, maxIncomingPeers = -1, maxOutgoingPeers = -1, scoreCheckCb: PeerScoreCheckCallback[A] = nil): PeerPool[A, B] = ## Create new PeerPool. ## ## ``maxPeers`` - maximum number of peers allowed. All the peers which ## exceeds this number will be rejected (``addPeer()`` procedure will return ## ``false``). By default this number is infinite. ## ## ``maxIncomingPeers`` - maximum number of incoming peers allowed. All the ## incoming peers exceeds this number will be rejected. By default this ## number is infinite. ## ## ``maxOutgoingPeers`` - maximum number of outgoing peers allowed. All the ## outgoing peers exceeds this number will be rejected. By default this ## number if infinite. ## ## ``scoreCheckCb`` - callback which will be called for all released peers. ## If callback procedure returns ``false`` peer will be removed from ## PeerPool. ## ## Please note, that if ``maxPeers`` is positive non-zero value, then equation ## ``maxPeers >= maxIncomingPeers + maxOutgoingPeers`` must be ``true``. var res = PeerPool[A, B]() if maxPeers != -1: doAssert(maxPeers >= maxIncomingPeers + maxOutgoingPeers) res.maxPeersCount = if maxPeers < 0: high(int) else: maxPeers res.maxIncPeersCount = if maxIncomingPeers < 0: high(int) else: maxIncomingPeers res.maxOutPeersCount = if maxOutgoingPeers < 0: high(int) else: maxOutgoingPeers res.incNotEmptyEvent = newAsyncEvent() res.outNotEmptyEvent = newAsyncEvent() res.incNotFullEvent = newAsyncEvent() res.outNotFullEvent = newAsyncEvent() res.incQueue = initHeapQueue[PeerIndex]() res.outQueue = initHeapQueue[PeerIndex]() res.registry = initTable[B, PeerIndex]() res.scoreCheck = scoreCheckCb res.storage = newSeq[PeerItem[A]]() proc peerCmp(a, b: PeerIndex): bool {.closure, gcsafe.} = let p1 = res.storage[a.data].data let p2 = res.storage[b.data].data result = p1 < p2 res.cmp = peerCmp result = res proc len*[A, B](pool: PeerPool[A, B]): int = ## Returns number of registered peers in PeerPool ``pool``. This number ## includes all the peers (acquired and available). result = len(pool.registry) proc lenAvailable*[A, B](pool: PeerPool[A, B], filter = {PeerType.Incoming, PeerType.Outgoing}): int {.inline.} = ## Returns number of available peers in PeerPool ``pool`` which satisfies ## filter ``filter``. if PeerType.Incoming in filter: result = result + len(pool.incQueue) if PeerType.Outgoing in filter: result = result + len(pool.outQueue) proc lenAcquired*[A, B](pool: PeerPool[A, B], filter = {PeerType.Incoming, PeerType.Outgoing}): int {.inline.} = ## Returns number of acquired peers in PeerPool ``pool`` which satisifies ## filter ``filter``. if PeerType.Incoming in filter: result = result + pool.acqIncPeersCount if PeerType.Outgoing in filter: result = result + pool.acqOutPeersCount proc checkPeerScore*[A, B](pool: PeerPool[A, B], peer: A): bool {.inline.} = ## Returns ``true`` if peer passing score check. if not(isNil(pool.scoreCheck)): if pool.scoreCheck(peer): result = true else: result = false else: result = true proc deletePeer*[A, B](pool: PeerPool[A, B], peer: A, force = false): bool = ## Remove ``peer`` from PeerPool ``pool``. ## ## Deletion occurs immediately only if peer is available, otherwise it will ## be deleted only when peer will be released. You can change this behavior ## with ``force`` option. mixin getKey var key = getKey(peer) if pool.registry.hasKey(key): let pindex = pool.registry[key].data var item = addr(pool.storage[pindex]) if (PeerFlags.Acquired in item[].flags): if not(force): item[].flags.incl(PeerFlags.DeleteOnRelease) else: if item[].peerType == PeerType.Incoming: dec(pool.curIncPeersCount) dec(pool.acqIncPeersCount) elif item[].peerType == PeerType.Outgoing: dec(pool.curOutPeersCount) dec(pool.acqOutPeersCount) pool.fireNotFullEvent(item[]) # Cleanup storage with default item, and removing key from hashtable. pool.storage[pindex] = PeerItem[A]() pool.registry.del(key) else: if item[].peerType == PeerType.Incoming: # If peer is available, then its copy present in heapqueue, so we need # to remove it. for i in 0 ..< len(pool.incQueue): if pool.incQueue[i].data == pindex: pool.incQueue.del(i) break dec(pool.curIncPeersCount) elif item[].peerType == PeerType.Outgoing: # If peer is available, then its copy present in heapqueue, so we need # to remove it. for i in 0 ..< len(pool.outQueue): if pool.outQueue[i].data == pindex: pool.outQueue.del(i) break dec(pool.curOutPeersCount) pool.fireNotFullEvent(item[]) # Cleanup storage with default item, and removing key from hashtable. pool.storage[pindex] = PeerItem[A]() pool.registry.del(key) result = true proc addPeerImpl[A, B](pool: PeerPool[A, B], peer: A, peerKey: B, peerType: PeerType): PeerIndex = proc onPeerClosed(udata: pointer) {.gcsafe.} = discard pool.deletePeer(peer) var item = PeerItem[A](data: peer, peerType: peerType, index: len(pool.storage)) pool.storage.add(item) var pitem = addr(pool.storage[^1]) let pindex = PeerIndex(data: item.index, cmp: pool.cmp) pool.registry[peerKey] = pindex pitem[].data.getFuture().addCallback(onPeerClosed) result = pindex proc addPeerNoWait*[A, B](pool: PeerPool[A, B], peer: A, peerType: PeerType): bool = ## Add peer ``peer`` of type ``peerType`` to PeerPool ``pool``. ## ## Procedure returns ``false`` in case ## * if ``peer`` is already closed. ## * if ``pool`` already has peer ``peer`` inside. ## * if ``pool`` currently has a maximum of peers. ## * if ``pool`` currently has a maximum of `Incoming` or `Outgoing` peers. ## ## Procedure returns ``true`` on success. mixin getKey, getFuture if not(pool.checkPeerScore(peer)): return false result = false let peerKey = getKey(peer) if not(pool.registry.hasKey(peerKey)) and not(peer.getFuture().finished): if len(pool.registry) < pool.maxPeersCount: if peerType == PeerType.Incoming: if pool.curIncPeersCount < pool.maxIncPeersCount: let pindex = pool.addPeerImpl(peer, peerKey, peerType) inc(pool.curIncPeersCount) pool.incQueue.push(pindex) pool.incNotEmptyEvent.fire() result = true elif peerType == PeerType.Outgoing: if pool.curOutPeersCount < pool.maxOutPeersCount: let pindex = pool.addPeerImpl(peer, peerKey, peerType) inc(pool.curOutPeersCount) pool.outQueue.push(pindex) pool.outNotEmptyEvent.fire() result = true proc addPeer*[A, B](pool: PeerPool[A, B], peer: A, peerType: PeerType): Future[bool] {.async.} = ## Add peer ``peer`` of type ``peerType`` to PeerPool ``pool``. ## ## This procedure will wait for an empty space in PeerPool ``pool``, if ## PeerPool ``pool`` is full. ## ## Procedure returns ``false`` in case: ## * if ``peer`` is already closed. ## * if ``pool`` already has peer ``peer`` inside. ## ## Procedure returns ``true`` on success. mixin getKey, getFuture if not(pool.checkPeerScore(peer)): return false var res = false let peerKey = getKey(peer) if not(pool.registry.hasKey(peerKey)) and not(peer.getFuture().finished): if len(pool.registry) >= pool.maxPeersCount: await pool.waitNotFullEvent(peerType) if peerType == PeerType.Incoming: if pool.curIncPeersCount >= pool.maxIncPeersCount: await pool.waitNotFullEvent(peerType) let pindex = pool.addPeerImpl(peer, peerKey, peerType) inc(pool.curIncPeersCount) pool.incQueue.push(pindex) pool.incNotEmptyEvent.fire() res = true elif peerType == PeerType.Outgoing: if pool.curOutPeersCount >= pool.maxOutPeersCount: await pool.waitNotFullEvent(peerType) let pindex = pool.addPeerImpl(peer, peerKey, peerType) inc(pool.curOutPeersCount) pool.outQueue.push(pindex) pool.outNotEmptyEvent.fire() res = true result = res proc addIncomingPeerNoWait*[A, B](pool: PeerPool[A, B], peer: A): bool {.inline.} = ## Add incoming peer ``peer`` to PeerPool ``pool``. ## ## Returns ``true`` on success. result = pool.addPeerNoWait(peer, PeerType.Incoming) proc addOutgoingPeerNoWait*[A, B](pool: PeerPool[A, B], peer: A): bool {.inline.} = ## Add outgoing peer ``peer`` to PeerPool ``pool``. ## ## Returns ``true`` on success. result = pool.addPeerNoWait(peer, PeerType.Outgoing) proc addIncomingPeer*[A, B](pool: PeerPool[A, B], peer: A): Future[bool] {.inline.} = ## Add incoming peer ``peer`` to PeerPool ``pool``. ## ## Returns ``true`` on success. result = pool.addPeer(peer, PeerType.Incoming) proc addOutgoingPeer*[A, B](pool: PeerPool[A, B], peer: A): Future[bool] {.inline.} = ## Add outgoing peer ``peer`` to PeerPool ``pool``. ## ## Returns ``true`` on success. result = pool.addPeer(peer, PeerType.Outgoing) proc acquire*[A, B](pool: PeerPool[A, B], filter = {PeerType.Incoming, PeerType.Outgoing}): Future[A] {.async.} = ## Acquire peer from PeerPool ``pool``, which match the filter ``filter``. doAssert(filter != {}, "Filter must not be empty") while true: var count = 0 if PeerType.Incoming in filter: count = count + len(pool.incQueue) if PeerType.Outgoing in filter: count = count + len(pool.outQueue) if count == 0: await pool.waitNotEmptyEvent(filter) else: var item = pool.getItem(filter) doAssert(PeerFlags.Acquired notin item[].flags) item[].flags.incl(PeerFlags.Acquired) result = item[].data break proc acquireNoWait*[A, B](pool: PeerPool[A, B], filter = {PeerType.Incoming, PeerType.Outgoing}): A = doAssert(filter != {}, "Filter must not be empty") var count = 0 if PeerType.Incoming in filter: count = count + len(pool.incQueue) if PeerType.Outgoing in filter: count = count + len(pool.outQueue) if count < 1: raise newException(PeerPoolError, "Not enough peers in pool") var item = pool.getItem(filter) doAssert(PeerFlags.Acquired notin item[].flags) item[].flags.incl(PeerFlags.Acquired) result = item[].data proc release*[A, B](pool: PeerPool[A, B], peer: A) = ## Release peer ``peer`` back to PeerPool ``pool`` mixin getKey var key = getKey(peer) var titem = pool.registry.getOrDefault(key, PeerIndex(data: -1)) if titem.data >= 0: let pindex = titem.data var item = addr(pool.storage[pindex]) if PeerFlags.Acquired in item[].flags: if not(pool.checkPeerScore(peer)): item[].flags.incl(DeleteOnRelease) item[].flags.excl(PeerFlags.Acquired) if PeerFlags.DeleteOnRelease in item[].flags: if item[].peerType == PeerType.Incoming: dec(pool.curIncPeersCount) dec(pool.acqIncPeersCount) elif item[].peerType == PeerType.Outgoing: dec(pool.curOutPeersCount) dec(pool.acqOutPeersCount) pool.storage[pindex] = PeerItem[A]() pool.registry.del(key) else: if item[].peerType == PeerType.Incoming: pool.incQueue.push(titem) dec(pool.acqIncPeersCount) elif item[].peerType == PeerType.Outgoing: pool.outQueue.push(titem) dec(pool.acqOutPeersCount) pool.fireNotEmptyEvent(item[]) proc release*[A, B](pool: PeerPool[A, B], peers: openarray[A]) {.inline.} = ## Release array of peers ``peers`` back to PeerPool ``pool``. for item in peers: pool.release(item) proc acquire*[A, B](pool: PeerPool[A, B], number: int, filter = {PeerType.Incoming, PeerType.Outgoing}): Future[seq[A]] {.async.} = ## Acquire ``number`` number of peers from PeerPool ``pool``, which match the ## filter ``filter``. doAssert(filter != {}, "Filter must not be empty") var peers = newSeq[A]() try: if number > 0: while true: if len(peers) >= number: break var count = 0 if PeerType.Incoming in filter: count = count + len(pool.incQueue) if PeerType.Outgoing in filter: count = count + len(pool.outQueue) if count == 0: await pool.waitNotEmptyEvent(filter) else: var item = pool.getItem(filter) doAssert(PeerFlags.Acquired notin item[].flags) item[].flags.incl(PeerFlags.Acquired) peers.add(item[].data) except CancelledError: # If we got cancelled, we need to return all the acquired peers back to # pool. for item in peers: pool.release(item) peers.setLen(0) raise result = peers proc acquireNoWait*[A, B](pool: PeerPool[A, B], number: int, filter = {PeerType.Incoming, PeerType.Outgoing}): seq[A] = ## Acquire ``number`` number of peers from PeerPool ``pool``, which match the ## filter ``filter``. doAssert(filter != {}, "Filter must not be empty") var peers = newSeq[A]() var count = 0 if PeerType.Incoming in filter: count = count + len(pool.incQueue) if PeerType.Outgoing in filter: count = count + len(pool.outQueue) if count < number: raise newException(PeerPoolError, "Not enough peers in pool") for i in 0 ..< number: var item = pool.getItem(filter) doAssert(PeerFlags.Acquired notin item[].flags) item[].flags.incl(PeerFlags.Acquired) peers.add(item[].data) result = peers proc acquireIncomingPeer*[A, B](pool: PeerPool[A, B]): Future[A] {.inline.} = ## Acquire single incoming peer from PeerPool ``pool``. pool.acquire({PeerType.Incoming}) proc acquireOutgoingPeer*[A, B](pool: PeerPool[A, B]): Future[A] {.inline.} = ## Acquire single outgoing peer from PeerPool ``pool``. pool.acquire({PeerType.Outgoing}) proc acquireIncomingPeers*[A, B](pool: PeerPool[A, B], number: int): Future[seq[A]] {.inline.} = ## Acquire ``number`` number of incoming peers from PeerPool ``pool``. pool.acquire(number, {PeerType.Incoming}) proc acquireOutgoingPeers*[A, B](pool: PeerPool[A, B], number: int): Future[seq[A]] {.inline.} = ## Acquire ``number`` number of outgoing peers from PeerPool ``pool``. pool.acquire(number, {PeerType.Outgoing}) iterator peers*[A, B](pool: PeerPool[A, B], filter = {PeerType.Incoming, PeerType.Outgoing}): A = ## Iterate over sorted list of peers. ## ## All peers will be sorted by equation `>`(Peer1, Peer2), so biggest values ## will be first. var sorted = initHeapQueue[PeerIndex]() for i in 0 ..< len(pool.storage): if pool.storage[i].peerType in filter: sorted.push(PeerIndex(data: i, cmp: pool.cmp)) while len(sorted) > 0: let pindex = sorted.pop().data yield pool.storage[pindex].data iterator availablePeers*[A, B](pool: PeerPool[A, B], filter = {PeerType.Incoming, PeerType.Outgoing}): A = ## Iterate over sorted list of available peers. ## ## All peers will be sorted by equation `>`(Peer1, Peer2), so biggest values ## will be first. var sorted = initHeapQueue[PeerIndex]() for i in 0 ..< len(pool.storage): if (PeerFlags.Acquired notin pool.storage[i].flags) and (pool.storage[i].peerType in filter): sorted.push(PeerIndex(data: i, cmp: pool.cmp)) while len(sorted) > 0: let pindex = sorted.pop().data yield pool.storage[pindex].data iterator acquiredPeers*[A, B](pool: PeerPool[A, B], filter = {PeerType.Incoming, PeerType.Outgoing}): A = ## Iterate over sorted list of acquired (non-available) peers. ## ## All peers will be sorted by equation `>`(Peer1, Peer2), so biggest values ## will be first. var sorted = initHeapQueue[PeerIndex]() for i in 0 ..< len(pool.storage): if (PeerFlags.Acquired in pool.storage[i].flags) and (pool.storage[i].peerType in filter): sorted.push(PeerIndex(data: i, cmp: pool.cmp)) while len(sorted) > 0: let pindex = sorted.pop().data yield pool.storage[pindex].data proc `[]`*[A, B](pool: PeerPool[A, B], key: B): A {.inline.} = ## Retrieve peer with key ``key`` from PeerPool ``pool``. let pindex = pool.registry[key] result = pool.storage[pindex.data] proc `[]`*[A, B](pool: var PeerPool[A, B], key: B): var A {.inline.} = ## Retrieve peer with key ``key`` from PeerPool ``pool``. let pindex = pool.registry[key] result = pool.storage[pindex.data].data proc hasPeer*[A, B](pool: PeerPool[A, B], key: B): bool {.inline.} = ## Returns ``true`` if peer with ``key`` present in PeerPool ``pool``. result = pool.registry.hasKey(key) proc getOrDefault*[A, B](pool: PeerPool[A, B], key: B): A {.inline.} = ## Retrieves the peer from PeerPool ``pool`` using key ``key``. If peer is ## not present, default initialization value for type ``A`` is returned ## (e.g. 0 for any integer type). let pindex = pool.registry.getOrDefault(key, PeerIndex(data: -1)) if pindex.data >= 0: result = pool.storage[pindex.data].data proc getOrDefault*[A, B](pool: PeerPool[A, B], key: B, default: A): A {.inline.} = ## Retrieves the peer from PeerPool ``pool`` using key ``key``. If peer is ## not present, default value ``default`` is returned. let pindex = pool.registry.getOrDefault(key, PeerIndex(data: -1)) if pindex.data >= 0: result = pool.storage[pindex.data].data else: result = default proc clear*[A, B](pool: PeerPool[A, B]) = ## Performs PeerPool's ``pool`` storage and counters reset. pool.incQueue.clear() pool.outQueue.clear() pool.registry.clear() for i in 0 ..< len(pool.storage): pool.storage[i] = PeerItem[A]() pool.storage.setLen(0) pool.curIncPeersCount = 0 pool.curOutPeersCount = 0 pool.acqIncPeersCount = 0 pool.acqOutPeersCount = 0 proc clearSafe*[A, B](pool: PeerPool[A, B]) {.async.} = ## Performs "safe" clear. Safe means that it first acquires all the peers ## in PeerPool, and only after that it will reset storage. var acquired = newSeq[A]() while len(pool.registry) > len(acquired): var peers = await pool.acquire(len(pool.registry) - len(acquired)) for item in peers: acquired.add(item) pool.clear() proc setScoreCheck*[A, B](pool: PeerPool[A, B], scoreCheckCb: PeerScoreCheckCallback[A]) = ## Add ScoreCheck callback. pool.scoreCheck = scoreCheckCb