31 lines
697 B
Rust
31 lines
697 B
Rust
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
use tokio::sync::Notify;
|
|
|
|
#[derive(Default)]
|
|
pub struct Latch {
|
|
count: AtomicUsize,
|
|
notify: Notify,
|
|
}
|
|
|
|
impl Latch {
|
|
pub fn increment(&self) {
|
|
self.count.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
|
|
pub fn decrement(&self) {
|
|
let previous = self.count.fetch_sub(1, Ordering::Relaxed);
|
|
if previous == 1 {
|
|
self.notify.notify_waiters();
|
|
}
|
|
}
|
|
|
|
pub async fn wait(&self) {
|
|
if self.count.load(Ordering::Relaxed) > 0 {
|
|
self.notify.notified().await;
|
|
}
|
|
}
|
|
}
|