Rework the Kademlia high-level system (#282)

* Kademlia high-level rework

* Some changes in the rework

* Some additional tweaks to kad rework

* Add update_kbuckets

* Rename a bunch of Kademlia stuff

* Add KadSystem::local_peer_id

* Some documentation update

* Concern

* Make the example compile

* Make things nicer

* Fix bug in UniqueConnec

* Add clear() to UniqueConnec

* Add UniqueConnec::poll

* Fix potential deadlock in UniqueConnec

* Add UniqueConnec::state()

* The future of get now contains a Weak

* Fix concerns

core/src/lib.rs

@@ -223,6 +223,7 @@ mod connection_reuse;
 mod keys_proto;
 mod peer_id;
 mod public_key;
+mod unique;
 
 pub mod either;
 pub mod muxing;
@@ -237,4 +238,5 @@ pub use self::peer_id::PeerId;
 pub use self::public_key::PublicKey;
 pub use self::swarm::{swarm, SwarmController, SwarmFuture};
 pub use self::transport::{MuxedTransport, Transport};
+pub use self::unique::{UniqueConnec, UniqueConnecFuture, UniqueConnecState};
 pub use self::upgrade::{ConnectionUpgrade, Endpoint};

core/src/unique.rs

@@ -0,0 +1,333 @@
+// Copyright 2018 Parity Technologies (UK) Ltd.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+use futures::{future, sync::oneshot, task, Async, Future, Poll, IntoFuture};
+use parking_lot::Mutex;
+use {Multiaddr, MuxedTransport, SwarmController, Transport};
+use std::io::{Error as IoError, ErrorKind as IoErrorKind};
+use std::mem;
+use std::sync::{Arc, Weak};
+
+/// Storage for a unique connection with a remote.
+pub struct UniqueConnec<T> {
+    inner: Arc<Mutex<UniqueConnecInner<T>>>,
+}
+
+enum UniqueConnecInner<T> {
+    /// The `UniqueConnec` was created, but nothing is in it.
+    Empty,
+    /// We started dialing, but no response has been obtained so far.
+    Pending {
+        /// Tasks that need to be awakened when the content of this object is set.
+        tasks_waiting: Vec<task::Task>,
+        /// Future that represents when `set_until` should have been called.
+        // TODO: Send + Sync bound is meh
+        dial_fut: Box<Future<Item = (), Error = IoError> + Send + Sync>,
+    },
+    /// The value of this unique connec has been set.
+    /// Can only transition to `Empty` when the future has expired.
+    Full {
+        /// Content of the object.
+        value: T,
+        /// Sender to trigger if the content gets cleared.
+        on_clear: oneshot::Sender<()>,
+    },
+    /// The `dial_fut` has errored.
+    Errored(IoError),
+}
+
+impl<T> UniqueConnec<T> {
+    /// Builds a new empty `UniqueConnec`.
+    #[inline]
+    pub fn empty() -> Self {
+        UniqueConnec {
+            inner: Arc::new(Mutex::new(UniqueConnecInner::Empty)),
+        }
+    }
+
+    /// Builds a new `UniqueConnec` that contains a value.
+    #[inline]
+    pub fn with_value(value: T) -> Self {
+        let (on_clear, _) = oneshot::channel();
+        UniqueConnec {
+            inner: Arc::new(Mutex::new(UniqueConnecInner::Full { value, on_clear })),
+        }
+    }
+
+    /// Instantly returns the value from the object if there is any.
+    pub fn poll(&self) -> Option<T>
+        where T: Clone,
+    {
+        let inner = self.inner.lock();
+        if let UniqueConnecInner::Full { ref value, .. } = &*inner {
+            Some(value.clone())
+        } else {
+            None
+        }
+    }
+
+    /// Loads the value from the object.
+    ///
+    /// If the object is empty, dials the given multiaddress with the given transport.
+    ///
+    /// The closure of the `swarm` is expected to call `set_until()` on the `UniqueConnec`. Failure
+    /// to do so will make the `UniqueConnecFuture` produce an error.
+    pub fn get_or_dial<S, Du>(&self, swarm: &SwarmController<S>, multiaddr: &Multiaddr,
+                              transport: Du) -> UniqueConnecFuture<T>
+        where T: Clone,
+              Du: Transport + 'static, // TODO: 'static :-/
+              Du::Output: Into<S::Output>,
+              S: Clone + MuxedTransport,
+    {
+        self.get(|| {
+            swarm.dial(multiaddr.clone(), transport)
+                .map_err(|_| IoError::new(IoErrorKind::Other, "multiaddress not supported"))
+                .into_future()
+                .flatten()
+        })
+    }
+
+    /// Loads the value from the object.
+    ///
+    /// If the object is empty, calls the closure. The closure should return a future that
+    /// should be signaled after `set_until` has been called. If the future produces an error,
+    /// then the object will empty itself again and the `UniqueConnecFuture` will return an error.
+    /// If the future is finished and `set_until` hasn't been called, then the `UniqueConnecFuture`
+    /// will return an error.
+    pub fn get<F, Fut>(&self, or: F) -> UniqueConnecFuture<T>
+        where F: FnOnce() -> Fut,
+              T: Clone,
+              Fut: IntoFuture<Item = (), Error = IoError>,
+              Fut::Future: Send + Sync + 'static, // TODO: 'static :-/
+    {
+        match &*self.inner.lock() {
+            UniqueConnecInner::Empty => (),
+            _ => return UniqueConnecFuture { inner: Arc::downgrade(&self.inner) },
+        };
+
+        // The mutex is unlocked when we call `or`, in order to avoid potential deadlocks.
+        let dial_fut = or().into_future();
+
+        let mut inner = self.inner.lock();
+        // Since we unlocked the mutex, it's possible that the object was filled in the meanwhile.
+        // Therefore we check again whether it's still `Empty`.
+        if let UniqueConnecInner::Empty = &mut *inner {
+            *inner = UniqueConnecInner::Pending {
+                tasks_waiting: Vec::new(),
+                dial_fut: Box::new(dial_fut),
+            };
+        }
+
+        UniqueConnecFuture { inner: Arc::downgrade(&self.inner) }
+    }
+
+    /// Puts `value` inside the object. The second parameter is a future whose completion will
+    /// clear up the content. Returns an adjusted version of that same future.
+    ///
+    /// If `clear()` is called, the returned future will automatically complete with an error.
+    ///
+    /// Has no effect if the object already contains something.
+    pub fn set_until<F>(&self, value: T, until: F) -> impl Future<Item = (), Error = F::Error>
+        where F: Future<Item = ()>
+    {
+        let mut tasks_to_notify = Vec::new();
+
+        let mut inner = self.inner.lock();
+        let (on_clear, on_clear_rx) = oneshot::channel();
+        match mem::replace(&mut *inner, UniqueConnecInner::Full { value, on_clear }) {
+            UniqueConnecInner::Empty => {},
+            UniqueConnecInner::Errored(_) => {},
+            UniqueConnecInner::Pending { tasks_waiting, .. } => {
+                tasks_to_notify = tasks_waiting;
+            },
+            old @ UniqueConnecInner::Full { .. } => {
+                // Keep the old value.
+                *inner = old;
+                return future::Either::B(until);
+            },
+        };
+        drop(inner);
+
+        // The mutex is unlocked when we notify the pending tasks.
+        for task in tasks_to_notify {
+            task.notify();
+        }
+
+        let inner = self.inner.clone();
+        let fut = until
+            .select(on_clear_rx.then(|_| Ok(())))
+            .map(|((), _)| ())
+            .map_err(|(err, _)| err)
+            .then(move |val| {
+                *inner.lock() = UniqueConnecInner::Empty;
+                val
+            });
+        future::Either::A(fut)
+    }
+
+    /// Clears the content of the object.
+    ///
+    /// Has no effect if the content is empty or pending.
+    /// If the node was full, calling `clear` will stop the future returned by `set_until`.
+    pub fn clear(&self) {
+        let mut inner = self.inner.lock();
+        match mem::replace(&mut *inner, UniqueConnecInner::Empty) {
+            UniqueConnecInner::Empty => {},
+            UniqueConnecInner::Errored(_) => {},
+            pending @ UniqueConnecInner::Pending { .. } => {
+                *inner = pending;
+            },
+            UniqueConnecInner::Full { on_clear, .. } => {
+                let _ = on_clear.send(());
+            },
+        };
+    }
+
+    /// Returns the state of the object.
+    ///
+    /// Note that this can be racy, as the object can be used at the same time. In other words,
+    /// the returned value may no longer reflect the actual state.
+    pub fn state(&self) -> UniqueConnecState {
+        match *self.inner.lock() {
+            UniqueConnecInner::Empty => UniqueConnecState::Empty,
+            UniqueConnecInner::Errored(_) => UniqueConnecState::Errored,
+            UniqueConnecInner::Pending { .. } => UniqueConnecState::Pending,
+            UniqueConnecInner::Full { .. } => UniqueConnecState::Full,
+        }
+    }
+}
+
+impl<T> Clone for UniqueConnec<T> {
+    #[inline]
+    fn clone(&self) -> UniqueConnec<T> {
+        UniqueConnec {
+            inner: self.inner.clone(),
+        }
+    }
+}
+
+impl<T> Default for UniqueConnec<T> {
+    #[inline]
+    fn default() -> Self {
+        UniqueConnec::empty()
+    }
+}
+
+/// Future returned by `UniqueConnec::get()`.
+pub struct UniqueConnecFuture<T> {
+    inner: Weak<Mutex<UniqueConnecInner<T>>>,
+}
+
+impl<T> Future for UniqueConnecFuture<T>
+    where T: Clone
+{    
+    type Item = T;
+    type Error = IoError;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        let inner = match self.inner.upgrade() {
+            Some(inner) => inner,
+            // All the `UniqueConnec` have been destroyed.
+            None => return Err(IoErrorKind::ConnectionAborted.into()),
+        };
+
+        let mut inner = inner.lock();
+        match mem::replace(&mut *inner, UniqueConnecInner::Empty) {
+            UniqueConnecInner::Empty => {
+                // This can happen if `set_until()` is called, and the future expires before the
+                // future returned by `get()` gets polled. This means that the connection has been
+                // closed.
+                Err(IoErrorKind::ConnectionAborted.into())
+            },
+            UniqueConnecInner::Pending { mut tasks_waiting, mut dial_fut } => {
+                match dial_fut.poll() {
+                    Ok(Async::Ready(())) => {
+                        // This happens if we successfully dialed a remote, but the callback
+                        // doesn't call `set_until`. This can be a logic error by the user,
+                        // but could also indicate that the user decided to filter out this
+                        // connection for whatever reason.
+                        *inner = UniqueConnecInner::Errored(IoErrorKind::ConnectionAborted.into());
+                        Err(IoErrorKind::ConnectionAborted.into())
+                    },
+                    Ok(Async::NotReady) => {
+                        tasks_waiting.push(task::current());
+                        *inner = UniqueConnecInner::Pending { tasks_waiting, dial_fut };
+                        Ok(Async::NotReady)
+                    }
+                    Err(err) => {
+                        let tr = IoError::new(IoErrorKind::ConnectionAborted, err.to_string());
+                        *inner = UniqueConnecInner::Errored(err);
+                        Err(tr)
+                    },
+                }
+            },
+            UniqueConnecInner::Full { value, on_clear } => {
+                *inner = UniqueConnecInner::Full {
+                    value: value.clone(),
+                    on_clear
+                };
+                Ok(Async::Ready(value))
+            },
+            UniqueConnecInner::Errored(err) => {
+                let tr = IoError::new(IoErrorKind::ConnectionAborted, err.to_string());
+                *inner = UniqueConnecInner::Errored(err);
+                Err(tr)
+            },
+        }
+    }
+}
+
+/// State of a `UniqueConnec`.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum UniqueConnecState {
+    /// The object is empty.
+    Empty,
+    /// `get_*` has been called and we are waiting for `set_until` to be called.
+    Pending,
+    /// `set_until` has been called.
+    Full,
+    /// The future returned by the closure of `get_*` has errored or has finished before
+    /// `set_until` has been called.
+    Errored,
+}
+
+#[cfg(test)]
+mod tests {
+    use futures::{future, Future};
+    use transport::DeniedTransport;
+    use {UniqueConnec, UniqueConnecState};
+    use swarm;
+
+    #[test]
+    fn invalid_multiaddr_produces_error() {
+        let unique = UniqueConnec::empty();
+        assert_eq!(unique.state(), UniqueConnecState::Empty);
+        let unique2 = unique.clone();
+        let (swarm_ctrl, _swarm_fut) = swarm(DeniedTransport, |_, _| {
+            unique2.set_until((), future::empty())
+        });
+        let fut = unique.get_or_dial(&swarm_ctrl, &"/ip4/1.2.3.4".parse().unwrap(),
+                                     DeniedTransport);
+        assert!(fut.wait().is_err());
+        assert_eq!(unique.state(), UniqueConnecState::Errored);
+    }
+
+    // TODO: more tests
+}