Структура Semaphore

Счётный семафор, выполняющий асинхронное получение разрешений.

Семафор поддерживает набор разрешений. Разрешения используются для синхронизации доступа к общему ресурсу. Семафор отличается от мьютекса тем, что позволяет нескольким вызывающим сторонам одновременно обращаться к общему ресурсу.

Когда вызывается acquire и у семафора остаются разрешения, функция немедленно возвращает разрешение. Однако если свободных разрешений нет, acquire асинхронно ожидает, пока одно из занятых разрешений не будет освобождено. После этого освобождённое разрешение передаётся вызывающей стороне.

Этот Semaphore справедлив: разрешения выдаются в порядке поступления запросов. Это правило справедливости действует и при использовании acquire_many, поэтому, если вызов acquire_many в начале очереди запрашивает больше разрешений, чем доступно в данный момент, он может помешать завершению вызова acquire, даже если у семафора достаточно разрешений, чтобы завершить вызов acquire.

Чтобы использовать Semaphore в функции poll, можно воспользоваться утилитой PollSemaphore.

Порядок операций с памятью

Если задача записывает какие-либо данные, а затем освобождает разрешение, любая задача, которая впоследствии получит разрешение, гарантированно увидит эти данные. Благодаря этому семафор можно безопасно использовать для передачи данных между задачами через общее состояние.

Если выразить это точнее в терминах порядков памяти атомарных операций: получение разрешения (через acquire, acquire_many, try_acquire, try_acquire_many или их варианты _owned), освобождение разрешений (путём сброса SemaphorePermit или OwnedSemaphorePermit либо вызова add_permits или forget_permits) и закрытие семафора (через close) — это операции AcqRel. Они полностью упорядочены, и каждая из них синхронизируется со всеми такими операциями, которые предшествуют ей, обеспечивая те же гарантии, что и операции AcqRel над одним атомарным объектом.

Неудачная попытка получить разрешение (включая TryAcquireError::NoPermits и TryAcquireError::Closed), а также методы available_permits и is_closed ведут себя как загрузка Acquire.

Примеры

Простое использование:

use tokio::sync::{Semaphore, TryAcquireError};

let semaphore = Semaphore::new(3);

let a_permit = semaphore.acquire().await.unwrap();
let two_permits = semaphore.acquire_many(2).await.unwrap();

assert_eq!(semaphore.available_permits(), 0);

let permit_attempt = semaphore.try_acquire();
assert_eq!(permit_attempt.err(), Some(TryAcquireError::NoPermits));

Ограничение количества одновременно открытых файлов в программе

В большинстве операционных систем есть ограничения на количество открытых файловых дескрипторов. Даже в системах без явно заданных ограничений ресурсные ограничения косвенно устанавливают верхний предел числа открытых файлов. Если программа попытается открыть слишком много файлов и превысит этот предел, произойдёт ошибка.

В этом примере используется Semaphore с 100 разрешениями. Получая разрешение от Semaphore перед обращением к файлу, вы гарантируете, что программа будет открывать не более 100 файлов одновременно. При попытке открыть 101-й файл программа будет ждать, пока не появится свободное разрешение, и лишь затем продолжит открывать файл.

use std::io::Result;
use tokio::fs::File;
use tokio::sync::Semaphore;
use tokio::io::AsyncWriteExt;

static PERMITS: Semaphore = Semaphore::const_new(100);

async fn write_to_file(message: &[u8]) -> Result<()> {
    let _permit = PERMITS.acquire().await.unwrap();
    let mut buffer = File::create("example.txt").await?;
    buffer.write_all(message).await?;
    Ok(()) // Permit goes out of scope here, and is available again for acquisition
}

Ограничение количества одновременно отправляемых исходящих запросов

В некоторых случаях может потребоваться ограничить количество исходящих запросов, отправляемых параллельно. Это может быть связано с ограничениями используемого API или сетевых ресурсов системы, на которой выполняется приложение.

