1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

use std::collections::HashMap;
use std::ops::Deref;
use std::sync::{Arc, Condvar, Mutex, RwLock};

struct Inner<T>
where
    T: Send + Clone,
{
    value: Mutex<Option<T>>,
    waker: Condvar,
}

/// Handle to a complete an associated promise
struct Completer<T>
where
    T: Send + Clone,
{
    inner: Arc<Inner<T>>,
}

impl<T> Completer<T>
where
    T: Send + Clone,
{
    /// Sets the associated promise as complete
    pub fn set(self, new_value: T) {
        let mut value_lock = self.inner.value.lock().unwrap();
        value_lock.replace(new_value);
        drop(value_lock);

        self.inner.waker.notify_all();
    }
}

/// Promise of a future computation
///
/// The Promise is initially incomplete. Once its been completed it will return clones of the
/// completed value until it's dropped.
#[derive(Clone)]
struct Promise<T>
where
    T: Send + Clone,
{
    inner: Arc<Inner<T>>,
}

impl<T> Promise<T>
where
    T: Send + Clone,
{
    /// Waits until the promise is complete and returns a clone of its value
    pub fn value(&self) -> T {
        let mut value_lock = self.inner.value.lock().unwrap();

        loop {
            match value_lock.deref() {
                Some(value) => break value.clone(),
                None => {
                    value_lock = self.inner.waker.wait(value_lock).unwrap();
                }
            }
        }
    }
}

/// Creates new completer and promise
fn promise<T>() -> (Completer<T>, Promise<T>)
where
    T: Send + Clone,
{
    let inner = Arc::new(Inner {
        value: Mutex::new(None),
        waker: Condvar::new(),
    });

    (
        Completer {
            inner: inner.clone(),
        },
        Promise { inner },
    )
}

/// Create an immediately completed promise
fn completed<T>(value: T) -> Promise<T>
where
    T: Send + Clone,
{
    Promise {
        inner: {
            Arc::new(Inner {
                value: Mutex::new(Some(value)),
                waker: Condvar::new(),
            })
        },
    }
}

/// Concurrent map of keys to values where each key is only calculated once
pub struct PromiseMap<K, V>
where
    K: std::hash::Hash,
    V: Send + Clone,
{
    promises: RwLock<HashMap<K, Promise<V>>>,
}

impl<K, V> PromiseMap<K, V>
where
    K: std::hash::Hash + Eq,
    V: Send + Clone,
{
    /// Creates a new `PromiseMap` with the passed values
    pub fn new(values: impl IntoIterator<Item = (K, V)>) -> Self {
        PromiseMap {
            promises: RwLock::new(values.into_iter().map(|(k, v)| (k, completed(v))).collect()),
        }
    }

    /// Fetches the value from the promise map or inserts it if it does not exist
    ///
    /// Each key will only be calculated once. If a calculation is already in progress on another
    /// thread the current thread will block until the existing calculation completes.
    pub fn get_or_insert_with<F>(&self, key: K, func: F) -> V
    where
        F: FnOnce() -> V,
    {
        // Opportunistically try to fetch the promise with a read lock
        let promises_read = self.promises.read().unwrap();
        if let Some(promise) = promises_read.get(&key) {
            let cloned_promise = promise.clone();
            drop(promises_read);

            return cloned_promise.value();
        }

        drop(promises_read);

        // Try again with a write lock to ensure another thread didn't already insert
        let mut promises_write = self.promises.write().unwrap();
        if let Some(promise) = promises_write.get(&key) {
            let cloned_promise = promise.clone();
            drop(promises_write);

            return cloned_promise.value();
        }

        let (completer, promise) = promise();
        promises_write.insert(key, promise);
        drop(promises_write);

        // Build a new value. This is presumably expensive
        let value = func();
        completer.set(value.clone());
        value
    }
}