blob: a3725ab3504fd5dc8cfe97ae92aab1d2b3c50491 [file] [log] [blame]
# Copyright 2019 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.
from contextlib import contextmanager
import attr
from gevent.queue import Channel
from ..types import ResourceCost
@attr.s
class ResourceWaiter(object):
"""Represents the machine's CPU, memory, disk and network as limited
resources.
Each cpu (according to multiprocessing.cpu_count()) is worth 1000 millicores.
Every subprocess that attempts to execute will first acquire its estimated
amount of millicores before launching the subprocess. As soon as the
subprocess completes, the held millicores are put back into the pool.
Similarly, memory is measured in megabytes of physical system memory (i.e. not
including swap). The assumption is that the recipe (and its subprocesses) is
really mostly the only thing running on the machine anyway.
Disk and net are different, however. They are unitless measures of 'percentage
of resource', where the absolute quantity of Disk (IOPS, read/write/seek
speed) and Network (bandwidth, latency) are not considered. Consider them more
a declaration that a given step contends on disk or network availability. If
you had many steps which each took `10%` disk, only 10 of them would run at
a time. Similarly, if you had steps which declared 50% of disk bandwidth
usage, only 2 of them would run at a time.
Because recipes are finite both in runtime and number of distinct steps, this
resource class unblocks other processes greedily. Whenever a subprocess
completes, this analyzes all the outstanding subprocesses and will unblock
whichever ones 'fit' in the now-freed resources. This is done in roughly FIFO
order (i.e. if two tasks could potentially fit, the first one to block will be
chosen over the second to unblock first).
This is different than what's deemed 'fair' in a typical scheduling scenario,
because in a mixed workload, heavy tasks could be forced to wait longer while
small tasks use the CPU. However, because the recipes typically run with
a hard finite timeout, it's better to use more of the CPU earlier than to
potentially waste time waiting for small tasks to finish in order to schedule
a heavy task earlier.
"""
# Required for __init__
_millicores_available = attr.ib()
_memory_available = attr.ib()
# Attrs with defaults.
_millicores_max = attr.ib()
@_millicores_max.default
def _millicores_max_default(self):
return self._millicores_available
_memory_max = attr.ib()
@_memory_max.default
def _memory_max_default(self):
return self._memory_available
_disk_available = attr.ib(default=100)
_disk_max = attr.ib(default=100)
_net_available = attr.ib(default=100)
_net_max = attr.ib(default=100)
# List[Tuple[amount, Channel]]
_waiters = attr.ib(factory=list)
def _fits(self, resources):
assert isinstance(resources, ResourceCost)
return resources.fits(self._millicores_available, self._memory_available,
self._disk_available, self._net_available)
def _decr(self, resources):
assert isinstance(resources, ResourceCost)
self._millicores_available -= resources.cpu
self._memory_available -= resources.memory
self._disk_available -= resources.disk
self._net_available -= resources.net
def _incr(self, resources):
assert isinstance(resources, ResourceCost)
self._millicores_available += resources.cpu
self._memory_available += resources.memory
self._disk_available += resources.disk
self._net_available += resources.net
@contextmanager
def wait_for(self, resources, call_if_blocking):
"""Block until `resources` are available.
Args:
* resources (ResourceCost|None) - The amount of various resources to
acquire before yielding. If any aspect of this exceeds the maximum
amount of resource available on the system, this will instead acquire
the system maximum. If resources is all 0's, or is None, this does not
block.
* call_if_blocking (None|func(ResourceCost)) - `wait_for` will invoke this
callback if we would end up blocking before yielding. This callback
should only be used for reporting/diagnostics (i.e. it shouldn't raise
an exception.)
Yields control once the requisite amount of resources are available. Exiting
the context frees up the resources.
"""
if resources is None:
yield
return
assert isinstance(resources, ResourceCost)
if resources.cpu > self._millicores_max:
resources = attr.evolve(resources, cpu=self._millicores_max)
if resources.memory > self._memory_max:
resources = attr.evolve(resources, memory=self._memory_max)
if resources and (self._waiters or not self._fits(resources)):
# we need some amount of resource AND
# someone else is already waiting, or there isn't enough resource.
if call_if_blocking:
call_if_blocking()
wake_me = Channel()
self._waiters.append((resources, wake_me))
self._waiters.sort(reverse=True) # stable sort
wake_me.get()
# At this point the greenlet that woke us already reserved our resources
# for us, and we're free to go.
else:
# Just directly take our cores.
assert self._fits(resources)
self._decr(resources)
try:
yield
finally:
self._incr(resources)
# We just added some resource back to the pot. Try to wake as many others
# as we can before proceeding.
to_wake, to_keep = [], []
for waiting_resources, chan in self._waiters:
if self._fits(waiting_resources):
to_wake.append(chan)
self._decr(waiting_resources)
else:
to_keep.append((waiting_resources, chan))
self._waiters = to_keep # _waiters was sorted before, so to_keep is also.
for chan in to_wake:
chan.put(None)