В этом примере используется Arc<Semaphore> с 10 разрешениями. Каждая созданная задача получает ссылку на семафор путём клонирования Arc<Semaphore>. Прежде чем отправить запрос, задача должна получить разрешение от семафора, вызвав Semaphore::acquire. Это гарантирует, что в любой момент времени параллельно отправляется не более 10 запросов. Отправив запрос, задача освобождает разрешение, чтобы другие задачи могли отправлять запросы.

use std::sync::Arc;
use tokio::sync::Semaphore;

// Define maximum number of parallel requests.
let semaphore = Arc::new(Semaphore::new(5));
// Spawn many tasks that will send requests.
let mut jhs = Vec::new();
for task_id in 0..50 {
    let semaphore = semaphore.clone();
    let jh = tokio::spawn(async move {
        // Acquire permit before sending request.
        let _permit = semaphore.acquire().await.unwrap();
        // Send the request.
        let response = send_request(task_id).await;
        // Drop the permit after the request has been sent.
        drop(_permit);
        // Handle response.
        // ...

        response
    });
    jhs.push(jh);
}
// Collect responses from tasks.
let mut responses = Vec::new();
for jh in jhs {
    let response = jh.await.unwrap();
    responses.push(response);
}
// Process responses.
// ...

Ограничение количества одновременно обрабатываемых входящих запросов

Как и количество одновременно открытых файлов, количество сетевых дескрипторов ограничено. Обработка неограниченного числа запросов может привести к отказу в обслуживании и многим другим проблемам.

В этом примере вместо глобальной переменной используется Arc<Semaphore>. Чтобы ограничить количество запросов, которые могут обрабатываться одновременно, мы получаем разрешение для каждой задачи до её создания. Получив разрешение, мы создаём новую задачу; после завершения задачи разрешение освобождается внутри неё, позволяя создавать другие задачи. Разрешения необходимо получать с помощью Semaphore::acquire_owned, чтобы их можно было перемещать между границами задач. (Поскольку наш семафор не является глобальной переменной — если бы он был глобальным, достаточно было бы acquire.)

use std::sync::Arc;
use tokio::sync::Semaphore;
use tokio::net::TcpListener;

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let semaphore = Arc::new(Semaphore::new(3));
    let listener = TcpListener::bind("127.0.0.1:8080").await?;

    loop {
        // Acquire permit before accepting the next socket.
        //
        // We use `acquire_owned` so that we can move `permit` into
        // other tasks.
        let permit = semaphore.clone().acquire_owned().await.unwrap();
        let (mut socket, _) = listener.accept().await?;

        tokio::spawn(async move {
            // Do work using the socket.
            handle_connection(&mut socket).await;
            // Drop socket while the permit is still live.
            drop(socket);
            // Drop the permit, so more tasks can be created.
            drop(permit);
        });
    }
}

Запрет параллельного выполнения тестов

По умолчанию Rust запускает тесты из одного файла параллельно. Однако в некоторых случаях параллельное выполнение двух тестов может привести к проблемам. Например, такое может произойти, если тесты используют одну и ту же базу данных.

Рассмотрим следующий сценарий:

1. test_insert: вставляет в базу данных пару «ключ-значение», а затем извлекает значение по тому же ключу, чтобы проверить вставку.
2. test_update: вставляет ключ, затем обновляет его, присваивая новое значение, и проверяет, что значение было обновлено правильно.
3. test_others: третий тест, который не изменяет состояние базы данных. Он может выполняться параллельно с другими тестами.

В этом примере test_insert и test_update должны выполняться последовательно, но порядок запуска тестов не имеет значения. Для решения этой задачи можно использовать семафор с одним разрешением.

use tokio::sync::Semaphore;

// Initialize a static semaphore with only one permit, which is used to
// prevent test_insert and test_update from running in parallel.
static PERMIT: Semaphore = Semaphore::const_new(1);

// Initialize the database that will be used by the subsequent tests.
static DB: Database = Database::setup();

