mirror of https://github.com/mongodb/mongo
59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
"""
|
|
A utility for acquiring exclusive access to named local resources from within a
|
|
locally run bazel action.
|
|
|
|
The resource should be defined in the bazel configuration:
|
|
# 4 possible port blocks, each representing the ports that can be used up until the next block.
|
|
--local_resources=port_block=4
|
|
--define=LOCAL_RESOURCES="port_block=20000,20250,20500,20750"
|
|
|
|
The target that needs the resource should add the tag:
|
|
tags = ["resources:port_block:1"] # This action needs one port range
|
|
And include LOCAL_RESOURCES in the action environment:
|
|
env = {
|
|
"LOCAL_RESOURCES": "$(LOCAL_RESOURCES)",
|
|
}
|
|
|
|
The action should use `acquire_local_resource` to get a lock on the local resource:
|
|
lock, port_block = acquire_local_resource("port_block")
|
|
# Use port_block, no other action will be using the same one.
|
|
lock.release()
|
|
"""
|
|
|
|
import os
|
|
import pathlib
|
|
import tempfile
|
|
from functools import cache
|
|
|
|
from filelock import FileLock, Timeout
|
|
|
|
|
|
@cache
|
|
def _parse_local_resources() -> dict:
|
|
resources = {}
|
|
for resource in os.environ.get("LOCAL_RESOURCES").split(";"):
|
|
name, values = resource.split("=", 1)
|
|
resources[name] = values.split(",")
|
|
return resources
|
|
|
|
|
|
def acquire_local_resource(resource_name: str) -> (FileLock, str):
|
|
local_resources = _parse_local_resources()
|
|
if resource_name not in local_resources:
|
|
raise Exception(
|
|
f"Resource {resource_name} not found in LOCAL_RESOURCES. LOCAL_RESOURCES are: {local_resources}"
|
|
)
|
|
|
|
lock_dir = pathlib.Path(tempfile.gettempdir(), "bazel_local_resources", resource_name)
|
|
lock_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
for resource in local_resources[resource_name]:
|
|
try:
|
|
lock = FileLock(lock_dir / f"{resource}.lock")
|
|
lock.acquire(timeout=0)
|
|
return lock, resource
|
|
except Timeout:
|
|
continue
|
|
|
|
raise Exception(f"Could not acquire a lock for resource {resource_name}")
|