    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

    //! Asynchronous request for an attribute.

    use core::borrow::Borrow;
    use core::future::Future;
    use core::hash::Hash;
    use core::pin::Pin;
    use core::task::{Context, Poll};
    use pin_project_lite::pin_project;
    use uuid::Uuid;

    use crate::db::AttributeValue;
    use crate::error::Error;

    use super::{AttributeValueRef, Database, LoadAttributesLog};

    pin_project! {
        /// Asynchronous request for an attribute in a specific partition.
        #[must_use = "futures do nothing unless you `.await` or poll them"]
        pub struct GetAttributeInPartition<'db, 'i, 'k, T, FS, K>
        where
            T: Send,
            FS: Send,
            K: ?Sized,
        {
            db: &'db Database<T, FS>,
            partition_index: usize,
            vector_id: &'i Uuid,
            key: &'k K,
            #[pin]
            load_attributes_log: Option<Pin<Box<
                dyn 'db + Future<Output = Result<(), Error>>,
            >>>,
            #[pin]
            get_attribute_internal: Option<Pin<Box<
                dyn 'db + Future<Output = Result<Option<AttributeValueRef<'db>>, Error>>,
            >>>,
        }
    }

    impl<'db, 'i, 'k, T, FS, K> GetAttributeInPartition<'db, 'i, 'k, T, FS, K>
    where
        T: Send,
        FS: Send,
        K: ?Sized,
    {
        /// Creates a new asynchronous request for an attribute in a specific
        /// partition.
        pub(super) fn new(
            db: &'db Database<T, FS>,
            partition_index: usize,
            vector_id: &'i Uuid,
            key: &'k K,
        ) -> Self {
            GetAttributeInPartition {
                db,
                partition_index,
                vector_id,
                key,
                load_attributes_log: None,
                get_attribute_internal: None,
            }
        }
    }

    impl<'db, 'i, 'k, T, FS, K> Future for GetAttributeInPartition<'db, 'i, 'k, T, FS, K>
    where
        T: Send,
        FS: Send,
        String: Borrow<K>,
        K: Hash + Eq + ?Sized,
        Database<T, FS>: LoadAttributesLog<'db>,
        'i: 'db,
        'k: 'db,
    {
        type Output = Result<Option<AttributeValue>, Error>;

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            let mut this = self.project();
            loop {
                if let Some(future) = this.get_attribute_internal
                    .as_mut()
                    .as_pin_mut()
                {
                    // 3. waits for the attribute value
                    match future.poll(cx) {
                        Poll::Ready(Ok(result)) => {
                            let value = result.map(|value| value.clone());
                            return Poll::Ready(Ok(value));
                        },
                        Poll::Pending => return Poll::Pending,
                        Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
                    }
                } else if let Some (future) = this.load_attributes_log
                    .as_mut()
                    .as_pin_mut()
                {
                    // 2. requests for the attribute value
                    match future.poll(cx) {
                        Poll::Ready(Ok(_)) => {
                            *this.get_attribute_internal = Some(Box::pin(
                                this.db.get_attribute_internal(
                                    this.vector_id,
                                    this.key,
                                ),
                            ));
                        },
                        Poll::Pending => return Poll::Pending,
                        Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
                    };
                } else {
                    // 1. loads the attributes log
                    *this.load_attributes_log = Some(
                        this.db.load_attributes_log(*this.partition_index),
                    );
                }
            }
        }
    }