#[tokio::test]
async fn test_insert() {
    // Acquire permit before proceeding. Since the semaphore has only one permit,
    // the test will wait if the permit is already acquired by other tests.
    let permit = PERMIT.acquire().await.unwrap();

    // Do the actual test stuff with database

    // Insert a key-value pair to database
    let (key, value) = ("name", 0);
    DB.insert(key, value).await;

    // Verify that the value has been inserted correctly.
    assert_eq!(DB.get(key).await, value);

    // Undo the insertion, so the database is empty at the end of the test.
    DB.delete(key).await;

    // Drop permit. This allows the other test to start running.
    drop(permit);
}

#[tokio::test]
async fn test_update() {
    // Acquire permit before proceeding. Since the semaphore has only one permit,
    // the test will wait if the permit is already acquired by other tests.
    let permit = PERMIT.acquire().await.unwrap();

    // Do the same insert.
    let (key, value) = ("name", 0);
    DB.insert(key, value).await;

    // Update the existing value with a new one.
    let new_value = 1;
    DB.update(key, new_value).await;

    // Verify that the value has been updated correctly.
    assert_eq!(DB.get(key).await, new_value);

    // Undo any modificattion.
    DB.delete(key).await;

    // Drop permit. This allows the other test to start running.
    drop(permit);
}

#[tokio::test]
async fn test_others() {
    // This test can run in parallel with test_insert and test_update,
    // so it does not use PERMIT.
}

Ограничение частоты с помощью корзины токенов

В этом примере демонстрируются методы add_permits и SemaphorePermit::forget.

Во многих приложениях и системах существуют ограничения на частоту выполнения определённых операций. Превышение этой частоты может привести к снижению производительности или даже к ошибкам.

В этом примере реализовано ограничение частоты с помощью корзины токенов. Корзина токенов — это способ ограничения частоты, который срабатывает не сразу, позволяя обрабатывать короткие всплески одновременно поступающих запросов.

При использовании корзины токенов каждый поступающий запрос расходует один токен, а токены пополняются с определённой частотой, задающей предел запросов. Когда поступает всплеск запросов, токены выдаются немедленно, пока корзина не опустеет. После этого запросам придётся ждать добавления новых токенов.

В отличие от примера, ограничивающего количество запросов, которые можно обрабатывать одновременно, мы не возвращаем токены после завершения обработки запроса. Вместо этого токены добавляются только задачей таймера.

Обратите внимание, что эта реализация неэффективна при малой длительности интервала, поскольку постоянно выполняет циклы и ожидание, расходуя много ресурсов CPU.

use std::sync::Arc;
use tokio::sync::Semaphore;
use tokio::time::{interval, Duration};

struct TokenBucket {
    sem: Arc<Semaphore>,
    jh: tokio::task::JoinHandle<()>,
}

impl TokenBucket {
    fn new(duration: Duration, capacity: usize) -> Self {
        let sem = Arc::new(Semaphore::new(capacity));

        // refills the tokens at the end of each interval
        let jh = tokio::spawn({
            let sem = sem.clone();
            let mut interval = interval(duration);
            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

            async move {
                loop {
                    interval.tick().await;

                    if sem.available_permits() < capacity {
                        sem.add_permits(1);
                    }
                }
            }
        });

        Self { jh, sem }
    }

    async fn acquire(&self) {
        // This can return an error if the semaphore is closed, but we
        // never close it, so this error can never happen.
        let permit = self.sem.acquire().await.unwrap();
        // To avoid releasing the permit back to the semaphore, we use
        // the `SemaphorePermit::forget` method.
        permit.forget();
    }
}

impl Drop for TokenBucket {
    fn drop(&mut self) {
        // Kill the background task so it stops taking up resources when we
        // don't need it anymore.
        self.jh.abort();
    }
}

let capacity = 5;
let update_interval = Duration::from_secs_f32(1.0 / capacity as f32);
let bucket = TokenBucket::new(update_interval, capacity);

for _ in 0..5 {
    bucket.acquire().await;

    // do the operation
}