97 lines
2.4 KiB
Nim
97 lines
2.4 KiB
Nim
# Nim-Libp2p
|
|
# Copyright (c) 2023 Status Research & Development GmbH
|
|
# Licensed under either of
|
|
# * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
|
|
# * MIT license ([LICENSE-MIT](LICENSE-MIT))
|
|
# at your option.
|
|
# This file may not be copied, modified, or distributed except according to
|
|
# those terms.
|
|
|
|
{.push raises: [].}
|
|
|
|
import sequtils
|
|
import chronos, chronicles
|
|
|
|
# TODO: this should probably go in chronos
|
|
|
|
logScope:
|
|
topics = "libp2p semaphore"
|
|
|
|
type
|
|
AsyncSemaphore* = ref object of RootObj
|
|
size*: int
|
|
count: int
|
|
queue: seq[Future[void]]
|
|
|
|
proc newAsyncSemaphore*(size: int): AsyncSemaphore =
|
|
AsyncSemaphore(size: size, count: size)
|
|
|
|
proc `count`*(s: AsyncSemaphore): int = s.count
|
|
|
|
proc tryAcquire*(s: AsyncSemaphore): bool =
|
|
## Attempts to acquire a resource, if successful
|
|
## returns true, otherwise false
|
|
##
|
|
|
|
if s.count > 0 and s.queue.len == 0:
|
|
s.count.dec
|
|
trace "Acquired slot", available = s.count, queue = s.queue.len
|
|
return true
|
|
|
|
proc acquire*(s: AsyncSemaphore): Future[void] =
|
|
## Acquire a resource and decrement the resource
|
|
## counter. If no more resources are available,
|
|
## the returned future will not complete until
|
|
## the resource count goes above 0.
|
|
##
|
|
|
|
let fut = newFuture[void]("AsyncSemaphore.acquire")
|
|
if s.tryAcquire():
|
|
fut.complete()
|
|
return fut
|
|
|
|
proc cancellation(udata: pointer) {.gcsafe.} =
|
|
fut.cancelCallback = nil
|
|
if not fut.finished:
|
|
s.queue.keepItIf( it != fut )
|
|
|
|
fut.cancelCallback = cancellation
|
|
|
|
s.queue.add(fut)
|
|
|
|
trace "Queued slot", available = s.count, queue = s.queue.len
|
|
return fut
|
|
|
|
proc forceAcquire*(s: AsyncSemaphore) =
|
|
## ForceAcquire will always succeed,
|
|
## creating a temporary slot if required.
|
|
## This temporary slot will stay usable until
|
|
## there is less `acquire`s than `release`s
|
|
s.count.dec
|
|
|
|
proc release*(s: AsyncSemaphore) =
|
|
## Release a resource from the semaphore,
|
|
## by picking the first future from the queue
|
|
## and completing it and incrementing the
|
|
## internal resource count
|
|
##
|
|
|
|
doAssert(s.count <= s.size)
|
|
|
|
if s.count < s.size:
|
|
trace "Releasing slot", available = s.count,
|
|
queue = s.queue.len
|
|
|
|
s.count.inc
|
|
while s.queue.len > 0:
|
|
var fut = s.queue[0]
|
|
s.queue.delete(0)
|
|
if not fut.finished():
|
|
s.count.dec
|
|
fut.complete()
|
|
break
|
|
|
|
trace "Released slot", available = s.count,
|
|
queue = s.queue.len
|
|
return
|