nim-libp2p-experimental/libp2p/utils/semaphore.nim

100 lines
2.5 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.
when (NimMajor, NimMinor) < (1, 4):
{.push raises: [Defect].}
else:
{.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